#!/usr/bin/env python3
"""Reference parser for .contract.md files (v1.0, v2.0, and v2.1).

Extracts all structured fields from a contract file without using an LLM.
Proves the format is machine-parseable by a simple tool.

Usage:
    python contract_parser.py <file.contract.md>
    python contract_parser.py --json <file.contract.md>     # machine output
    python contract_parser.py --validate <file.contract.md> # check required fields
    python contract_parser.py --list <file.contract.md>     # show function signatures only
    python contract_parser.py --dir <directory>             # parse all .contract.md files
    python contract_parser.py --dir <directory> --validate  # batch validate
"""
from __future__ import annotations

import json
import re
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import yaml


# ── Data classes ──────────────────────────────────────────────────────────────


@dataclass
class ErrorSpec:
    error_type: str
    recovery: str


@dataclass
class Precondition:
    id: str
    severity: str          # "hard" | "soft"
    condition: str
    validation: str = ""


@dataclass
class Postcondition:
    id: str
    category: str          # "return_value" | "state_change" | "side_effect" | "exception"
    condition: str
    validation: str = ""


@dataclass
class TestSpec:
    name: str
    category: str          # "happy" | "error" | "boundary" | "edge" | "security"
    description: str       # full text after category
    tags: tuple[str, ...] = ()  # modifiers: ("tracer",) etc.


@dataclass
class Step:
    number: int
    text: str
    step_type: str = ""    # "setup" | "sequential" | "branch" | "loop" | "error_handler" | "cleanup"
    sub_steps: list[str] = field(default_factory=list)
    branch_if: str | None = None
    branch_else: list[str] = field(default_factory=list)


@dataclass
class FunctionBlock:
    name: str
    params: str
    return_type: str
    brief: str
    requires: str | None = None          # v1 single-line REQUIRES
    ensures: str | None = None           # v1 single-line ENSURES
    preconditions: list[Precondition] = field(default_factory=list)    # v2
    postconditions: list[Postcondition] = field(default_factory=list)  # v2
    errors: list[ErrorSpec] = field(default_factory=list)
    state_transitions: list[str] = field(default_factory=list)
    steps: list[Step] = field(default_factory=list)
    tests: list[TestSpec] = field(default_factory=list)                # v2


@dataclass
class Contract:
    frontmatter: dict[str, Any]
    body_sections: dict[str, str]
    functions: list[FunctionBlock]
    source_path: Path | None = None

    @property
    def version(self) -> str:
        return str(self.frontmatter.get("contract_version", "1.0"))


# ── Validation ────────────────────────────────────────────────────────────────

REQUIRED_FRONTMATTER_V1 = [
    "contract_version", "module", "purpose", "language", "min_complexity",
]
REQUIRED_FRONTMATTER_V2_MODULE = [
    "contract_version", "module", "purpose", "language", "complexity",
]
REQUIRED_FRONTMATTER_V2_ISSUE = [
    "contract_version", "target_module", "scope", "language", "complexity", "prd",
]
RECOMMENDED_FRONTMATTER_V1 = ["depends_on", "used_by", "estimated_loc"]
RECOMMENDED_FRONTMATTER_V2 = ["depends_on", "used_by", "estimated_loc", "confidence"]
REQUIRED_BODY_SECTIONS = ["Context", "Data flow", "Invariants"]
VALID_COMPLEXITIES_V1 = {"trivial", "simple", "medium", "complex", "expert"}
VALID_COMPLEXITIES_V2 = {"low", "medium", "high"}
VALID_PRE_SEVERITIES = {"hard", "soft"}
VALID_POST_CATEGORIES = {"return_value", "state_change", "side_effect", "exception"}
VALID_STEP_TYPES = {"setup", "sequential", "branch", "loop", "error_handler", "cleanup"}
VALID_TEST_CATEGORIES_V20 = {"happy", "error", "boundary", "edge", "security"}
VALID_TEST_CATEGORIES_V21 = VALID_TEST_CATEGORIES_V20 | {"scenario", "trace", "adversarial", "property"}


_ISSUE_PATH_RE = re.compile(r"docs/contracts/issues/\d+\.contract\.md$")


