import logging
import os
import re
from typing import List, Optional

import requests

try:
    from dotenv import load_dotenv

    load_dotenv()
except ImportError:
    pass

logger = logging.getLogger(__name__)


def extract_json(text: str) -> Optional[str]:
    """
    Robustly extracts the first JSON object or array from a string that may contain
    natural language prefix/suffix. Also purges <think> blocks and markdown ticks.
    """
    if not text:
        return None

    # 100% Autonomous: Scrub model's internal thinking process
    if "<think>" in text:
        text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
        logger.debug("🧠 [LLM] Scoped thinking block detected and purged.")

    # Remove markdown code block formats
    text = re.sub(r"^```json\s*", "", text, flags=re.MULTILINE)
    text = re.sub(r"^```\s*", "", text, flags=re.MULTILINE)

    # Try perfect json block extraction first
    match = re.search(r"(\{.*\}|\[.*\])", text, re.DOTALL)
    if match:
        candidate = match.group(0)
        try:
            import json

            json.loads(candidate)
            return candidate
        except Exception:
            pass

    # Smart Fallback: Truncated JSON Healing
    # If standard validation fails (e.g., due to EOF truncation by local models),
    # run a regex extraction pass over the raw generated text to safely salvage
    # all key-value pairs that *were* successfully completed before the truncation.
    import json

    matches = re.findall(r'"([a-zA-Z0-9_]+)"\s*:\s*(?:([0-9.-]+)|"([^"\\]*(?:\\.[^"\\]*)*)")', text)
    if matches:
        res = {}
        for k, num, obj in matches:
            if num:
                try:
                    res[k] = float(num) if "." in num else int(num)
                except ValueError:
                    res[k] = num
            else:
                res[k] = obj.replace('\\"', '"')

        recovered_json = json.dumps(res)
        logger.warning(f"🔧 [Fuzzy Parse] Successfully salvaged {len(res)} keys from heavily truncated LLM output.")
        return recovered_json

    return None


_MODEL_PRICING_CACHE = None


def get_model_pricing(model_id: str) -> dict:
    global _MODEL_PRICING_CACHE
    if _MODEL_PRICING_CACHE is None:
        try:
            r = requests.get("https://openrouter.ai/api/v1/models", timeout=5)
            if r.status_code == 200:
                models = r.json().get("data", [])
                _MODEL_PRICING_CACHE = {m["id"]: m.get("pricing", {}) for m in models}
            else:
                _MODEL_PRICING_CACHE = {}
        except Exception:
            _MODEL_PRICING_CACHE = {}

    # Check if exact match exists, if not, try partial matches (e.g., if version suffixes differ)
    if _MODEL_PRICING_CACHE and model_id not in _MODEL_PRICING_CACHE:
        for k, v in _MODEL_PRICING_CACHE.items():
            if model_id in k or k in model_id:
                return v

    return _MODEL_PRICING_CACHE.get(model_id, {})


def prewarm_ollama_models(configs):
    """
    Sends a dummy request to the configured local Ollama API endpoints via a background thread
    to force the models to load into VRAM during bot startup, minimizing initial connection latency
    and avoiding timeouts downstream.
    """
    args = configs.args

    def _warmup():
        models_to_warm = set()

        # Collect unique local models
        for attr, url_attr in [
            ("ai_telepathic_model", "ai_telepathic_url"),
            ("ai_fallback_model", "ai_fallback_url"),
            ("ai_condenser_model", "ai_condenser_url"),
            ("ai_model", "ai_model_url"),
        ]:
            url = getattr(args, url_attr, "")
            model = getattr(args, attr, "")
            if model and url and ("localhost" in url or "127.0.0.1" in url):
                models_to_warm.add((url, model))

        for url, model in models_to_warm:
            logger.info(
                f"🔥 [VRAM Pre-Warm] Instructing local Ollama engine to load {model} into memory in the background..."
            )
            try:
                # Fire an ultra-short generation to force it into VRAM
                requests.post(
                    url,
                    json={"model": model, "prompt": "Hi", "stream": False, "options": {"num_predict": 1}},
                    timeout=120,
                )
            except Exception:
                pass

    if hasattr(args, "ai_telepathic_model"):
        import threading

        threading.Thread(target=_warmup, daemon=True).start()


