"""Job handlers (registered with `@handler`): LLM extraction of stored snapshots, embeddings, reprocessing, quality recompute."""
from __future__ import annotations

import logging
from datetime import UTC
from typing import Any

from aiatlas.db import execute, fetch_one, transaction
from aiatlas.schemas import schema_for
from aiatlas.sdk import archive
from aiatlas.sdk.facts import EntityRef, Facts
from aiatlas.sdk.writer import FactWriter
from aiatlas.services.jobs import handler
from aiatlas.services.llm import LLMUnavailable, gateway

log = logging.getLogger(__name__)


@handler("llm_extract")
async def llm_extract(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:
    """Stage 2–4: run the local LLM on a stored snapshot's cleaned text and write facts with extractor='llm'."""
    if not gateway.available:
        raise LLMUnavailable("LLM not configured; leaving snapshot llm_pending")
    snapshot_id = payload["snapshot_id"]
    async with transaction() as conn:
        snap = await fetch_one(conn, """select s.*, d.doc_type, d.entity_id as doc_entity_id, d.connector_name, d.source_id, e.entity_type, e.canonical_name
                                        from snapshots s join documents d on d.id = s.document_id left join entities e on e.id = d.entity_id where s.id = :id""", id=snapshot_id)
    if not snap or not snap["text_path"]:
        return {"skipped": "no text"}
    text = archive.load_text(snap["text_path"])
    if len(text) < 200:
        async with transaction() as conn:
            await execute(conn, "update snapshots set processing_status = 'llm_done' where id = :id", id=snapshot_id)
        return {"skipped": "too short"}
    task = payload.get("task") or "auto"
    if task == "auto":
        task = _guess_task(snap["doc_type"], snap["entity_type"])
    if task == "classify_then_extract":
        label = await gateway.classify(text=text, labels=["model_release", "pricing", "research_paper", "company_news", "framework_release", "hardware", "other"],
                                       snapshot_id=snapshot_id)
        task = {"model_release": "release_announcement", "pricing": "pricing", "research_paper": "paper_passport", "company_news": "release_announcement",
                "framework_release": "release_announcement", "hardware": "hardware_spec"}.get(label or "", "")
        if not task:
            async with transaction() as conn:
                await execute(conn, "update snapshots set processing_status = 'llm_done' where id = :id", id=snapshot_id)
            return {"classified": label, "extracted": False}
    schema, stage = schema_for(task)
    res = await gateway.extract(task_type=task, document=text, schema=schema, stage=stage, snapshot_id=snapshot_id, entity_id=snap["doc_entity_id"], job_id=job["id"])
    if not res.ok or not res.data:
        async with transaction() as conn:
            await execute(conn, "update snapshots set processing_status = 'failed' where id = :id", id=snapshot_id)
        raise RuntimeError(f"llm extraction failed: {res.error}")
    facts = facts_from_llm(task, res.data, entity_id=snap["doc_entity_id"], entity_type=snap["entity_type"], entity_name=snap["canonical_name"], url=snap["url"])
    async with transaction() as conn:
        tier = await fetch_one(conn, "select tier from sources where id = :id", id=snap["source_id"]) if snap["source_id"] else None
        # LLM output never outranks a deterministic statement from the same source: one tier lower (disagreement → conflicting + review)
        writer = FactWriter(conn, source_id=snap["source_id"], snapshot_id=snapshot_id, source_url=snap["url"], tier=min(4, (tier["tier"] if tier else 2) + 1),
                            connector_name=snap["connector_name"], extractor="llm", extractor_version=res.model, observed_at=snap["observed_at"].astimezone(UTC),
                            run_id=snap["run_id"] or job["id"])
        ws = await writer.write(facts)
        await execute(conn, "update snapshots set processing_status = 'llm_done' where id = :id", id=snapshot_id)
    return {"task": task, "model": res.model, **ws.as_dict()}


def _guess_task(doc_type: str | None, entity_type: str | None) -> str:
    if doc_type in ("model_card", "model_page") or entity_type == "model":
        return "model_passport"
    if doc_type == "pricing":
        return "pricing"
    if doc_type in ("paper", "pdf"):
        return "paper_passport"
    if doc_type == "leaderboard":
        return "benchmark_results"
    if entity_type in ("company", "organization"):
        return "company_passport"
    if entity_type == "hardware" or doc_type == "hardware":
        return "hardware_spec"
    if doc_type in ("release", "feed_item", "news"):
        return "release_announcement"
    return "classify_then_extract"


GENERIC_MODEL_NAMES = {"claude", "gpt", "gemini", "llama", "mistral", "qwen", "deepseek", "grok", "gemma", "phi", "command", "codex", "sora", "veo", "imagen",
                       "model", "models", "the model", "new model", "ai model", "llm", "openai", "anthropic", "google", "meta", "microsoft", "nvidia"}
STATUS_ALIASES = {"available": "active", "ga": "active", "generally available": "active", "live": "active", "released": "active", "beta": "preview",
                  "experimental": "preview", "coming soon": "announced", "sunset": "retired", "discontinued": "retired", "legacy": "deprecated"}


def plausible_model_name(name: str | None) -> bool:
    """Reject family/vendor names and vague phrases the LLM sometimes returns as a model name."""
    if not name:
        return False
    n = name.strip()
    low = n.lower()
    if low in GENERIC_MODEL_NAMES or len(n) < 3 or len(n) > 80:
        return False
    return any(ch.isdigit() for ch in n) or len(n.split()) >= 2 or "-" in n


def _norm_status(v: Any) -> Any:
    return STATUS_ALIASES.get(v.strip().lower(), v.strip().lower()) if isinstance(v, str) else v


def facts_from_llm(task: str, data: dict[str, Any], *, entity_id: str | None, entity_type: str | None, entity_name: str | None, url: str) -> Facts:
    """Map a validated LLM output onto facts. LLM claims default to 'medium' confidence and are written one tier below their source,
    so they never outrank deterministic statements (the writer stores disagreements as conflicting)."""
    facts = Facts()
    conf = "medium"
    if "status" in data:
        data = {**data, "status": _norm_status(data.get("status"))}

    def model_ref(name: str, developer: str | None = None) -> EntityRef:
        org = facts.entity("company", developer) if developer else None
        return facts.entity("model", name, organization=org)

    if task == "model_passport":
        name = data.get("name") or entity_name
        if not name or (not (entity_id and entity_type == "model") and not plausible_model_name(name)):
            return facts
        ref = EntityRef(entity_type="model", name=name, id=entity_id) if entity_id and entity_type == "model" else model_ref(name, data.get("developer"))
        if ref not in facts.entities:
            facts.entities.append(ref)
        for prop in ("family", "version", "release_date", "status", "openness", "license", "architecture", "parameter_count", "active_parameter_count", "is_moe",
                     "context_length", "max_output_tokens", "knowledge_cutoff", "tool_calling", "structured_output", "reasoning", "vision", "audio",
                     "fine_tuning_available", "tokenizer", "hardware_requirements", "safety_notes", "training_data_notes"):
            facts.claim(ref, prop, data.get(prop), confidence=conf)
        mods = sorted(set((data.get("modalities_input") or []) + (data.get("modalities_output") or [])))
        facts.claim(ref, "modalities", mods, confidence=conf)
        facts.claim(ref, "modalities_input", data.get("modalities_input"), confidence=conf)
        facts.claim(ref, "modalities_output", data.get("modalities_output"), confidence=conf)
        facts.claim(ref, "languages", data.get("languages"), confidence=conf)
        for key, prop in (("paper_url", "paper_url"), ("model_card_url", "model_card_url"), ("repository_url", "repository_url"), ("official_page_url", "official_url")):
            facts.claim(ref, prop, data.get(key), confidence=conf)
        if data.get("developer") and not (entity_id and entity_type == "model"):
            pass
        elif data.get("developer"):
            org = facts.entity("company", data["developer"])
            facts.relate(org, "develops", ref, confidence=conf)
        if data.get("base_model"):
            base = facts.entity("model", data["base_model"])
            facts.relate(ref, "derived_from", base, confidence="low")
        if data.get("predecessor"):
            pred = facts.entity("model", data["predecessor"])
            facts.relate(pred, "superseded_by", ref, confidence="low")
        facts.document_entity = ref
    elif task == "pricing":
        provider_name = data.get("provider")
        if not provider_name:
            return facts
        provider = facts.entity("provider", provider_name)
        for line in data.get("prices") or []:
            if not line.get("model"):
                continue
            model = facts.entity("model", line["model"])
            facts.price(model=model, provider=provider, provider_model_id=line.get("provider_model_id"), input_per_mtok=line.get("input_per_mtok"),
                        output_per_mtok=line.get("output_per_mtok"), cached_input_per_mtok=line.get("cached_input_per_mtok"),
                        cache_write_per_mtok=line.get("cache_write_per_mtok"), batch_input_per_mtok=line.get("batch_input_per_mtok"),
                        batch_output_per_mtok=line.get("batch_output_per_mtok"), per_image=line.get("per_image"), currency=data.get("currency") or "USD",
                        context_length=line.get("context_length"), max_output_tokens=line.get("max_output_tokens"), meta={"extractor": "llm", "notes": line.get("notes")})
    elif task == "company_passport":
        name = data.get("name") or entity_name
        if not name:
            return facts
        ref = EntityRef(entity_type="company", name=name, id=entity_id) if entity_id and entity_type in ("company", "organization") else facts.entity("company", name)
        if ref not in facts.entities:
            facts.entities.append(ref)
        for prop in ("legal_name", "country", "headquarters", "founded", "founders", "leadership", "website", "employee_count", "funding_total_usd"):
            facts.claim(ref, prop, data.get(prop), confidence=conf)
        if data.get("description"):
            facts.claim(ref, "description", data["description"], confidence=conf)
        for m in data.get("models") or []:
            if plausible_model_name(m):
                facts.relate(ref, "develops", facts.entity("model", m, organization=ref), confidence="low")
        for inv in data.get("investors") or []:
            facts.relate(ref, "funded_by", facts.entity("company", inv), confidence="low")
        if data.get("parent_company"):
            facts.relate(facts.entity("company", data["parent_company"]), "owns", ref, confidence="low")
        facts.document_entity = ref
    elif task == "paper_passport":
        title = data.get("title") or entity_name
        if not title:
            return facts
        ref = EntityRef(entity_type="paper", name=title, id=entity_id) if entity_id and entity_type == "paper" else facts.entity("paper", title)
        if ref not in facts.entities:
            facts.entities.append(ref)
        for prop in ("authors", "affiliations", "date", "field", "summary", "methods", "key_claims", "results", "limitations", "code_url"):
            facts.claim(ref, prop, data.get(prop), confidence=conf)
        for m in data.get("models") or []:
            if plausible_model_name(m):
                facts.relate(facts.entity("model", m), "described_by", ref, confidence="low")
        for d in data.get("datasets") or []:
            facts.relate(ref, "uses_dataset", facts.entity("dataset", d), confidence="low")
        for b in data.get("benchmarks") or []:
            facts.relate(ref, "evaluates_on", facts.entity("benchmark", b), confidence="low")
        facts.document_entity = ref
    elif task == "benchmark_results":
        bname = data.get("benchmark") or entity_name
        if not bname:
            return facts
        bench = facts.entity("benchmark", bname)
        for row in data.get("rows") or []:
            if row.get("model") and row.get("score") is not None:
                facts.result(model=facts.entity("model", row["model"]), benchmark=bench, score=float(row["score"]), metric=row.get("metric") or data.get("metric"),
                             higher_is_better=bool(data.get("higher_is_better", True)), config={"config": row.get("config"), "extractor": "llm"}, confidence="low")
    elif task == "hardware_spec":
        name = data.get("name") or entity_name
        if not name:
            return facts
        org = facts.entity("company", data["manufacturer"]) if data.get("manufacturer") else None
        ref = EntityRef(entity_type="hardware", name=name, id=entity_id) if entity_id and entity_type == "hardware" else facts.entity("hardware", name, organization=org)
        if ref not in facts.entities:
            facts.entities.append(ref)
        for prop in ("kind", "architecture", "release_date", "memory_gb", "memory_type", "memory_bandwidth_gbs", "compute_fp16_tflops", "compute_fp8_tflops",
                     "compute_int8_tops", "tdp_watts", "form_factor", "price_usd", "interconnect"):
            facts.claim(ref, prop, data.get(prop), confidence=conf)
        if org:
            facts.relate(org, "manufactures", ref, confidence=conf)
        facts.document_entity = ref
    elif task == "release_announcement":
        org = facts.entity("company", data["organization"]) if data.get("organization") else None
        for m in data.get("models") or []:
            if not plausible_model_name(m):
                continue
            ref = facts.entity("model", m, organization=org)
            if org:
                facts.relate(org, "develops", ref, confidence="low")
        if data.get("model_passport"):
            facts.extend(facts_from_llm("model_passport", data["model_passport"], entity_id=None, entity_type=None, entity_name=None, url=url))
        if data.get("pricing"):
            facts.extend(facts_from_llm("pricing", data["pricing"], entity_id=None, entity_type=None, entity_name=None, url=url))
    return facts


@handler("embed_entity")
async def embed_entity(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:
    from aiatlas.services.embeddings import embed_entities

    return await embed_entities(payload.get("entity_ids") or [payload["entity_id"]])


@handler("reprocess_snapshot")
async def reprocess_snapshot(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:
    from aiatlas.connectors import get

    connector = get(payload["connector"])
    ctx = await connector.run(reprocess=True, only_urls=[payload["url"]] if payload.get("url") else None, force=True)
    return {k: v for k, v in ctx.stats.__dict__.items() if k != "meta"}


@handler("recompute_quality")
async def recompute_quality(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:
    from aiatlas.services.quality import recompute

    return await recompute(entity_ids=payload.get("entity_ids"))


__all__ = ["embed_entity", "facts_from_llm", "llm_extract", "recompute_quality", "reprocess_snapshot"]