def _is_issue_scoped(contract: Contract) -> bool:
    """Detect issue-scoped contracts per CONTRACT-FORMAT § 2.1.I."""
    if contract.source_path and _ISSUE_PATH_RE.search(str(contract.source_path)):
        return True
    return "prd" in contract.frontmatter


@dataclass
class ValidationIssue:
    severity: str  # "error" | "warning"
    message: str


def validate_contract(contract: Contract) -> list[ValidationIssue]:
    """Check a parsed contract for required fields and structural issues."""
    issues: list[ValidationIssue] = []
    fm = contract.frontmatter
    is_v2 = contract.version.startswith("2")
    is_v21 = contract.version.startswith("2.1")
    issue_scoped = is_v2 and _is_issue_scoped(contract)

    if is_v2:
        required_fm = REQUIRED_FRONTMATTER_V2_ISSUE if issue_scoped else REQUIRED_FRONTMATTER_V2_MODULE
    else:
        required_fm = REQUIRED_FRONTMATTER_V1
    recommended_fm = RECOMMENDED_FRONTMATTER_V2 if is_v2 else RECOMMENDED_FRONTMATTER_V1
    valid_complexities = VALID_COMPLEXITIES_V2 if is_v2 else VALID_COMPLEXITIES_V1
    complexity_key = "complexity" if is_v2 else "min_complexity"

    # Required frontmatter
    for key in required_fm:
        if key not in fm or fm[key] is None:
            issues.append(ValidationIssue("error", f"Missing required frontmatter field: {key!r}"))

    # Complexity value
    complexity = fm.get(complexity_key, "")
    if complexity and complexity not in valid_complexities:
        issues.append(ValidationIssue(
            "error",
            f"{complexity_key} {complexity!r} is not one of: {sorted(valid_complexities)}"
        ))

    # Confidence range (v2)
    if is_v2 and "confidence" in fm:
        conf = fm["confidence"]
        if isinstance(conf, (int, float)) and not (0.0 <= conf <= 1.0):
            issues.append(ValidationIssue("error", f"confidence {conf} out of range [0.0, 1.0]"))

    # Recommended frontmatter
    for key in recommended_fm:
        if key not in fm:
            issues.append(ValidationIssue("warning", f"Missing recommended frontmatter field: {key!r}"))

    # Required body sections
    for section in REQUIRED_BODY_SECTIONS:
        if section not in contract.body_sections:
            issues.append(ValidationIssue("warning", f"Missing recommended body section: {section!r}"))

    # Function block count
    n = len(contract.functions)
    if n == 0:
        issues.append(ValidationIssue("warning", "No function blocks found"))
    elif n > 8:
        issues.append(ValidationIssue("warning", f"{n} function blocks — consider splitting the module (spec recommends 3–8)"))

    # Per-function checks
    for fn in contract.functions:
        prefix = f"Function {fn.name!r}"
        if not fn.brief:
            issues.append(ValidationIssue("error", f"{prefix}: missing BRIEF"))
        if not fn.steps:
            issues.append(ValidationIssue("error", f"{prefix}: missing STEPS"))

        if is_v2:
            # v2: validate preconditions
            for pre in fn.preconditions:
                if pre.severity not in VALID_PRE_SEVERITIES:
                    issues.append(ValidationIssue(
                        "warning", f"{prefix}: PRE {pre.id} severity {pre.severity!r} not in {sorted(VALID_PRE_SEVERITIES)}"
                    ))

            # v2: validate postconditions
            for post in fn.postconditions:
                if post.category not in VALID_POST_CATEGORIES:
                    issues.append(ValidationIssue(
                        "warning", f"{prefix}: POST {post.id} category {post.category!r} not in {sorted(VALID_POST_CATEGORIES)}"
                    ))

            # v2: validate step types
            for step in fn.steps:
                if step.step_type and step.step_type not in VALID_STEP_TYPES:
                    issues.append(ValidationIssue(
                        "warning", f"{prefix}: step {step.number} type {step.step_type!r} not in {sorted(VALID_STEP_TYPES)}"
                    ))

            # v2: validate test categories
            valid_test_cats = VALID_TEST_CATEGORIES_V21 if is_v21 else VALID_TEST_CATEGORIES_V20
            for test in fn.tests:
                if test.category not in valid_test_cats:
                    issues.append(ValidationIssue(
                        "warning", f"{prefix}: test {test.name!r} category {test.category!r} not in {sorted(valid_test_cats)}"
                    ))
                # Modifier tags (e.g. "tracer") get the same vocabulary check.
                for tag in test.tags:
                    if tag not in _VALID_TEST_TAGS:
                        issues.append(ValidationIssue(
                            "warning",
                            f"{prefix}: test {test.name!r} tag {tag!r} not in {sorted(_VALID_TEST_TAGS)}",
                        ))

            # v2: at most one tracer per function block — pairs with the tdd skill
            tracer_tests = [t for t in fn.tests if "tracer" in t.tags]
            if len(tracer_tests) > 1:
                names = ", ".join(repr(t.name) for t in tracer_tests)
                issues.append(ValidationIssue(
                    "warning",
                    f"{prefix}: multiple tracer tests ({names}) — only one test should carry the [tracer] tag",
                ))

            # v2: warn if no preconditions
            if not fn.preconditions and not fn.requires:
                issues.append(ValidationIssue("warning", f"{prefix}: no preconditions (PRE lines)"))

            # v2: warn if no postconditions
            if not fn.postconditions and not fn.ensures:
                issues.append(ValidationIssue("warning", f"{prefix}: no postconditions (POST lines)"))

    return issues