def unload_ollama_models(configs):
    """
    Sends keep_alive: 0 to all configured local Ollama API endpoints via a background thread
    to force the models to unload from VRAM during bot shutdown.
    """
    args = configs.args

    def _unload():
        models_to_unload = set()

        # Collect unique local models
        for attr, url_attr in [
            ("ai_telepathic_model", "ai_telepathic_url"),
            ("ai_fallback_model", "ai_fallback_url"),
            ("ai_condenser_model", "ai_condenser_url"),
            ("ai_model", "ai_model_url"),
        ]:
            url = getattr(args, url_attr, "")
            model = getattr(args, attr, "")
            if model and url and ("localhost" in url or "127.0.0.1" in url):
                models_to_unload.add((url, model))

        for url, model in models_to_unload:
            logger.info(f"❄️ [VRAM Cleanup] Instructing local Ollama engine to unload {model} from memory...")
            try:
                # Fire keep_alive: 0 to unload it from VRAM
                requests.post(url, json={"model": model, "keep_alive": 0}, timeout=5)
            except Exception as e:
                logger.debug(f"Failed to unload {model}: {e}")

    if hasattr(args, "ai_telepathic_model"):
        import threading

        threading.Thread(target=_unload, daemon=True).start()


def log_openrouter_burn():
    """Fetches and logs the current OpenRouter API key usage (money burned) ONLY if OpenRouter is actively used."""
    key = os.environ.get("OPENROUTER_API_KEY")
    if not key:
        return

    try:
        from GramAddict.core.config import Config

        args = Config().args
        uses_openrouter = False

        # Check all possible model/url endpoints for 'openrouter'
        for attr in [
            "ai_model",
            "ai_model_url",
            "ai_telepathic_model",
            "ai_telepathic_url",
            "ai_fallback_model",
            "ai_fallback_url",
            "ai_condenser_model",
            "ai_condenser_url",
        ]:
            val = getattr(args, attr, "")
            if val and "openrouter" in str(val).lower():
                uses_openrouter = True
                break

        if not uses_openrouter:
            return
    except Exception:
        pass

    try:
        r = requests.get("https://openrouter.ai/api/v1/auth/key", headers={"Authorization": f"Bearer {key}"}, timeout=5)
        if r.status_code == 200:
            data = r.json().get("data", {})
            total_spent = data.get("usage", 0.0)
            daily_spent = data.get("usage_daily", 0.0)
            limit = data.get("limit")

            logger.info(
                f"🔥 [OpenRouter Burn Rate] Daily: ${daily_spent:.4f} | Total: ${total_spent:.4f}"
                + (f" | Limit: ${limit}" if limit else ""),
                extra={"color": "\x1b[38;5;208m\x1b[1m"},
            )
    except Exception as e:
        logger.debug(f"Could not fetch OpenRouter burn rate: {e}")


