"""Three-way verifier (Anthropic-style):
  LLM proposer ↔ Linter guardian ↔ Regex secondary verifier
  Three-way agreement → ship
  Two-way agreement → ship low-priority
  One-way → escalate (abstain)

Inputs:
  - LLM output: {"segments": [{"intent","slots","text"}, ...], "abstain_reason": null|str}
  - Linter check: schema valid? intent in 51-enum? slots in lowercase whitelist? FMM-resolvable?
  - Regex parser: best-effort intent classification on raw text

Output verdict:
  - "ACCEPT_HIGH" — all three agree on (intent or compound structure)
  - "ACCEPT_LOW" — two-way agreement
  - "ESCALATE" — one-way or zero — abstain
  - "REJECT" — LLM emits invalid schema or abstain_reason set

Usage:
    from three_way_verifier import verify
    verdict = verify(llm_output, raw_text, regex_parse_fn=run_regex)
"""
from __future__ import annotations
import json
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

ROOT = Path(__file__).resolve().parents[3]
INTENTS_50 = sorted(json.loads((ROOT / "poc/deberta_intent/checkpoints-base/label_mapping.json").read_text())["intent2id"].keys())
INTENTS_51 = set(INTENTS_50 + ["unknown"])
SLOTS_LOWER = {"altimeter_setting", "altitude", "approach_type", "call_sign", "clock_position",
               "direction", "distance", "facility", "fix", "frequency", "heading", "pattern_leg",
               "route", "runway", "speed", "taxiway", "time", "transponder_code", "turn_direction",
               "sequence"}


@dataclass
class VerdictResult:
    verdict: str  # ACCEPT_HIGH | ACCEPT_LOW | ESCALATE | REJECT
    score: float  # 0.0 - 1.0 confidence
    llm_intent: str | None = None
    regex_intent: str | None = None
    linter_passed: bool = False
    agreements: int = 0  # how many of the 3 votes agree on intent
    reasons: list[str] = field(default_factory=list)
    payload: dict | None = None


def _strip_think(text: str) -> str:
    return re.sub(r'<think>.*?</think>\s*', '', text, flags=re.DOTALL).strip()


def parse_llm_output(raw: str) -> dict | None:
    """Strip thinking tokens + code fences, json.loads."""
    text = _strip_think(raw)
    if text.startswith("```"):
        text = text.split("```")[1]
        if text.startswith("json"):
            text = text[4:]
    try:
        return json.loads(text.strip())
    except Exception:
        return None


def lint_segment(seg: dict) -> tuple[bool, list[str]]:
    """L1+L3+L6 minimal: schema, intent enum, slot key whitelist."""
    errs = []
    if not isinstance(seg, dict):
        return False, ["segment_not_dict"]
    intent = seg.get("intent")
    if intent not in INTENTS_51:
        errs.append(f"intent_not_in_enum:{intent}")
    slots = seg.get("slots") or {}
    if not isinstance(slots, dict):
        errs.append("slots_not_dict")
    else:
        for k in slots:
            if k.lower() not in SLOTS_LOWER:
                errs.append(f"unknown_slot_key:{k}")
    return len(errs) == 0, errs


def lint_payload(payload: dict | None) -> tuple[bool, list[str]]:
    if not isinstance(payload, dict):
        return False, ["payload_not_dict"]
    segs = payload.get("segments")
    if not isinstance(segs, list) or len(segs) == 0:
        return False, ["segments_missing_or_empty"]
    errs = []
    for i, s in enumerate(segs):
        ok, e = lint_segment(s)
        if not ok:
            errs.extend([f"seg{i}.{x}" for x in e])
    return len(errs) == 0, errs