# ── Parsing ────────────────────────────────────────────────────────────────────


def _parse_frontmatter(text: str) -> tuple[dict, str]:
    """Extract YAML frontmatter. Returns (metadata_dict, remaining_text)."""
    if not text.startswith("---"):
        raise ValueError("Contract must start with YAML frontmatter (---)")
    end = text.find("---", 3)
    if end == -1:
        raise ValueError("Unterminated frontmatter: no closing ---")
    meta = yaml.safe_load(text[3:end])
    body = text[end + 3:].strip()
    return meta, body


def _parse_body_sections(body: str) -> dict[str, str]:
    """Split body into named sections by ## headers."""
    sections: dict[str, str] = {}
    current = None
    lines: list[str] = []
    for line in body.split("\n"):
        if line.startswith("## "):
            if current is not None:
                sections[current] = "\n".join(lines).strip()
            current = line[3:].strip()
            lines = []
        elif current is not None:
            lines.append(line)
    if current is not None:
        sections[current] = "\n".join(lines).strip()
    return sections


def _extract_function_blocks(body: str) -> list[str]:
    """Extract raw function block text from ```contract fences."""
    pattern = r"```contract\s*\n(.*?)```"
    return re.findall(pattern, body, re.DOTALL)


_PRE_RE = re.compile(r"^\[(\S+)\s+(\S+)\]\s+(.*?)(?:\s+--\s+(.*))?$")
_POST_RE = re.compile(r"^\[(\S+)\s+(\S+)\]\s+(.*?)(?:\s+--\s+(.*))?$")
_TEST_RE = re.compile(r"^(\S+)\s+\[([\w,\s]+)\]:\s+(.*)$")
_VALID_TEST_TAGS = {"tracer"}
_STEP_TYPE_RE = re.compile(r"^\[(\w+)\]\s+(.*)")


