"""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.sdk import archive
from aiatlas.sdk.facts import EntityRef, Facts
from aiatlas.sdk.writer import FactWriter
from aiatlas.schemas import schema_for
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
        writer = FactWriter(conn, source_id=snap["source_id"], snapshot_id=snapshot_id, source_url=snap["url"], tier=(tier["tier"] if tier else 2),
                            connector_name=snap["connector_name"], extractor="llm", extractor_version=res.model, observed_at=snap["observed_at"].astimezone(UTC))
        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"


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 never outrank tier-1 deterministic ones
    (the writer stores disagreements as conflicting)."""
    facts = Facts()
    conf = "medium"

    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:
            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 []:
            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 []:
            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 []:
            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__ = ["llm_extract", "facts_from_llm", "embed_entity", "reprocess_snapshot", "recompute_quality"]