"""
_stage.py - Shared helpers for the LLM QAQC stages (Stages 0-11).

Every LLM stage follows the same shape:
  1. render a user instruction by replacing {placeholders}
     (str.replace, NOT str.format -- the prompts carry literal JSON braces),
  2. call the multimodal model via llm.call_json,
  3. validate / coerce the result so malformed model output degrades
     gracefully (drop the item) instead of crashing the pipeline.

This module factors that out so each stage module stays thin, and defines the
canonical "issue" schema shared by Stages 6-9 (qaqc / code / constructability /
dedup-validate). Mirrors the proven shape of conflict_checker._valid_conflict.
"""

from typing import Any, Callable, Dict, List, Optional

from backend.llm import call_json

_SEVERITY = ("critical", "high", "medium", "low")
_CONFIDENCE = ("high", "medium", "low")


def render(instruction: str, subs: Dict[str, str]) -> str:
    """Substitute {key} placeholders via str.replace (brace-safe for JSON prompts)."""
    out = instruction
    for k, v in (subs or {}).items():
        out = out.replace("{" + k + "}", "" if v is None else str(v))
    return out


def call_stage(
    system_prompt: str,
    user_instruction: str,
    subs: Optional[Dict[str, str]] = None,
    images_b64: Optional[List[str]] = None,
    max_tokens: int = 4096,
    model: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
    """Run one LLM call for a stage; returns the parsed JSON dict or None."""
    user_text = render(user_instruction, subs or {})
    return call_json(
        system_prompt=system_prompt,
        user_text=user_text,
        images_b64=images_b64,
        max_tokens=max_tokens,
        model=model,
    )


def collect_list(
    parsed: Optional[Any],
    key: str,
    validate: Optional[Callable[[Dict], Optional[Dict]]] = None,
) -> List[Dict]:
    """
    Pull a list of items out of a model response, optionally validating each.

    Tolerant of shape drift: the model may return {key: [...]} or a bare [...].
    """
    if isinstance(parsed, list):
        items = parsed
    elif isinstance(parsed, dict):
        items = parsed.get(key) or []
    else:
        return []
    if not isinstance(items, list):
        return []
    out: List[Dict] = []
    for it in items:
        if not isinstance(it, dict):
            continue
        v = validate(it) if validate else it
        if v:
            out.append(v)
    return out


def validate_issue(c: Dict, source_stage: str = "") -> Optional[Dict]:
    """
    Coerce + validate one canonical QAQC issue; drop if it has no description.

    Tolerant of schema drift between prompts: accepts evidence[].sheet or
    evidence[].sheet_number, and recommended_resolution or recommended_next_step.
    """
    if not isinstance(c, dict):
        return None
    description = (c.get("description") or c.get("summary") or "").strip()
    if not description:
        return None  # an issue with no description is noise

    severity = (c.get("severity") or "medium").strip().lower()
    if severity not in _SEVERITY:
        severity = "medium"
    confidence = (c.get("confidence") or "medium").strip().lower()
    if confidence not in _CONFIDENCE:
        confidence = "medium"

    evidence: List[Dict] = []
    for ev in (c.get("evidence") or []):
        if isinstance(ev, dict):
            evidence.append({
                "discipline": ev.get("discipline") or "?",
                "sheet": ev.get("sheet") or ev.get("sheet_number") or "?",
                "source_text": ev.get("source_text") or "",
                "asserted_value": ev.get("asserted_value") or "",
            })

    return {
        "issue_id": (c.get("issue_id") or "").strip(),
        "source_stage": (c.get("source_stage") or source_stage).strip(),
        "category": (c.get("category") or c.get("issue_type") or "uncategorized").strip(),
        "severity": severity,
        "confidence": confidence,
        "location": (c.get("location") or "").strip(),
        "disciplines": c.get("disciplines") or c.get("disciplines_involved") or [],
        "sheets": c.get("sheets") or [],
        "description": description,
        "evidence": evidence,
        "recommended_resolution": (
            c.get("recommended_resolution") or c.get("recommended_next_step") or ""
        ).strip(),
        "code_reference": c.get("code_reference"),
        "risk_score": c.get("risk_score"),
    }