def _parse_function_block(text: str) -> FunctionBlock:
    """Parse a single function block into structured fields."""
    lines = text.strip().split("\n")

    # Collect FN signature lines until we find one containing ") ->"
    fn_lines: list[str] = []
    body_start = 0
    for i, line in enumerate(lines):
        fn_lines.append(line)
        if ")" in line and "->" in line:
            body_start = i + 1
            break

    fn_text = " ".join(l.strip() for l in fn_lines)
    fn_match = re.match(r"FN\s+([\w.]+)\((.*)\)\s*->\s*(.+)", fn_text)
    if not fn_match:
        raise ValueError(f"Invalid FN signature: {fn_text!r}")
    name = fn_match.group(1)
    params = fn_match.group(2).strip()
    return_type = fn_match.group(3).strip()

    brief = ""
    requires = None
    ensures = None
    preconditions: list[Precondition] = []
    postconditions: list[Postcondition] = []
    errors: list[ErrorSpec] = []
    state_transitions: list[str] = []
    steps: list[Step] = []
    tests: list[TestSpec] = []
    in_errors = False
    in_steps = False
    in_tests = False
    in_else = False
    current_step: Step | None = None

    for line in lines[body_start:]:
        stripped = line.strip()

        if stripped.startswith("BRIEF:"):
            brief = stripped[len("BRIEF:"):].strip()
            in_errors = in_steps = in_tests = False
            continue

        if stripped.startswith("REQUIRES:"):
            requires = stripped[len("REQUIRES:"):].strip()
            in_errors = in_steps = in_tests = False
            continue

        if stripped.startswith("ENSURES:"):
            ensures = stripped[len("ENSURES:"):].strip()
            in_errors = in_steps = in_tests = False
            continue

        if stripped.startswith("PRE:"):
            in_errors = in_steps = in_tests = False
            rest = stripped[len("PRE:"):].strip()
            m = _PRE_RE.match(rest)
            if m:
                preconditions.append(Precondition(
                    id=m.group(1), severity=m.group(2),
                    condition=m.group(3).strip(),
                    validation=(m.group(4) or "").strip(),
                ))
            else:
                preconditions.append(Precondition(id="", severity="hard", condition=rest))
            continue

        if stripped.startswith("POST:"):
            in_errors = in_steps = in_tests = False
            rest = stripped[len("POST:"):].strip()
            m = _POST_RE.match(rest)
            if m:
                postconditions.append(Postcondition(
                    id=m.group(1), category=m.group(2),
                    condition=m.group(3).strip(),
                    validation=(m.group(4) or "").strip(),
                ))
            else:
                postconditions.append(Postcondition(id="", category="return_value", condition=rest))
            continue

        if stripped.startswith("STATE:"):
            state_transitions.append(stripped[len("STATE:"):].strip())
            in_errors = in_steps = in_tests = False
            continue

        if stripped == "STEPS:":
            in_steps = True
            in_errors = in_tests = False
            continue

        if stripped == "ERRORS:":
            in_errors = True
            in_steps = in_tests = False
            continue

        if stripped == "TESTS:":
            in_tests = True
            in_errors = in_steps = False
            # Flush last step
            if current_step is not None:
                steps.append(current_step)
                current_step = None
            continue

        if in_errors and "->" in stripped:
            parts = stripped.split("->", 1)
            errors.append(ErrorSpec(
                error_type=parts[0].strip(),
                recovery=parts[1].strip(),
            ))
            continue

        if in_tests and stripped:
            m = _TEST_RE.match(stripped)
            if m:
                # Bracket may hold "happy" or "happy,tracer". First token is
                # the category; remaining tokens are modifier tags.
                bracket_tokens = [t.strip() for t in m.group(2).split(",") if t.strip()]
                category = bracket_tokens[0] if bracket_tokens else ""
                tags = tuple(bracket_tokens[1:])
                tests.append(TestSpec(
                    name=m.group(1), category=category,
                    description=m.group(3).strip(),
                    tags=tags,
                ))
            continue

        # Numbered step
        step_match = re.match(r"(\d+)\.\s+(.*)", stripped)
        if step_match and in_steps:
            if current_step is not None:
                steps.append(current_step)
            step_text = step_match.group(2).strip()
            step_type = ""
            type_m = _STEP_TYPE_RE.match(step_text)
            if type_m:
                step_type = type_m.group(1)
                step_text = type_m.group(2).strip()
            current_step = Step(
                number=int(step_match.group(1)),
                text=step_text,
                step_type=step_type,
            )
            in_else = False
            continue

        # Sub-step or branch under a step
        if in_steps and current_step is not None and stripped:
            if re.match(r"IF\s+", stripped) or re.match(r"ELSE IF\s+", stripped):
                current_step.branch_if = stripped
                in_else = False
            elif stripped == "ELSE:" or stripped.startswith("ELSE "):
                in_else = True
            elif stripped.startswith("- "):
                content = stripped[2:]
                if in_else:
                    current_step.branch_else.append(content)
                else:
                    current_step.sub_steps.append(content)
            elif re.match(r"[A-Z]", stripped):
                if in_else:
                    current_step.branch_else.append(stripped)
                else:
                    current_step.sub_steps.append(stripped)
            elif re.match(r"[a-z]", stripped):
                if in_else:
                    current_step.branch_else.append(stripped)
                else:
                    current_step.sub_steps.append(stripped)

    if current_step is not None:
        steps.append(current_step)

    return FunctionBlock(
        name=name,
        params=params,
        return_type=return_type,
        brief=brief,
        requires=requires,
        ensures=ensures,
        preconditions=preconditions,
        postconditions=postconditions,
        errors=errors,
        state_transitions=state_transitions,
        steps=steps,
        tests=tests,
    )


