"""Shared helpers for agentic pipelines (agentic_v2.py, agentic_lev.py).

Keeps provider abstraction, prompt hashing, _meta provenance, PDF loading,
and canonical output writing in one place so the two pipelines stay aligned
with the single-call baseline in scripts/experiments/model_comparison/.
"""

from __future__ import annotations

import hashlib
import json
import os
import re
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from config import SYSTEM_PROMPT  # noqa: E402
from validation.validate_extraction import CANONICAL_FIELDS, validate_extraction  # noqa: E402

# ── Provenance constants (computed once at import) ─────────────────
PROMPT_HASH = hashlib.sha256(SYSTEM_PROMPT.encode()).hexdigest()[:16]
try:
    GIT_COMMIT = subprocess.check_output(
        ["git", "rev-parse", "--short", "HEAD"],
        stderr=subprocess.DEVNULL,
        cwd=str(ROOT),
    ).decode().strip()
except Exception:
    GIT_COMMIT = "unknown"


# ── Backbone catalogue ─────────────────────────────────────────────
# Each entry is self-contained so the calling script only needs the key.
# `output_slug` is used to build data/extractions/agentic_{v2,lev}_<slug>/.

MODELS: dict[str, dict[str, Any]] = {
    "sonnet-4-5": {
        "name": "Claude Sonnet 4.5 (gold)",
        "model_id": "claude-sonnet-4-5-20250929",
        "provider": "anthropic",
        "output_slug": "sonnet_4_5",
        "input_price_per_mtok": 3.0,
        "output_price_per_mtok": 15.0,
    },
    "gpt-5.4": {
        "name": "GPT-5.4 full",
        "model_id": "gpt-5.4-2026-03-05",
        "provider": "openai",
        "output_slug": "gpt5_4_full",
        "max_tokens_param": "max_completion_tokens",
        "input_price_per_mtok": 1.25,
        "output_price_per_mtok": 10.0,
    },
    "gpt-5.4-mini": {
        "name": "GPT-5.4 mini",
        "model_id": "gpt-5.4-mini-2026-03-17",
        "provider": "openai",
        "output_slug": "gpt5_4_mini",
        "max_tokens_param": "max_completion_tokens",
        "input_price_per_mtok": 0.15,
        "output_price_per_mtok": 0.60,
    },
    "sonnet-4-6": {
        "name": "Claude Sonnet 4.6",
        "model_id": "claude-sonnet-4-6",
        "provider": "anthropic",
        "output_slug": "sonnet_4_6",
        "input_price_per_mtok": 3.0,
        "output_price_per_mtok": 15.0,
    },
    "gemini-3.1-pro": {
        "name": "Gemini 3.1 Pro Preview",
        "model_id": "gemini-3.1-pro-preview",
        "provider": "google",
        "output_slug": "gemini_3_1_pro",
        "input_price_per_mtok": 1.25,
        "output_price_per_mtok": 10.0,
    },
    "llama-4-scout": {
        "name": "Llama 4 Scout 17B-16E (vLLM)",
        "model_id": "meta-llama/Llama-4-Scout-17B-16E-Instruct",
        "provider": "local_vllm",
        "output_slug": "llama_4_scout",
        "base_url": os.getenv("VLLM_BASE_URL", "http://localhost:8000/v1"),
        "api_key": os.getenv("VLLM_API_KEY", "EMPTY"),
        "input_price_per_mtok": 0.0,
        "output_price_per_mtok": 0.0,
    },
}


def resolve_backbone(key: str) -> dict[str, Any]:
    if key not in MODELS:
        raise SystemExit(
            f"Unknown backbone '{key}'. Available: {sorted(MODELS)}"
        )
    return MODELS[key]


# ── Text / JSON utilities ─────────────────────────────────────────


def sanitize_utf8(text: str) -> str:
    """Drop lone UTF-16 surrogates PyPDF2 occasionally emits."""
    return text.encode("utf-8", errors="replace").decode("utf-8")


