"""
llm.py - Thin OpenRouter client + a JSON-returning chat helper.

Mirrors the call pattern proven in the IronBid pipeline: OpenAI SDK pointed at
OpenRouter, base64 image data URLs, markdown-fence stripping, and a single
retry on JSON parse failure. One helper serves both the extraction (Stage 1,
with an image) and conflict-reasoning (Stage 3, with images) calls.
"""

import os
import json
import hashlib
import threading
from typing import List, Optional, Dict, Any

from openai import OpenAI

from backend import config

_clients: Dict[str, OpenAI] = {}

# Per-run text backend choice (hybrid). Set by the runner at job start; read by
# call_json when routing a no-image (text) call. Module-global mirrors the
# set_stage/cost pattern (single-user tool).
_text_local = False

# Per-job model override (user picked a model in the UI). Same module-global
# pattern: set by the job runner before the pipeline starts, cleared after.
_model_override: Optional[str] = None


def set_model_override(model: Optional[str]) -> None:
    """Override the model for all OpenRouter calls (vision + text), or None to clear."""
    global _model_override
    _model_override = (model or "").strip() or None


def set_text_backend(local: bool) -> None:
    """Choose whether text (no-image) calls go to the local endpoint this run."""
    global _text_local
    _text_local = bool(local)

# --- per-job cost accounting -------------------------------------------------
# OpenRouter returns the real USD cost of each call when we request usage
# accounting. We accumulate it in a module-level counter; the runner resets it
# at the start of a job and reads it at the end. (Single-user tool: overlapping
# jobs would share this counter, which is acceptable here.)
_cost_lock = threading.Lock()
_cost = {"usd": 0.0, "calls": 0, "cached": 0}
_by_stage: Dict[str, Dict[str, Any]] = {}
_current_stage = "?"

# Model usage tracking - records which models were actually called this run.
_models: Dict[str, Any] = {
    "vision": set(),       # models used for vision (image) calls
    "text_local": set(),   # local model names attempted
    "text_cloud": set(),   # cloud text model names used
    "fallback_count": 0,   # times local fell back to cloud
}


def set_stage(name: str) -> None:
    """Label subsequent calls so cost can be attributed per pipeline stage."""
    global _current_stage
    _current_stage = name


def reset_cost() -> None:
    with _cost_lock:
        _cost["usd"] = 0.0
        _cost["calls"] = 0
        _cost["cached"] = 0
        _by_stage.clear()
        _models["vision"] = set()
        _models["text_local"] = set()
        _models["text_cloud"] = set()
        _models["fallback_count"] = 0


def get_cost() -> Dict[str, Any]:
    with _cost_lock:
        out = dict(_cost)
        out["by_stage"] = {k: dict(v) for k, v in _by_stage.items()}
        out["models"] = {
            "vision":         sorted(_models["vision"]),
            "text_local":     sorted(_models["text_local"]),
            "text_cloud":     sorted(_models["text_cloud"]),
            "fallback_count": _models["fallback_count"],
        }
        return out


def _record_model(be: Dict[str, Any], has_images: bool, fell_back: bool) -> None:
    """Thread-safe recording of which model was actually used for a call."""
    model = be["model"]
    with _cost_lock:
        if has_images:
            _models["vision"].add(model)
        elif fell_back or not be.get("local"):
            _models["text_cloud"].add(model)
        else:
            _models["text_local"].add(model)


def _stage_bucket() -> Dict[str, Any]:
    return _by_stage.setdefault(_current_stage, {"usd": 0.0, "calls": 0, "cached": 0})


def _add_cost(usd: float) -> None:
    with _cost_lock:
        _cost["usd"] += usd
        _cost["calls"] += 1
        b = _stage_bucket()
        b["usd"] += usd
        b["calls"] += 1


def _add_cached() -> None:
    with _cost_lock:
        _cost["cached"] += 1
        _stage_bucket()["cached"] += 1


# --- disk-backed response cache (testing loop; opt-in via config.LLM_CACHE) ---

