"""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_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)]