def parse_json_response(raw_text: str) -> dict:
    text = raw_text.strip()
    if "```json" in text:
        text = text.split("```json", 1)[1].split("```", 1)[0].strip()
    elif text.startswith("```"):
        text = text[3:]
        if text.endswith("```"):
            text = text[:-3]
    first = text.find("{")
    last = text.rfind("}")
    if first != -1 and last != -1:
        text = text[first:last + 1]
    # Strip unescaped ASCII control characters (JSON spec disallows them
    # inside strings). Models occasionally emit raw \x00-\x1f bytes,
    # which json.loads rejects. Whitelist \t, \n, \r since those are
    # valid JSON whitespace at structural positions.
    text = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f]", "", text)
    return json.loads(text.strip())


_PREFIX_MAP = {
    "sc:name": "name", "sc:description": "description", "sc:url": "url",
    "sc:license": "license", "sc:creator": "creator", "sc:publisher": "publisher",
    "sc:datePublished": "datePublished", "sc:inLanguage": "inLanguage",
    "cr:citeAs": "citeAs", "cr:isLiveDataset": "isLiveDataset",
    "sc:citeAs": "citeAs", "sc:isLiveDataset": "isLiveDataset",
    "cr:name": "name", "cr:description": "description",
}


def normalize_field_names(metadata: dict) -> dict:
    """Strip sc:/cr: prefixes the RAI fields don't carry."""
    return {_PREFIX_MAP.get(k, k): v for k, v in metadata.items()}


def fill_canonical(metadata: dict) -> dict:
    """Ensure all 30 canonical keys exist (null where missing)."""
    for field in CANONICAL_FIELDS:
        metadata.setdefault(field, None)
    return metadata


# ── PDF loading ───────────────────────────────────────────────────


def get_paper_text(ds_id: str) -> str | None:
    """Load + clean + sanitize paper text. Mirrors extract_all_models.py."""
    from croissantminer.pdf.reader import extract_text_from_pdf as _extract
    from croissantminer.pdf.processor import clean_text as _clean

    raw_dir = ROOT / "data" / "raw"
    candidates = [
        raw_dir / f"{ds_id}.pdf",
        raw_dir / f"{ds_id.replace('_', '/')}.pdf",
    ]
    links_path = ROOT / "data" / "paper_links.json"
    if links_path.exists():
        with open(links_path) as f:
            links = json.load(f)
        match = re.search(r"(\d{4}\.\d{4,5})", links.get(ds_id, ""))
        if match:
            arxiv_id = match.group(1)
            for suffix in ("", "v1", "v2", "v3", "v4", "v5"):
                candidates.append(raw_dir / f"{arxiv_id}{suffix}.pdf")

    for path in candidates:
        if path.exists():
            return sanitize_utf8(_clean(_extract(str(path))))
    return None


def strip_references(text: str) -> str:
    """Cut off the references section (anywhere past the paper midpoint)."""
    patterns = (
        r"\n\s*References\s*\n",
        r"\n\s*REFERENCES\s*\n",
        r"\n\s*Bibliography\s*\n",
    )
    for pattern in patterns:
        match = re.search(pattern, text)
        if match and match.start() > len(text) * 0.5:
            return text[: match.start()]
    return text


# ── Provider-agnostic LLM call ────────────────────────────────────


def _call_anthropic(cfg, system_prompt, user_content, max_tokens):
    import anthropic

    client = anthropic.Anthropic()
    kwargs = {
        "model": cfg["model_id"],
        "max_tokens": max_tokens,
        "system": system_prompt,
        "messages": [{"role": "user", "content": user_content}],
    }
    if not cfg.get("skip_temperature"):
        kwargs["temperature"] = 0.0

    response = client.messages.create(**kwargs)
    raw = ""
    for block in response.content:
        if getattr(block, "type", None) == "text":
            raw += block.text
        elif hasattr(block, "text"):
            raw += block.text
    usage = {
        "input_tokens": response.usage.input_tokens,
        "output_tokens": response.usage.output_tokens,
    }
    return raw, usage