def query_llm(
    url: str,
    model: str,
    prompt: str,
    images_b64: Optional[List[str]] = None,
    system: Optional[str] = None,
    format_json: bool = False,
    timeout: int = 180,
    fallback_model: Optional[str] = None,
    fallback_url: Optional[str] = None,
    temperature: Optional[float] = None,
    max_tokens: Optional[int] = None,
) -> Optional[dict]:
    """
    Unified LLM API Caller with configurable fallback.
    """
    openrouter_key = os.environ.get("OPENROUTER_API_KEY")

    # URL-based provider detection (not model-name based — works for any model)
    is_openai_compat = "/v1/chat/completions" in url or "openrouter.ai" in url.lower() or "openai.com" in url.lower()

    # If using a cloud model but a local URL was passed, fix it
    if not is_openai_compat and ("openrouter" in model.lower() or "/" in model):
        # Model looks like "org/model-name" which is OpenRouter format
        is_openai_compat = True
        url = "https://openrouter.ai/api/v1/chat/completions"

    headers = {"Content-Type": "application/json"}

    if is_openai_compat:
        if openrouter_key:
            headers["Authorization"] = f"Bearer {openrouter_key}"

        messages = []
        if system:
            messages.append({"role": "system", "content": system})

        user_content = []
        if prompt:
            user_content.append({"type": "text", "text": prompt})

        if images_b64:
            for img in images_b64:
                user_content.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img}"}})

        messages.append({"role": "user", "content": user_content if len(user_content) > 1 else prompt})

        req_data = {"model": model, "messages": messages, "stream": False}
        if format_json:
            req_data["response_format"] = {"type": "json_object"}
        if temperature is not None:
            req_data["temperature"] = temperature
        if max_tokens is not None:
            req_data["max_tokens"] = max_tokens

    else:
        # Ollama /generate API
        req_data = {"model": model, "prompt": prompt, "stream": False}
        if system:
            req_data["system"] = system
        if images_b64:
            req_data["images"] = images_b64
        if format_json:
            req_data["format"] = "json"
        else:
            # For free-text calls (Brain action extraction), explicitly disable
            # thinking mode. Reasoning models like qwen3.5 put EVERYTHING in
            # the thinking block and return response='', which is useless for
            # action extraction. think=false forces a direct response.
            req_data["think"] = False

        # Ollama passes configs inside 'options'
        if temperature is not None or max_tokens is not None:
            req_data["options"] = {}
            if temperature is not None:
                req_data["options"]["temperature"] = temperature
            if max_tokens is not None:
                req_data["options"]["num_predict"] = max_tokens

    try:
        response = requests.post(url, json=req_data, headers=headers, timeout=timeout)
        response.raise_for_status()
        resp_json = response.json()

        # Normalize response payload so callers don't have to distinguish
        if is_openai_compat:
            # OpenRouter returns choices[0].message.content
            content = resp_json.get("choices", [{}])[0].get("message", {}).get("content", "")

            usage = resp_json.get("usage", {})
            if usage:
                cost_str = ""
                # Attempt to get precise cost sent by OpenRouter or calculate it manually
                if "total_cost" in usage:
                    cost_str = f" | 💸 Cost: ${usage['total_cost']:.6f}"
                else:
                    pricing = get_model_pricing(model)
                    if pricing:
                        try:
                            p_cost = float(pricing.get("prompt", 0)) * usage.get("prompt_tokens", 0)
                            c_cost = float(pricing.get("completion", 0)) * usage.get("completion_tokens", 0)
                            calc_cost = p_cost + c_cost
                            if calc_cost > 0:
                                cost_str = f" | 💸 Cost: ${calc_cost:.6f}"
                        except Exception:
                            pass

                p_tokens = usage.get("prompt_tokens", 0)
                c_tokens = usage.get("completion_tokens", 0)
                t_tokens = usage.get("total_tokens", 0)

                # Make it stand out!
                logger.info(
                    f"🪙  [LLM Burn] {model} -> In: {p_tokens} | Out: {c_tokens} | Total: {t_tokens}{cost_str}",
                    extra={"color": "\x1b[38;5;208m\x1b[1m"},
                )

            # Validation: if JSON was expected, try to extract it
            if format_json:
                extracted = extract_json(content)
                if not extracted:
                    raise ValueError(f"OpenRouter returned non-JSON content when JSON was expected: {content[:100]}...")
                content = extracted

            return {"response": content}
        else:
            # Ollama returns response OR thinking (for reasoning models)
            raw_response = resp_json.get("response", "")
            raw_thinking = resp_json.get("thinking", "")

            logger.debug(f"DEBUG LLM PAYLOAD: response='{raw_response}', thinking='{raw_thinking}'")

            # CRITICAL: For free-text mode (format_json=False), do NOT substitute
            # thinking for empty response. The thinking block is REASONING, not
            # a decision. The Brain parser would extract random actions from it.
            # For JSON mode (format_json=True), falling back to thinking IS correct
            # because reasoning models may place structured output in the thinking block.
            if format_json:
                content = raw_response or raw_thinking or ""
                extracted = extract_json(content)
                if not extracted:
                    logger.warning(f"Failed to extract JSON from content: {content[:100]}")
                else:
                    content = extracted
            else:
                content = raw_response

            return {"response": content}
    except requests.exceptions.ConnectionError:
        logger.error(f"⚠️ [LLM Provider] Connection refused for {model} at {url}. Is the service running?")
    except Exception as e:
        logger.error(f"LLM Provider Error with {model}: {e}")

        # Prevent infinite fallback loops
        if getattr(query_llm, "_is_fallback", False):
            return None

        # Decide on fallback model/url
        f_model = fallback_model
        f_url = fallback_url

        # Read fallback config from args if available
        if not f_model or not f_url:
            from GramAddict.core.config import Config

            try:
                args = Config().args
                f_model = f_model or getattr(args, "ai_fallback_model", None)
                f_url = f_url or getattr(args, "ai_fallback_url", None)
            except Exception:
                pass

        # Last resort: If no fallback config exists, don't silently use hardcoded defaults.
        # Config() defines these via argparse defaults — if they're missing, there's nothing to fallback to.
        if not f_model or not f_url:
            logger.warning("⚠️ [Circuit Breaker] No fallback model/URL configured. Cannot retry.")
            return None

        # Circuit Breaker: If fallback is identical to primary, don't waste time retrying
        if f_model == model and f_url == url:
            logger.warning(f"⚠️ [Circuit Breaker] Fallback AI matches Primary AI ({model}). Skipping redundant retry.")
            return None

        query_llm._is_fallback = True
        try:
            logger.warning(f"🔄 Primary AI ({model}) failed. Attempting fallback to {f_model}...")
            return query_llm(
                url=f_url,
                model=f_model,
                prompt=prompt,
                images_b64=images_b64,
                system=system,
                format_json=format_json,
                timeout=timeout,
                temperature=temperature,
                max_tokens=max_tokens,
            )
        finally:
            query_llm._is_fallback = False
        return None


