"""Wave 5b: vision fact-check of extracted evidence against cited sheet images."""

from backend import config
from backend.agents.base import AgentResult, AgentScope, AgentUsage, failure
from backend.llm import call_json
from backend.pipeline._serialize import dumps
from backend.pipeline._stage import collect_list, render
from backend.prompts import VERIFY_SYSTEM_PROMPT, VERIFY_USER_INSTRUCTION

_SEVERITY_RANK = {"critical": 0, "high": 1, "medium": 2, "low": 3}
_VERDICTS = ("confirmed", "corrected", "not_found")


def select_findings(findings, clusters, max_checks, severities):
    """Severity-gated selection plus any finding tied to a disputed cluster."""
    disputed_keys = {c.get("key") for c in clusters if c.get("disputed_attributes")}
    selected = [f for f in findings
                if str(f.get("severity") or "").lower() in severities
                or f.get("cluster_key") in disputed_keys]
    selected.sort(key=lambda f: _SEVERITY_RANK.get(
        str(f.get("severity") or "").lower(), 9))
    return selected[:max_checks]


def _valid_verdict(item):
    if not isinstance(item, dict):
        return None
    verdict = str(item.get("verdict") or "").lower()
    if verdict not in _VERDICTS:
        return None
    return {"sheet": item.get("sheet") or "",
            "source_text": item.get("source_text") or "",
            "verdict": verdict,
            "actual_text": item.get("actual_text"),
            "notes": item.get("notes")}


def _status(verdicts):
    if not verdicts:
        return "unverified"
    confirmed = sum(1 for v in verdicts if v["verdict"] == "confirmed")
    if confirmed == len(verdicts):
        return "confirmed"
    if confirmed == 0:
        return "refuted"
    return "mixed"


class EvidenceVerifierAgent:
    name = "verify"

    def __init__(self, usage: AgentUsage) -> None:
        self.usage = usage

    def run(self, scope: AgentScope) -> AgentResult:
        try:
            finding = scope.payload["finding"]
            instruction = render(VERIFY_USER_INSTRUCTION, {
                "finding": dumps(finding),
                "text_layer": scope.payload.get("text_layer_excerpt")
                              or "(no text layer available for the cited sheets)",
            })
            parsed = call_json(
                system_prompt=VERIFY_SYSTEM_PROMPT,
                user_text=instruction,
                images_b64=scope.payload.get("images_b64") or [],
                max_tokens=config.VERIFY_MAX_TOKENS,
                model=config.AGENT_VERIFY_MODEL,
                reasoning_effort=config.AGENT_VERIFY_REASONING_EFFORT or None,
                usage_tracker=self.usage,
                usage_stage="agent.verify",
            )
            verdicts = collect_list(parsed, "verdicts", _valid_verdict)
            return AgentResult(scope_id=scope.scope_id, artifacts=[{
                "finding_index": scope.payload["finding_index"],
                "status": _status(verdicts),
                "verdicts": verdicts,
            }])
        except Exception as exc:
            return failure(scope, exc)


def apply_verdicts(findings, verify_results):
    """Annotate findings with verification; return refuted ones to suppress."""
    by_index = {}
    for result in verify_results:
        for artifact in result.artifacts:
            by_index[artifact["finding_index"]] = artifact
    suppressed = []
    for index, finding in enumerate(findings):
        artifact = by_index.get(index)
        if not artifact:
            continue
        finding["verification"] = {"status": artifact["status"],
                                   "verdicts": artifact["verdicts"]}
        if artifact["status"] == "refuted":
            finding["confidence"] = "low"
            suppressed.append(finding)
    return suppressed