def _call_openai(cfg, system_prompt, user_content, max_tokens, base_url=None, api_key=None):
    from openai import OpenAI

    kwargs = {}
    if base_url:
        kwargs["base_url"] = base_url
    if api_key:
        kwargs["api_key"] = api_key
    client = OpenAI(**kwargs)

    params = {
        "model": cfg["model_id"],
        "temperature": 0.0,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_content},
        ],
    }
    params[cfg.get("max_tokens_param", "max_completion_tokens")] = max_tokens

    response = client.chat.completions.create(**params)
    raw = response.choices[0].message.content or ""
    usage = {
        "input_tokens": response.usage.prompt_tokens,
        "output_tokens": response.usage.completion_tokens,
    }
    return raw, usage


_OPENROUTER_MODEL_MAP = {
    "gemini-3.1-pro-preview": "google/gemini-3.1-pro-preview",
    "gemini-2.5-pro": "google/gemini-2.5-pro",
}


def _call_via_openrouter(cfg, system_prompt, user_content, max_tokens):
    """OpenAI-compatible call against OpenRouter, used to bypass Google's
    250-RPD AI-Studio cap on Gemini 3.1 Pro preview. Same model, different
    routing pool. Pricier per token but no daily request cap."""
    from openai import OpenAI

    api_key = os.environ["OPENROUTER_API_KEY"]
    model_id = _OPENROUTER_MODEL_MAP.get(cfg["model_id"], cfg["model_id"])
    client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=api_key)
    response = client.chat.completions.create(
        model=model_id,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_content},
        ],
        temperature=0.0,
        max_tokens=max_tokens,
    )
    raw = response.choices[0].message.content or ""
    usage = {
        "input_tokens": response.usage.prompt_tokens,
        "output_tokens": response.usage.completion_tokens,
    }
    return raw, usage


def _call_google(cfg, system_prompt, user_content, max_tokens):
    import requests

    # 2026-05-28: route via OpenRouter when OPENROUTER_API_KEY is set, to
    # bypass Google AI Studio's 250 RPD cap on gemini-3.1-pro. Same model,
    # OpenAI-compatible endpoint; ~60% pricier per token than direct.
    if os.getenv("OPENROUTER_API_KEY"):
        return _call_via_openrouter(cfg, system_prompt, user_content, max_tokens)

    api_key = os.getenv("GEMINI_API_KEY")
    url = f"https://generativelanguage.googleapis.com/v1beta/models/{cfg['model_id']}:generateContent"
    gen_config = {"temperature": 0.0, "maxOutputTokens": max_tokens}
    # Gemini 3.x burns most output budget on internal reasoning before
    # emitting JSON, truncating extraction mid-field. Set thinkingLevel
    # to "low" so the budget goes to the actual JSON. Gemini 2.5 (and
    # earlier) reject thinkingLevel, so apply only to 3.x.
    if cfg.get("model_id", "").startswith("gemini-3"):
        gen_config["thinkingConfig"] = {"thinkingLevel": "low"}
    payload = {
        "contents": [{"parts": [{"text": user_content}]}],
        "systemInstruction": {"parts": [{"text": system_prompt}]},
        "generationConfig": gen_config,
    }
    resp = requests.post(url, params={"key": api_key}, json=payload, timeout=180)
    if resp.status_code != 200:
        raise RuntimeError(f"Gemini API error {resp.status_code}: {resp.text[:300]}")
    data = resp.json()
    try:
        raw = data["candidates"][0]["content"]["parts"][0]["text"]
    except (KeyError, IndexError):
        raise RuntimeError(f"Unexpected Gemini response: {json.dumps(data)[:300]}")
    usage_meta = data.get("usageMetadata", {})
    usage = {
        "input_tokens": usage_meta.get("promptTokenCount", 0),
        "output_tokens": usage_meta.get("candidatesTokenCount", 0),
    }
    return raw, usage


