"""
code_review.py - Stage 7: code / ADA / TDLR / Municode review (LLM + retrieval).

Flags likely code and accessibility issues, grounded in BOTH the drawing
evidence and retrieved excerpts of the applicable standards (see code_refs.py).
The model may only cite section numbers that appear in the retrieved excerpts;
with an empty corpus it runs reasoning-only and cites nothing.

Sheets are processed in batches (bounded assertion count per call) so a single
oversized whole-set call can't truncate and zero out the whole stage. Each batch
retrieves excerpts targeted to its own assertions. Emits the canonical issue
schema; returns [] on failure.
"""

from concurrent.futures import ThreadPoolExecutor
from typing import Dict, List

from backend import config
from backend.pipeline import code_refs
from backend.pipeline.jurisdiction import active_review_paths
from backend.pipeline._serialize import dumps, slim_sheets
from backend.pipeline._stage import call_stage, collect_list, validate_issue
from backend.prompts import CODE_REVIEW_SYSTEM_PROMPT, CODE_REVIEW_USER_INSTRUCTION


def _sheet_batches(sheets: List[Dict], max_assertions: int) -> List[List[Dict]]:
    """Group whole sheets so each batch holds <= max_assertions facts."""
    batches: List[List[Dict]] = []
    cur: List[Dict] = []
    n = 0
    for s in sheets:
        a = len(s.get("assertions", []))
        if cur and n + a > max_assertions:
            batches.append(cur)
            cur, n = [], 0
        cur.append(s)
        n += a
    if cur:
        batches.append(cur)
    return batches


def _review_batch(batch: List[Dict], jurisdiction: Dict, sheet_index: Dict,
                  paths: Dict) -> List[Dict]:
    assertions = [a for s in batch for a in s.get("assertions", [])]
    excerpts = code_refs.retrieve(paths, assertions)
    parsed = call_stage(
        CODE_REVIEW_SYSTEM_PROMPT,
        CODE_REVIEW_USER_INSTRUCTION,
        subs={
            "jurisdiction": dumps(jurisdiction or {}),
            "sheet_index": dumps(sheet_index or {}),
            "assertions": dumps(slim_sheets(batch)),
            "code_references": code_refs.format_excerpts(excerpts),
        },
        max_tokens=config.CODE_MAX_TOKENS,
    )
    return collect_list(parsed, "issues", lambda c: validate_issue(c, "code"))


def code_review(jurisdiction: Dict, sheets: List[Dict], sheet_index: Dict) -> List[Dict]:
    paths = active_review_paths(jurisdiction)
    if not code_refs.load_corpus():
        print("[Code] no corpus -> reasoning-only (no citations)")
    batches = _sheet_batches(sheets, config.CODE_BATCH_SIZE)

    issues: List[Dict] = []
    with ThreadPoolExecutor(max_workers=config.CODE_CONCURRENCY) as pool:
        for res in pool.map(lambda b: _review_batch(b, jurisdiction, sheet_index, paths), batches):
            issues.extend(res)
    print(f"[Code] {len(issues)} code/accessibility issue(s) across {len(batches)} batch(es)")
    return issues