"""
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 re
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 overrides (user picked models in the UI). Same module-global
# pattern: set by the job runner before the pipeline starts, cleared after.
# Vision applies to image calls, text to no-image calls on OpenRouter (and to
# the local->cloud fallback). The LOCAL endpoint's model name is never taken
# from these overrides - hybrid local keeps LOCAL_TEXT_MODEL.
_vision_model_override: Optional[str] = None
_text_model_override: Optional[str] = None

# Per-job raw request/response dumps (missed/hallucinated-item debugging).
# Set by the job runner to <out_dir>/llm_raw at job start, cleared after.
# Same module-global pattern as the model overrides (single-user tool).
_raw_dump_dir: Optional[str] = None
_seq_lock = threading.Lock()
_call_seq = 0


def set_raw_dump_dir(path: Optional[str]) -> None:
    """Point raw LLM request/response dumps at a directory. None disables."""
    global _raw_dump_dir, _call_seq
    with _seq_lock:
        _raw_dump_dir = path
        _call_seq = 0


def _next_seq() -> int:
    global _call_seq
    with _seq_lock:
        _call_seq += 1
        return _call_seq


def _summarize_parsed(parsed: Any) -> str:
    """
    Compact digest of a parsed response for the job log. List values become
    item counts (e.g. conflicts[3]) so a stage that returned nothing (miss)
    or invented items (hallucination) is visible without opening the raw dump.
    """
    if isinstance(parsed, list):
        return f"list[{len(parsed)}]"
    if not isinstance(parsed, dict):
        return type(parsed).__name__
    parts = []
    for k, v in parsed.items():
        if isinstance(v, list):
            parts.append(f"{k}[{len(v)}]")
        elif isinstance(v, dict):
            parts.append(f"{k}{{{len(v)}}}")
        else:
            s = str(v)
            parts.append(f"{k}={s[:40]!r}{'...' if len(s) > 40 else ''}")
    out = ", ".join(parts)
    return out[:300] + ("..." if len(out) > 300 else "")


def _dump_raw(seq: int, be: Dict[str, Any], usage_stage: str,
              system_prompt: str, user_text: str,
              images_b64: Optional[List[str]], max_tokens: int,
              raw: str, parsed: Any,
              finish_reason: Optional[str] = None) -> None:
    """Write the full request/response for one call to the job's llm_raw dir."""
    if not _raw_dump_dir:
        return
    try:
        os.makedirs(_raw_dump_dir, exist_ok=True)
        safe_stage = re.sub(r"[^A-Za-z0-9_.-]+", "_", usage_stage)[:40]
        safe_model = re.sub(r"[^A-Za-z0-9_.-]+", "_", be["model"])
        payload = {
            "seq": seq,
            "stage": usage_stage,
            "model": be["model"],
            "backend": "local" if be.get("local") else "cloud",
            "max_tokens": max_tokens,
            # base64 image payloads deliberately excluded (multi-MB each);
            # the count + source.pdf in the job dir identify what was sent.
            "n_images": len(images_b64 or []),
            "system_prompt": system_prompt,
            "user_text": user_text,
            "finish_reason": finish_reason,
            "raw_response": raw,
            "parsed": parsed,
        }
        path = os.path.join(_raw_dump_dir, f"{seq:04d}_{safe_stage}_{safe_model}.json")
        tmp = f"{path}.{threading.get_ident()}.tmp"
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(payload, f, indent=2)
        os.replace(tmp, path)        # atomic so thread-pooled stages can't tear it
    except OSError as e:
        print(f"[LLM] raw dump failed: {e}")


def _log_call(seq: int, be: Dict[str, Any], usage_stage: str,
              user_text: str, images_b64: Optional[List[str]],
              raw: str, parsed: Any, usd: Optional[float],
              cached: bool = False,
              reasoning_tokens: Optional[int] = None) -> None:
    """One verbose per-call line for the job log (tee'd by job_log.py)."""
    if not config.LLM_VERBOSE:
        return
    backend = "local" if be.get("local") else "cloud"
    if cached:
        cost_str = "cache hit"
    elif usd is not None:
        cost_str = f"${usd:.4f}"
    else:
        cost_str = "cost n/a"
    think_str = f"think {reasoning_tokens}tk | " if reasoning_tokens is not None else ""
    print(f"[LLM] #{seq:04d} {usage_stage} | {be['model']} ({backend}) | "
          f"in {len(user_text)}ch+{len(images_b64 or [])}img | "
          f"out {len(raw)}ch | {think_str}{cost_str} | {_summarize_parsed(parsed)}")


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)