def call_llm(
    cfg: dict,
    system_prompt: str,
    user_content: str,
    max_tokens: int = 8192,
    retries: int = 3,
) -> tuple[str, dict]:
    """Provider-agnostic call with retry on rate limits / 5xx.

    Returns (raw_text, usage). Parsing is left to the caller so each phase
    can handle its own schema.
    """
    provider = cfg["provider"]
    backoffs = [5, 15, 45]
    last_exc: Exception | None = None

    for attempt in range(retries):
        try:
            if provider == "anthropic":
                return _call_anthropic(cfg, system_prompt, user_content, max_tokens)
            if provider == "openai":
                return _call_openai(cfg, system_prompt, user_content, max_tokens)
            if provider == "google":
                return _call_google(cfg, system_prompt, user_content, max_tokens)
            if provider == "local_vllm":
                return _call_openai(
                    cfg, system_prompt, user_content, max_tokens,
                    base_url=cfg.get("base_url"),
                    api_key=cfg.get("api_key") or "EMPTY",
                )
            raise ValueError(f"Unknown provider: {provider}")
        except Exception as exc:  # noqa: BLE001 — retry on anything transient
            last_exc = exc
            msg = str(exc).lower()
            retryable = (
                "rate" in msg or "429" in msg or "timeout" in msg
                or str(getattr(exc, "status_code", "")).startswith("5") or "overloaded" in msg
            )
            if attempt == retries - 1 or not retryable:
                raise
            wait = backoffs[attempt] if attempt < len(backoffs) else 60
            time.sleep(wait)

    raise last_exc  # unreachable but keeps type-checkers quiet


def estimate_cost(cfg: dict, usage: dict, batch_discount: bool = False) -> float:
    mult = 0.5 if batch_discount else 1.0
    inp = usage.get("input_tokens", 0) / 1e6 * cfg["input_price_per_mtok"] * mult
    out = usage.get("output_tokens", 0) / 1e6 * cfg["output_price_per_mtok"] * mult
    return round(inp + out, 4)


# ── Canonical output ──────────────────────────────────────────────


def meta_block(
    cfg: dict,
    script: str,
    mode: str,
    paper_set: str = "102_dev_test_split",
    extra: dict | None = None,
) -> dict:
    block = {
        "parser": "pypdf2",
        "parser_version": "3.0.1",
        "prompt_sha256_prefix": PROMPT_HASH,
        "prompt_chars": len(SYSTEM_PROMPT),
        "model_id": cfg["model_id"],
        "provider": cfg["provider"],
        "temperature": None if cfg.get("skip_temperature") else 0.0,
        "timestamp": datetime.utcnow().isoformat() + "Z",
        "git_commit": GIT_COMMIT,
        "paper_set": paper_set,
        "script": script,
        "mode": mode,
    }
    if extra:
        block.update(extra)
    return block


def write_canonical_output(
    out_dir: Path,
    ds_id: str,
    extraction: dict,
    usage: dict,
    cfg: dict,
    meta: dict,
) -> Path:
    """Write one paper's extraction in the canonical {dataset_id, model,
    extraction, usage, valid, _meta} schema. Returns the output path."""
    out_dir.mkdir(parents=True, exist_ok=True)
    safe_id = ds_id.replace("/", "__")
    path = out_dir / f"{safe_id}.json"

    extraction = normalize_field_names(extraction)
    extraction = fill_canonical(extraction)
    is_valid, errors = validate_extraction(extraction, ds_id)

    payload = {
        "dataset_id": ds_id,
        "model": cfg["model_id"],
        "extraction": extraction,
        "usage": usage,
        "valid": is_valid,
        "_meta": meta,
    }
    if not is_valid:
        payload["validation_errors"] = errors

    with open(path, "w") as f:
        json.dump(payload, f, indent=2, ensure_ascii=False)
    return path


def output_dir_for(pipeline: str, backbone_key: str) -> Path:
    """Canonical output directory per pipeline/backbone pair."""
    cfg = resolve_backbone(backbone_key)
    return ROOT / "data" / "extractions" / f"agentic_{pipeline}_{cfg['output_slug']}"