def parse_contract(path: Path) -> Contract:
    """Parse a .contract.md file into a structured Contract object."""
    text = path.read_text(encoding="utf-8")
    frontmatter, body = _parse_frontmatter(text)
    body_sections = _parse_body_sections(body)
    raw_blocks = _extract_function_blocks(body)
    functions: list[FunctionBlock] = []
    for b in raw_blocks:
        try:
            functions.append(_parse_function_block(b))
        except ValueError:
            # Skip blocks that don't match FN signature syntax
            # (e.g. DATACLASS blocks, non-standard formats)
            pass
    return Contract(
        frontmatter=frontmatter,
        body_sections=body_sections,
        functions=functions,
        source_path=path,
    )


# ── Output ─────────────────────────────────────────────────────────────────────


def _to_dict(contract: Contract) -> dict[str, Any]:
    """Convert to JSON-serialisable dict."""
    return {
        "frontmatter": contract.frontmatter,
        "body_sections": contract.body_sections,
        "functions": [
            {
                "name": fn.name,
                "params": fn.params,
                "return_type": fn.return_type,
                "brief": fn.brief,
                "requires": fn.requires,
                "ensures": fn.ensures,
                "preconditions": [
                    {"id": p.id, "severity": p.severity, "condition": p.condition, "validation": p.validation}
                    for p in fn.preconditions
                ],
                "postconditions": [
                    {"id": p.id, "category": p.category, "condition": p.condition, "validation": p.validation}
                    for p in fn.postconditions
                ],
                "errors": [
                    {"error_type": e.error_type, "recovery": e.recovery}
                    for e in fn.errors
                ],
                "state_transitions": fn.state_transitions,
                "steps": [
                    {
                        "number": s.number,
                        "text": s.text,
                        "step_type": s.step_type,
                        "sub_steps": s.sub_steps,
                        "branch_if": s.branch_if,
                        "branch_else": s.branch_else,
                    }
                    for s in fn.steps
                ],
                "tests": [
                    {"name": t.name, "category": t.category, "description": t.description}
                    for t in fn.tests
                ],
            }
            for fn in contract.functions
        ],
    }


def print_human(contract: Contract) -> None:
    """Pretty-print contract summary for human reading."""
    fm = contract.frontmatter
    is_v2 = contract.version.startswith("2")
    complexity_key = "complexity" if is_v2 else "min_complexity"

    issue_scoped = is_v2 and _is_issue_scoped(contract)
    if issue_scoped:
        print(f"Target:     {fm.get('target_module', '?')}")
        print(f"Scope:      {fm.get('scope', '?')}")
    else:
        print(f"Module:     {fm.get('module', '?')}")
        print(f"Purpose:    {fm.get('purpose', '?')}")
    print(f"Version:    {contract.version}")
    print(f"Complexity: {fm.get(complexity_key, '?')}")
    print(f"Est. LOC:   {fm.get('estimated_loc', '?')}")
    if is_v2 and "confidence" in fm:
        print(f"Confidence: {fm['confidence']}")
    print(f"Depends:    {', '.join(fm.get('depends_on', []))}")
    print(f"Used by:    {', '.join(fm.get('used_by', []))}")
    if is_v2 and fm.get("assumptions"):
        print(f"Assumptions: {len(fm['assumptions'])}")
    if is_v2 and fm.get("open_questions"):
        print(f"Open Qs:    {len(fm['open_questions'])}")
    print()

    for name, text in contract.body_sections.items():
        print(f"## {name}")
        preview = text[:300] + ("..." if len(text) > 300 else "")
        print(preview)
        print()

    print(f"Function blocks: {len(contract.functions)}")
    for fn in contract.functions:
        print(f"\n  FN {fn.name}({fn.params}) -> {fn.return_type}")
        print(f"     BRIEF: {fn.brief}")
        if fn.requires:
            print(f"     REQUIRES: {fn.requires}")
        if fn.ensures:
            print(f"     ENSURES: {fn.ensures}")
        for pre in fn.preconditions:
            print(f"     PRE [{pre.id} {pre.severity}]: {pre.condition}")
        for post in fn.postconditions:
            print(f"     POST [{post.id} {post.category}]: {post.condition}")
        if fn.errors:
            print(f"     ERRORS ({len(fn.errors)}):")
            for e in fn.errors:
                print(f"       {e.error_type} -> {e.recovery}")
        typed_steps = sum(1 for s in fn.steps if s.step_type)
        print(f"     STEPS: {len(fn.steps)} ({typed_steps} typed)")
        if fn.tests:
            print(f"     TESTS: {len(fn.tests)}")
            for t in fn.tests:
                print(f"       [{t.category}] {t.name}: {t.description[:60]}")