def query_telepathic_llm(
    model: str,
    url: str,
    system_prompt: str,
    user_prompt: str,
    temperature: float = 0.0,
    use_local_edge: bool = False,
    images_b64: Optional[List[str]] = None,
) -> str:
    """
    Routes UI Telepathic requests purely based on textual interpretation of the screen's XML nodes.
    If use_local_edge is manually enabled, routes to localhost:11434.
    Otherwise honors the provided URL and model (e.g. OpenRouter).
    Passes optional Base64 screenshots (images_b64) if ai-vision-navigation is enabled.
    """
    target_url = url
    target_model = model

    if use_local_edge:
        is_already_local = "localhost" in url or "127.0.0.1" in url
        if is_already_local:
            logger.debug(
                f"⚡ [Edge Inference] Primary model {model} is already local. Using it directly to prevent VRAM thrashing."
            )
        else:
            logger.info("⚡ [Edge Inference] Routing telepathic request to local Ollama host (0ms latency target).")
            from GramAddict.core.config import Config

            try:
                args = Config().args
                target_url = getattr(args, "ai_fallback_url")
                target_model = getattr(args, "ai_fallback_model")
            except Exception:
                raise RuntimeError(
                    "Config().args not initialized — cannot resolve fallback AI model. Fail Fast."
                )

    is_local = "localhost" in target_url or "127.0.0.1" in target_url
    calc_timeout = 180 if is_local else 45

    ans = query_llm(
        url=target_url,
        model=target_model,
        prompt=user_prompt,
        images_b64=images_b64,
        system=system_prompt,
        format_json=True,
        timeout=calc_timeout,  # Navigation VLM must fail fast for Cloud, but wait for Local VRAM loads
        temperature=temperature,
        max_tokens=150,  # Hard stop to prevent VLM from endlessly hallucinating UI elements
    )
    if ans and "response" in ans:
        return ans["response"]
    return "{}"
