"""Strict parser for the joint output contract (spec §5.1).

Never raises: any malformed input yields `ParsedOutput(format_ok=False, ...)`
with `parse_errors` populated. This is required by the Phase 3 gate (fuzz-tested
on >=100 malformed inputs; must never crash) since parse failure feeds the R5
reward guard at training/eval time, not a Python exception.
"""
from __future__ import annotations

import re

from src.contracts.schema import ParsedOutput, TraceStep

_ANALYSIS_RE = re.compile(r"<analysis>(.*?)</analysis>", re.DOTALL)
_LABEL_RE = re.compile(r"<label>\s*(.*?)\s*</label>", re.DOTALL)
_CWE_RE = re.compile(r"<cwe>\s*(.*?)\s*</cwe>", re.DOTALL)
_STATEMENTS_RE = re.compile(r"<statements>\s*(.*?)\s*</statements>", re.DOTALL)
_TRACE_BLOCK_RE = re.compile(r"<trace>(.*?)</trace>", re.DOTALL)
_STEP_RE = re.compile(r"<step>\s*(.*?)\s*</step>", re.DOTALL)
_EXPLANATION_RE = re.compile(r"<explanation>\s*(.*?)\s*</explanation>", re.DOTALL)

_STEP_LINE_RE = re.compile(
    r"^\s*line\s+(\d+)\s*:?\s*(.*?)\s*(?:\[\s*(.*?)\s*\])?\s*$", re.IGNORECASE
)
_VALID_LABELS = {"VULNERABLE", "SAFE"}
_CWE_VALUE_RE = re.compile(r"^(CWE-\d+|NONE)$", re.IGNORECASE)
_RANGE_RE = re.compile(r"^\s*(\d+)\s*(?:-\s*(\d+))?\s*$")


def parse_statement_lines(text: str) -> tuple[set[int], list[str]]:
    """Parses 'comma-separated line ranges, e.g. 12-14, 27' -> a set of line numbers."""
    lines: set[int] = set()
    errors: list[str] = []
    text = (text or "").strip()
    if not text:
        return lines, errors
    for chunk in text.split(","):
        chunk = chunk.strip()
        if not chunk:
            continue
        m = _RANGE_RE.match(chunk)
        if not m:
            errors.append(f"unparseable statement range token: {chunk!r}")
            continue
        start = int(m.group(1))
        end = int(m.group(2)) if m.group(2) else start
        if end < start:
            errors.append(f"inverted range: {chunk!r}")
            continue
        if end - start > 10_000:
            errors.append(f"implausibly large range (rejected): {chunk!r}")
            continue
        lines.update(range(start, end + 1))
    return lines, errors


def _parse_trace(trace_block: str) -> tuple[list[TraceStep], list[str]]:
    steps: list[TraceStep] = []
    errors: list[str] = []
    for step_text in _STEP_RE.findall(trace_block):
        m = _STEP_LINE_RE.match(step_text)
        if not m:
            errors.append(f"unparseable trace step: {step_text!r}")
            continue
        line_no, code, tag = m.group(1), m.group(2) or "", m.group(3) or ""
        try:
            steps.append(TraceStep(line=int(line_no), code=code, tag=tag))
        except (TypeError, ValueError):
            errors.append(f"non-integer line in trace step: {step_text!r}")
    return steps, errors


def parse_output(raw: str, contract_version: int = 1) -> ParsedOutput:
    """Parses one model completion against the joint output contract.

    Returns format_ok=False (never raises) on any structural violation:
    missing <analysis> block, missing/invalid <label>, missing/invalid <cwe>,
    or an empty <explanation>. Statement/trace parse errors are logged but do
    not by themselves flip format_ok to False, since partial localization or
    trace output is still scoreable by R2/R3 on whatever did parse.
    """
    if contract_version == 2:
        from src.contracts.v2 import parse_v2
        return parse_v2(raw)
    if contract_version != 1:
        return ParsedOutput(format_ok=False, parse_errors=["unsupported contract version"])
    raw = raw if isinstance(raw, str) else ""
    errors: list[str] = []

    analysis_match = _ANALYSIS_RE.search(raw)
    if not analysis_match:
        return ParsedOutput(format_ok=False, raw=raw, parse_errors=["no <analysis> block found"])
    body = analysis_match.group(1)

    label_match = _LABEL_RE.search(body)
    label = label_match.group(1).strip().upper() if label_match else None
    if label not in _VALID_LABELS:
        errors.append(f"invalid or missing <label>: {label!r}")
        return ParsedOutput(format_ok=False, raw=raw, label=label, parse_errors=errors)

    cwe_match = _CWE_RE.search(body)
    cwe = cwe_match.group(1).strip().upper() if cwe_match else None
    if not cwe or not _CWE_VALUE_RE.match(cwe):
        errors.append(f"invalid or missing <cwe>: {cwe!r}")
        return ParsedOutput(format_ok=False, raw=raw, label=label, cwe=cwe, parse_errors=errors)

    statements_match = _STATEMENTS_RE.search(body)
    statement_lines, stmt_errors = parse_statement_lines(
        statements_match.group(1) if statements_match else ""
    )
    errors.extend(stmt_errors)

    trace_block_match = _TRACE_BLOCK_RE.search(body)
    trace, trace_errors = (
        _parse_trace(trace_block_match.group(1)) if trace_block_match else ([], [])
    )
    errors.extend(trace_errors)

    explanation_match = _EXPLANATION_RE.search(body)
    explanation = explanation_match.group(1).strip() if explanation_match else ""
    if not explanation:
        errors.append("missing or empty <explanation>")
        return ParsedOutput(
            format_ok=False,
            raw=raw,
            label=label,
            cwe=cwe,
            statement_lines=statement_lines,
            trace=trace,
            parse_errors=errors,
        )

    return ParsedOutput(
        format_ok=True,
        raw=raw,
        label=label,
        cwe=cwe,
        statement_lines=statement_lines,
        trace=trace,
        explanation=explanation,
        parse_errors=errors,
    )