"""Scoped code/accessibility review agents."""

from typing import Dict, List

from backend import config
from backend.agents.base import AgentResult, AgentScope, AgentUsage, failure
from backend.llm import call_json
from backend.pipeline import code_refs
from backend.pipeline._serialize import dumps, slim_sheets
from backend.pipeline._stage import collect_list, validate_issue
from backend.pipeline.jurisdiction import active_review_paths
from backend.prompts import CODE_REVIEW_SYSTEM_PROMPT, CODE_REVIEW_USER_INSTRUCTION


def build_code_scopes(
    sheets: List[Dict], jurisdiction: Dict, sheet_index: Dict
) -> List[AgentScope]:
    cap = max(1, config.AGENT_CODE_BATCH_SIZE)
    fragments: List[Dict] = []
    for sheet in sheets:
        assertions = sheet.get("assertions") or []
        if not assertions:
            continue
        for offset in range(0, len(assertions), cap):
            fragments.append({
                **sheet,
                "assertions": assertions[offset:offset + cap],
            })
    batches: List[List[Dict]] = []
    current: List[Dict] = []
    count = 0
    for fragment in fragments:
        size = len(fragment["assertions"])
        if current and count + size > cap:
            batches.append(current)
            current, count = [], 0
        current.append(fragment)
        count += size
    if current:
        batches.append(current)
    return [
        AgentScope(
            scope_id=f"code:{index + 1}",
            payload={
                "sheets": batch,
                "jurisdiction": jurisdiction,
                "sheet_index": sheet_index,
            },
        )
        for index, batch in enumerate(batches)
    ]


class CodeAgent:
    name = "code_reviewer"

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

    def run(self, scope: AgentScope) -> AgentResult:
        try:
            sheets = scope.payload.get("sheets") or []
            jurisdiction = scope.payload.get("jurisdiction") or {}
            sheet_index = scope.payload.get("sheet_index") or {}
            assertions = [
                assertion for sheet in sheets
                for assertion in sheet.get("assertions", [])
            ]
            excerpts = code_refs.retrieve(
                active_review_paths(jurisdiction), assertions
            )
            instruction = CODE_REVIEW_USER_INSTRUCTION
            for key, value in {
                "jurisdiction": dumps(jurisdiction),
                "sheet_index": dumps(sheet_index),
                "assertions": dumps(slim_sheets(sheets)),
                "code_references": code_refs.format_excerpts(excerpts),
            }.items():
                instruction = instruction.replace("{" + key + "}", value)
            parsed = call_json(
                system_prompt=CODE_REVIEW_SYSTEM_PROMPT,
                user_text=instruction,
                max_tokens=config.CODE_MAX_TOKENS,
                model=config.AGENT_CODE_MODEL,
                usage_tracker=self.usage,
                usage_stage="agent.code",
            )
            findings = collect_list(
                parsed, "issues", lambda item: validate_issue(item, "code")
            )
            for finding in findings:
                finding.update(agent=self.name, scope_id=scope.scope_id)
            return AgentResult(scope_id=scope.scope_id, artifacts=findings)
        except Exception as exc:
            return failure(scope, exc)