def print_list(contract: Contract) -> None:
    """Print function signatures only — one per line."""
    fm = contract.frontmatter
    is_v2 = contract.version.startswith("2")
    complexity_key = "complexity" if is_v2 else "min_complexity"
    issue_scoped = is_v2 and _is_issue_scoped(contract)
    label = fm.get('target_module', '?') if issue_scoped else fm.get('module', '?')
    print(f"{label}  [{fm.get(complexity_key, '?')}]")
    for fn in contract.functions:
        params_short = fn.params[:60] + ("..." if len(fn.params) > 60 else "")
        print(f"  FN {fn.name}({params_short}) -> {fn.return_type}")
        print(f"     {fn.brief}")


def print_validation(contract: Contract, issues: list[ValidationIssue]) -> None:
    """Print validation results."""
    path_str = str(contract.source_path) if contract.source_path else "?"
    errors = [i for i in issues if i.severity == "error"]

    if not issues:
        print(f"\u2713 {path_str} \u2014 OK")
        return

    status = "FAIL" if errors else "WARN"
    marker = "\u2717" if errors else "!"
    print(f"{marker} {path_str} \u2014 {status}")
    for issue in issues:
        marker = "  ERROR  " if issue.severity == "error" else "  warn   "
        print(f"{marker}{issue.message}")


# ── CLI ────────────────────────────────────────────────────────────────────────


def main() -> None:
    if len(sys.argv) < 2:
        print("Usage: contract_parser.py [--json|--validate|--list] <file.contract.md>")
        print("       contract_parser.py --dir <directory> [--json|--validate|--list]")
        sys.exit(1)

    args = sys.argv[1:]
    json_output = "--json" in args
    validate_mode = "--validate" in args
    list_mode = "--list" in args
    dir_mode = "--dir" in args
    flags = {"--json", "--validate", "--list", "--dir"}
    positional = [a for a in args if a not in flags]

    if dir_mode:
        if not positional:
            print("--dir requires a directory path")
            sys.exit(1)
        directory = Path(positional[0])
        if not directory.is_dir():
            print(f"Not a directory: {directory}")
            sys.exit(1)
        files = sorted(directory.rglob("*.contract.md"))
        if not files:
            print(f"No .contract.md files found in {directory}")
            sys.exit(0)
        any_errors = False
        for path in files:
            try:
                contract = parse_contract(path)
                if json_output:
                    print(json.dumps({str(path): _to_dict(contract)}, indent=2))
                elif validate_mode:
                    issues = validate_contract(contract)
                    print_validation(contract, issues)
                    if any(i.severity == "error" for i in issues):
                        any_errors = True
                elif list_mode:
                    print_list(contract)
                    print()
                else:
                    print_human(contract)
                    print("\u2500" * 60)
            except Exception as exc:
                print(f"\u2717 {path} \u2014 PARSE ERROR: {exc}")
                any_errors = True
        sys.exit(1 if any_errors else 0)

    if not positional:
        print("Provide a .contract.md file path")
        sys.exit(1)

    path = Path(positional[0])
    if not path.exists():
        print(f"File not found: {path}")
        sys.exit(1)

    try:
        contract = parse_contract(path)
    except Exception as exc:
        print(f"Parse error: {exc}")
        sys.exit(1)

    if json_output:
        print(json.dumps(_to_dict(contract), indent=2))
    elif validate_mode:
        issues = validate_contract(contract)
        print_validation(contract, issues)
        if any(i.severity == "error" for i in issues):
            sys.exit(1)
    elif list_mode:
        print_list(contract)
    else:
        print_human(contract)


if __name__ == "__main__":
    main()