def _cache_key(model: str, system_prompt: str, user_text: str,
               images_b64: Optional[List[str]], max_tokens: int,
               json_mode: bool) -> str:
    h = hashlib.sha256()
    parts = [model, str(max_tokens), str(json_mode), system_prompt, user_text]
    for b in (images_b64 or []):
        parts.append(b)
    for p in parts:
        h.update(b"\x00")
        h.update(p.encode("utf-8", "ignore"))
    return h.hexdigest()


def _cache_get(key: str) -> Optional[Dict[str, Any]]:
    path = os.path.join(config.LLM_CACHE_DIR, key + ".json")
    try:
        with open(path, encoding="utf-8") as f:
            return json.load(f)
    except (OSError, json.JSONDecodeError):
        return None


def _cache_set(key: str, value: Dict[str, Any]) -> None:
    try:
        os.makedirs(config.LLM_CACHE_DIR, exist_ok=True)
        path = os.path.join(config.LLM_CACHE_DIR, key + ".json")
        tmp = f"{path}.{os.getpid()}.{threading.get_ident()}.tmp"
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(value, f)
        os.replace(tmp, path)        # atomic so thread-pooled stages can't tear it
    except OSError as e:
        print(f"[LLM] cache write failed: {e}")


def get_client(base_url: str, api_key: str) -> OpenAI:
    """Lazily build + cache one OpenAI-compatible client per endpoint."""
    client = _clients.get(base_url)
    if client is None:
        client = OpenAI(base_url=base_url, api_key=api_key)
        _clients[base_url] = client
    return client


def _resolve_backend(has_images: bool, model_override: Optional[str]) -> Dict[str, Any]:
    """
    Pick (base_url, api_key, model, usage) for this call.

    Vision (has_images) always uses OpenRouter. Text (no images) uses the local
    endpoint when the run chose hybrid AND it's configured; otherwise OpenRouter.
    """
    if not has_images and _text_local and config.LOCAL_BASE_URL:
        return {
            "base_url": config.LOCAL_BASE_URL,
            "api_key": config.LOCAL_API_KEY,
            "model": model_override or config.LOCAL_TEXT_MODEL or config.TEXT_MODEL,
            "usage": False,          # local has no OpenRouter usage accounting
            "local": True,
        }
    # Vision, or text-on-OpenRouter (default / fallback). A per-job override
    # (user's UI model pick) wins over per-call and env defaults.
    default_model = config.MODEL if has_images else config.TEXT_MODEL
    return {
        "base_url": config.AI_BASE_URL,
        "api_key": config.AI_API_KEY,
        "model": _model_override or model_override or default_model,
        "usage": True,
        "local": False,
    }


def _strip_fences(raw: str) -> str:
    raw = raw.strip()
    if raw.startswith("```"):
        # ```json ... ``` or ``` ... ```
        raw = raw.split("```")[1]
        if raw.startswith("json"):
            raw = raw[4:]
        raw = raw.strip()
    return raw


def _repair_truncated(raw: str) -> Optional[Dict[str, Any]]:
    """
    Best-effort recovery of a truncated JSON object/array.

    Walks the text tracking bracket depth and string state, finds the last
    point where we can safely cut (after a complete element), drops the partial
    trailing element, and closes the still-open brackets. Lets a response that
    got cut off by max_tokens still yield its complete leading items instead of
    nothing. Returns None if nothing salvageable.
    """
    stack: List[str] = []
    in_str = esc = False
    cut = None  # (index_exclusive, stack_snapshot) of a safe truncation point
    for i, ch in enumerate(raw):
        if in_str:
            if esc:
                esc = False
            elif ch == "\\":
                esc = True
            elif ch == '"':
                in_str = False
            continue
        if ch == '"':
            in_str = True
        elif ch in "{[":
            stack.append("}" if ch == "{" else "]")
        elif ch in "}]":
            if stack:
                stack.pop()
            cut = (i + 1, list(stack))     # after a complete sub-structure
        elif ch == ",":
            cut = (i, list(stack))         # before an (incomplete) next element

    if not cut:
        return None
    idx, open_brackets = cut
    candidate = raw[:idx].rstrip().rstrip(",")
    candidate += "".join(reversed(open_brackets))
    try:
        return json.loads(candidate)
    except json.JSONDecodeError:
        return None


