"""Central merge, judge, and prioritization agent."""

import json
import re
from typing import Dict, List, Tuple

from backend import config
from backend.agents.base import AgentUsage
from backend.agents.prompts import (
    BRAIN_CLARIFY_SYSTEM_PROMPT,
    BRAIN_CLARIFY_USER_PROMPT,
    BRAIN_SYSTEM_PROMPT,
    BRAIN_USER_PROMPT,
)
from backend.llm import call_json
from backend.pipeline._stage import collect_list, validate_issue


def _finding_ref(finding: Dict, index: int) -> str:
    return (
        finding.get("issue_id")
        or f"{finding.get('agent', 'agent')}:{finding.get('scope_id', '?')}:{index + 1}"
    )


def _signature(finding: Dict) -> Tuple[str, str, str]:
    norm = lambda value: re.sub(r"[^a-z0-9]+", " ", str(value).lower()).strip()
    description = " ".join(norm(finding.get("description")).split()[:12])
    return (
        norm(finding.get("category")),
        norm(finding.get("location")),
        description,
    )


def _fallback(findings: List[Dict]) -> Tuple[List[Dict], List[Dict]]:
    """Conservative local consolidation when the Brain call fails."""
    kept: Dict[Tuple[str, str, str], Dict] = {}
    refs: Dict[Tuple[str, str, str], List[str]] = {}
    decisions: List[Dict] = []
    severity_rank = {"critical": 4, "high": 3, "medium": 2, "low": 1}
    for index, finding in enumerate(findings):
        ref = _finding_ref(finding, index)
        supported = bool(finding.get("evidence")) or finding.get("agent") == "completeness"
        if not supported or not finding.get("description"):
            decisions.append({
                "finding_refs": [ref],
                "action": "dropped",
                "reason": "missing actionable support",
                "kept_issue_id": None,
            })
            continue
        signature = _signature(finding)
        if signature not in kept:
            kept[signature] = dict(finding)
            refs[signature] = [ref]
        else:
            refs[signature].append(ref)
            existing = kept[signature]
            if severity_rank.get(finding.get("severity"), 2) > severity_rank.get(
                existing.get("severity"), 2
            ):
                existing["severity"] = finding.get("severity")
            existing["evidence"] = (
                existing.get("evidence") or []
            ) + (finding.get("evidence") or [])

    issues = list(kept.values())
    for index, (signature, issue) in enumerate(kept.items()):
        issue["issue_id"] = issue.get("issue_id") or f"AGENT-{index + 1:04d}"
        issue["risk_score"] = {
            "critical": 95, "high": 75, "medium": 50, "low": 25
        }.get(issue.get("severity"), 50)
        issue["recommended_priority"] = {
            "critical": "immediate",
            "high": "before_bid",
            "medium": "before_construction",
            "low": "track_only",
        }.get(issue.get("severity"), "before_construction")
        decisions.append({
            "finding_refs": refs[signature],
            "action": "merged" if len(refs[signature]) > 1 else "kept",
            "reason": "conservative deterministic fallback",
            "kept_issue_id": issue["issue_id"],
        })
    issues.sort(key=lambda item: -int(item.get("risk_score") or 0))
    return issues, decisions


class BrainAgent:
    name = "brain"

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

    def run(
        self,
        findings: List[Dict],
        sheet_index: Dict,
        jurisdiction: Dict,
    ) -> Tuple[List[Dict], List[Dict]]:
        instruction = BRAIN_USER_PROMPT
        for key, value in {
            "sheet_index": sheet_index,
            "jurisdiction": jurisdiction,
            "findings": findings,
        }.items():
            instruction = instruction.replace(
                "{" + key + "}", json.dumps(value, ensure_ascii=True)
            )
        parsed = call_json(
            system_prompt=BRAIN_SYSTEM_PROMPT,
            user_text=instruction,
            max_tokens=config.AGENT_BRAIN_MAX_TOKENS,
            model=config.AGENT_BRAIN_MODEL,
            usage_tracker=self.usage,
            usage_stage="agent.brain",
        )
        issues = collect_list(
            parsed, "issues", lambda item: validate_issue(item, item.get("source_stage", ""))
        )
        if not issues:
            return _fallback(findings)
        raw_issues = parsed.get("issues") if isinstance(parsed, dict) else []
        for index, issue in enumerate(issues):
            raw = raw_issues[index] if index < len(raw_issues) else {}
            issue["issue_id"] = issue.get("issue_id") or f"AGENT-{index + 1:04d}"
            issue["risk_score"] = raw.get("risk_score") or issue.get("risk_score") or 50
            issue["recommended_priority"] = (
                raw.get("recommended_priority") or "before_construction"
            )
        issues.sort(key=lambda item: -int(item.get("risk_score") or 0))
        decisions = parsed.get("decisions") or []
        return issues, [item for item in decisions if isinstance(item, dict)]

    def plan_clarifications(self, prioritized: List[Dict]) -> List[Dict]:
        """Wave 6.5 planning call: name kept findings the Brain wants to
        double-check before publishing, as typed clarification requests.

        Returns a capped list of {issue_id, request_type, reason}. Only
        findings that carry an issue_id and do NOT already have a verification
        result are offered to the model; anything the model names outside that
        set, or with an unknown request_type, is dropped by the caller/router.
        Never raises — a failed/empty plan just yields no requests.
        """
        max_requests = config.BRAIN_CLARIFY_MAX_REQUESTS
        if not prioritized or max_requests <= 0:
            return []
        candidates = [
            {
                "issue_id": f.get("issue_id"),
                "severity": f.get("severity"),
                "confidence": f.get("confidence"),
                "source_stage": f.get("source_stage"),
                "description": (f.get("description") or "")[:400],
                "evidence": f.get("evidence") or [],
                "already_verified": bool(f.get("verification")),
            }
            for f in prioritized
            if f.get("issue_id") and not f.get("verification")
        ]
        if not candidates:
            return []
        instruction = (
            BRAIN_CLARIFY_USER_PROMPT
            .replace("{max_requests}", str(max_requests))
            .replace("{findings}", json.dumps(candidates, ensure_ascii=True))
        )
        try:
            parsed = call_json(
                system_prompt=BRAIN_CLARIFY_SYSTEM_PROMPT,
                user_text=instruction,
                max_tokens=config.BRAIN_CLARIFY_MAX_TOKENS,
                model=config.AGENT_BRAIN_MODEL,
                usage_tracker=self.usage,
                usage_stage="agent.brain_clarify",
            )
        except Exception:
            return []
        raw = parsed.get("requests") if isinstance(parsed, dict) else parsed
        if not isinstance(raw, list):
            return []
        valid_ids = {c["issue_id"] for c in candidates}
        requests: List[Dict] = []
        seen: set = set()
        for item in raw:
            if not isinstance(item, dict):
                continue
            issue_id = item.get("issue_id")
            if issue_id not in valid_ids or issue_id in seen:
                continue
            requests.append({
                "issue_id": issue_id,
                "request_type": (item.get("request_type") or "verify_evidence").strip(),
                "reason": (item.get("reason") or "").strip(),
            })
            seen.add(issue_id)
            if len(requests) >= max_requests:
                break
        return requests