def set_model_overrides(vision: Optional[str] = None, text: Optional[str] = None) -> None:
    """Per-run OpenRouter vision/text model picks. None/blank clears to defaults."""
    global _vision_model_override, _text_model_override
    _vision_model_override = (vision or "").strip() or None
    _text_model_override = (text or "").strip() or None

# --- 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, reasoning_effort: Optional[str] = None,
               reasoning_max_tokens: Optional[int] = None) -> str:
    h = hashlib.sha256()
    parts = [model, str(max_tokens), str(json_mode), str(reasoning_effort),
             str(reasoning_max_tokens), 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,
            # Local model name comes from per-call args or LOCAL_TEXT_MODEL —
            # never the UI's OpenRouter picks, which a local server won't serve.
            "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.
    if has_images:
        model = _vision_model_override or model_override or config.MODEL
    else:
        model = _text_model_override or model_override or config.TEXT_MODEL
    return {
        "base_url": config.AI_BASE_URL,
        "api_key": config.AI_API_KEY,
        "model": 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 _reasoning_tokens(response) -> Optional[int]:
    """Hidden thinking tokens for this call (OpenRouter usage details).

    Direct evidence of whether the reasoning knob is working - without it,
    thinking burn can only be inferred from char counts vs the token cap."""
    try:
        dump = response.model_dump()
    except Exception:
        return None
    usage = dump.get("usage") or {}
    details = usage.get("completion_tokens_details") or {}
    n = details.get("reasoning_tokens")
    return int(n) if isinstance(n, (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 = "?",
    reasoning_effort: Optional[str] = None,
    reasoning_max_tokens: Optional[int] = None,
) -> 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. reasoning_effort: optional OpenRouter reasoning knob ("low"/"medium"/
    "high") - keeps thinking models from burning the output budget on hidden
    reasoning. reasoning_max_tokens: optional hard thinking-token budget
    (OpenRouter reasoning max_tokens -> Gemini thinking_budget); stronger than
    effort, and takes precedence when both are given. Returns the parsed dict,
    or None on a hard failure (caller degrades).
    """
    has_images = bool(images_b64)
    be = _resolve_backend(has_images, model)
    seq = _next_seq()

    cache_key = None
    if config.LLM_CACHE:
        cache_key = _cache_key(be["model"], system_prompt, user_text,
                               images_b64, max_tokens, json_mode=True,
                               reasoning_effort=reasoning_effort,
                               reasoning_max_tokens=reasoning_max_tokens)
        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
                )
            _log_call(seq, be, usage_stage, user_text, images_b64,
                      "", hit, None, cached=True)
            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[str, Any] = dict(model=be["model"], messages=messages,
                                          max_tokens=max_tokens, timeout=config.LLM_TIMEOUT)
            extra_body: Dict[str, Any] = {}
            if be["usage"]:
                extra_body["usage"] = {"include": True}
            # OpenRouter reasoning knob; only sent to cloud backends (local
            # servers reject unknown fields). A hard thinking budget wins over
            # the vaguer effort tier - OpenRouter treats them as exclusive.
            if not be.get("local"):
                if reasoning_max_tokens:
                    extra_body["reasoning"] = {"max_tokens": int(reasoning_max_tokens)}
                elif reasoning_effort:
                    extra_body["reasoning"] = {"effort": reasoning_effort}
            if extra_body:
                kwargs["extra_body"] = extra_body
            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
                )
            choice = response.choices[0]
            finish_reason = getattr(choice, "finish_reason", None)
            raw = _strip_fences(choice.message.content or "")
            reasoning_tokens = _reasoning_tokens(response)
            if finish_reason == "length":
                # Hit max_tokens (thinking tokens included on reasoning models).
                # Logged explicitly so silent truncation isn't mistaken for a
                # parse problem; _parse below still salvages what it can.
                print(f"[LLM] output hit max_tokens (finish_reason=length, "
                      f"{len(raw)}ch returned, "
                      f"thinking={reasoning_tokens}tk)")
            parsed = _parse(raw)
            if parsed is not None:
                _record_model(be, has_images, fell_back)
                if cache_key:
                    _cache_set(cache_key, parsed)
                _log_call(seq, be, usage_stage, user_text, images_b64,
                          raw, parsed, usd, reasoning_tokens=reasoning_tokens)
                if config.LLM_RAW_DUMP:
                    _dump_raw(seq, be, usage_stage, system_prompt, user_text,
                              images_b64, max_tokens, raw, parsed,
                              finish_reason=finish_reason)
                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": _text_model_override or 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