def _response_cost(response) -> Optional[float]:
    """Pull OpenRouter's per-call USD cost out of the usage object, if present."""
    try:
        dump = response.model_dump()
    except Exception:
        return None
    usage = dump.get("usage") or {}
    cost = usage.get("cost")
    if cost is None:
        cost = (usage.get("cost_details") or {}).get("upstream_inference_cost")
    return float(cost) if isinstance(cost, (int, float)) else None


def _parse(raw: str) -> Optional[Dict[str, Any]]:
    """Parse model JSON, falling back to truncation repair."""
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        repaired = _repair_truncated(raw)
        if repaired is not None:
            print("[LLM] recovered truncated JSON (dropped partial trailing item)")
        return repaired


def call_json(
    system_prompt: str,
    user_text: str,
    images_b64: Optional[List[str]] = None,
    max_tokens: int = 4096,
    model: Optional[str] = None,
    usage_tracker: Optional[Any] = None,
    usage_stage: str = "?",
) -> Optional[Dict[str, Any]]:
    """
    Send one chat completion expecting a JSON object back.

    Uses the provider's JSON mode (response_format) so the model returns a bare
    JSON object instead of prose/empty text, and recovers from max_tokens
    truncation. images_b64: optional base64 JPEGs attached as high-detail image
    parts. Returns the parsed dict, or None on a hard failure (caller degrades).
    """
    has_images = bool(images_b64)
    be = _resolve_backend(has_images, model)

    cache_key = None
    if config.LLM_CACHE:
        cache_key = _cache_key(be["model"], system_prompt, user_text,
                               images_b64, max_tokens, json_mode=True)
        hit = _cache_get(cache_key)
        if hit is not None:
            _add_cached()
            if usage_tracker:
                usage_tracker.record(
                    usage_stage, be["model"], cached=True, has_images=has_images
                )
            return hit

    content: List[Dict[str, Any]] = []
    for b64 in images_b64 or []:
        content.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{b64}", "detail": "high"},
        })
    content.append({"type": "text", "text": user_text})

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": content},
    ]

    use_json_mode = True
    fell_back = False
    for attempt in range(2):
        try:
            client = get_client(be["base_url"], be["api_key"])
            kwargs = dict(model=be["model"], messages=messages,
                          max_tokens=max_tokens, timeout=config.LLM_TIMEOUT)
            if be["usage"]:
                kwargs["extra_body"] = {"usage": {"include": True}}
            if use_json_mode:
                kwargs["response_format"] = {"type": "json_object"}
            response = client.chat.completions.create(**kwargs)
            usd = _response_cost(response)
            if usd is not None:
                _add_cost(usd)
            if usage_tracker:
                usage_tracker.record(
                    usage_stage, be["model"], usd=usd or 0.0, has_images=has_images
                )
            raw = _strip_fences(response.choices[0].message.content or "")
            parsed = _parse(raw)
            if parsed is not None:
                _record_model(be, has_images, fell_back)
                if cache_key:
                    _cache_set(cache_key, parsed)
                return parsed
            if attempt == 0:
                print("[LLM] JSON parse error (retrying)")
                continue
            print(f"[LLM] JSON parse error (giving up); raw head: {raw[:300]}")
            return None
        except Exception as e:
            # Some models reject response_format -> drop it and retry once.
            if use_json_mode and "response_format" in str(e).lower():
                print("[LLM] model rejected JSON mode; retrying without it")
                use_json_mode = False
                continue
            # Local endpoint unreachable -> fall back to OpenRouter text backend once.
            if be["local"] and not fell_back:
                print(f"[LLM] local endpoint failed ({e}); falling back to OpenRouter")
                with _cost_lock:
                    _models["text_local"].add(be["model"])
                    _models["fallback_count"] += 1
                be = {"base_url": config.AI_BASE_URL, "api_key": config.AI_API_KEY,
                      "model": config.TEXT_MODEL, "usage": True, "local": False}
                cache_key = None       # don't cache fallback under the local-model key
                fell_back = True
                continue
            print(f"[LLM] call error: {e}")
            if attempt == 0:
                continue
            return None
    return None