def verify(llm_raw_output: str, raw_text: str, regex_parse_fn: Callable[[str], str | None] | None = None) -> VerdictResult:
    """Run three-way verification. Returns VerdictResult."""
    payload = parse_llm_output(llm_raw_output)

    # Layer 1: schema/parse
    if payload is None:
        return VerdictResult(verdict="REJECT", score=0.0, reasons=["llm_output_unparseable"])

    # Layer 2: linter
    linter_ok, linter_errs = lint_payload(payload)
    if not linter_ok:
        return VerdictResult(verdict="REJECT", score=0.0, linter_passed=False,
                             reasons=linter_errs, payload=payload)

    # Abstain check
    if payload.get("abstain_reason"):
        return VerdictResult(verdict="ESCALATE", score=0.2, linter_passed=True,
                             reasons=[f"llm_abstained:{payload['abstain_reason']}"], payload=payload)

    # Pull primary intent (first segment's intent)
    segs = payload["segments"]
    llm_intent = segs[0].get("intent")
    if llm_intent == "unknown":
        return VerdictResult(verdict="ESCALATE", score=0.3, llm_intent="unknown",
                             linter_passed=True, reasons=["llm_chose_unknown"], payload=payload)

    # Layer 3: regex secondary
    regex_intent = None
    if regex_parse_fn:
        try:
            regex_intent = regex_parse_fn(raw_text)
        except Exception as e:
            regex_intent = None

    # Compute agreement
    agreements = 1  # LLM always votes once
    if linter_ok:
        agreements += 1
    if regex_intent and regex_intent == llm_intent:
        agreements += 1

    if agreements >= 3:
        return VerdictResult(verdict="ACCEPT_HIGH", score=0.95, llm_intent=llm_intent,
                             regex_intent=regex_intent, linter_passed=True,
                             agreements=agreements, payload=payload)
    if agreements == 2:
        return VerdictResult(verdict="ACCEPT_LOW", score=0.7, llm_intent=llm_intent,
                             regex_intent=regex_intent, linter_passed=True,
                             agreements=agreements,
                             reasons=["regex_disagreement"] if regex_intent and regex_intent != llm_intent else [],
                             payload=payload)
    return VerdictResult(verdict="ESCALATE", score=0.4, llm_intent=llm_intent,
                         regex_intent=regex_intent, linter_passed=linter_ok,
                         agreements=agreements, reasons=["only_llm_voted"], payload=payload)


# Convenience: simplistic regex stub for testing
def _mock_regex_parse(text: str) -> str | None:
    """Tiny placeholder; real one would use modules/intent/parser.py."""
    t = text.lower()
    if "cleared to land" in t or "cleared for landing" in t:
        return "landing_clearance"
    if "cleared for takeoff" in t:
        return "takeoff_clearance"
    if "contact" in t and ("tower" in t or "ground" in t or "approach" in t):
        return "frequency_change"
    if "squawk" in t:
        return "squawk_code_set"
    if "taxi" in t:
        return "taxi_instruction"
    return None


if __name__ == "__main__":
    # Self-test
    cases = [
        # case 1: clean accept
        ('{"segments":[{"intent":"landing_clearance","slots":{"runway":"two seven"},"text":"cleared to land runway two seven"}],"abstain_reason":null}',
         "Cleared to land runway two seven"),
        # case 2: regex disagreement
        ('{"segments":[{"intent":"informational","slots":{},"text":"foo"}],"abstain_reason":null}',
         "Cleared to land runway two seven"),
        # case 3: invalid intent
        ('{"segments":[{"intent":"made_up_intent","slots":{},"text":"foo"}],"abstain_reason":null}',
         "test"),
        # case 4: abstain
        ('{"segments":[{"intent":"unknown","slots":{},"text":"static"}],"abstain_reason":"garbled_transcript"}',
         "zzzkkrrr"),
        # case 5: malformed
        ('not json at all', "test"),
    ]
    for i, (raw, text) in enumerate(cases):
        v = verify(raw, text, regex_parse_fn=_mock_regex_parse)
        print(f"\n=== case {i+1} ===")
        print(f"  verdict: {v.verdict}  score={v.score:.2f}")
        print(f"  llm_intent={v.llm_intent}  regex_intent={v.regex_intent}")
        print(f"  agreements={v.agreements}  linter_passed={v.linter_passed}")
        print(f"  reasons={v.reasons}")