"""
State Machine Pipeline — generated by aiwg nlp new
Pattern: state-machine
Dependencies: anthropic

Install: pip install anthropic
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from pathlib import Path
from typing import Any

import anthropic

PROMPTS_DIR = Path(__file__).parent.parent / "prompts"
FSM_CONFIG_PATH = Path(__file__).parent.parent / "fsm.config.yaml"
AUDIT_LOG = Path(__file__).parent.parent / "audit" / "transitions.jsonl"
MODEL = "claude-haiku-4-5"
TIMEOUT_SECONDS = 30
MAX_RETRIES = 3
BACKOFF_SECONDS = 1.0


# ---------------------------------------------------------------------------
# FSM State and Transition types
# ---------------------------------------------------------------------------

class TerminalResult(Enum):
    ACCEPT = "accept"
    REJECT = "reject"
    ESCALATE = "escalate"


@dataclass
class PipelineContext:
    """Mutable context passed through FSM states."""
    input_document: str
    fields: dict[str, Any] = field(default_factory=dict)
    retry_count: int = 0
    state_history: list[str] = field(default_factory=list)
    raw_outputs: dict[str, str] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# LLM call
# ---------------------------------------------------------------------------

def call_llm(client: anthropic.Anthropic, system: str, user: str, model: str = MODEL) -> str:
    last_error = None
    for attempt in range(MAX_RETRIES):
        try:
            response = client.messages.create(
                model=model,
                max_tokens=512,
                system=system,
                messages=[{"role": "user", "content": user}],
                timeout=TIMEOUT_SECONDS,
            )
            return response.content[0].text
        except anthropic.RateLimitError as e:
            last_error = e
            time.sleep(BACKOFF_SECONDS * (2 ** attempt))
        except anthropic.APIStatusError as e:
            if e.status_code in {502, 503}:
                last_error = e
                time.sleep(BACKOFF_SECONDS * (2 ** attempt))
            else:
                raise
    raise RuntimeError(f"LLM call failed after {MAX_RETRIES} attempts") from last_error


def load_prompt(filename: str, variables: dict[str, str]) -> tuple[str, str]:
    path = PROMPTS_DIR / filename
    content = path.read_text(encoding="utf-8")
    if content.startswith("---"):
        _, _, content = content.split("---", 2)
    system, user, section = "", "", None
    for line in content.splitlines():
        if line.strip() == "## System":
            section = "system"
        elif line.strip() == "## User":
            section = "user"
        elif section == "system":
            system += line + "\n"
        elif section == "user":
            user += line + "\n"
    for k, v in variables.items():
        system = system.replace("{{" + k + "}}", v)
        user = user.replace("{{" + k + "}}", v)
    return system.strip(), user.strip()


# ---------------------------------------------------------------------------
# Audit logging
# ---------------------------------------------------------------------------

def log_transition(from_state: str, to_state: str, guard: str, ctx: PipelineContext) -> None:
    AUDIT_LOG.parent.mkdir(parents=True, exist_ok=True)
    record = {
        "from": from_state,
        "to": to_state,
        "guard": guard,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "retry_count": ctx.retry_count,
    }
    with AUDIT_LOG.open("a") as f:
        f.write(json.dumps(record) + "\n")


# ---------------------------------------------------------------------------
# State handlers
# ---------------------------------------------------------------------------

def state_extract(client: anthropic.Anthropic, ctx: PipelineContext) -> str:
    """EXTRACT: LLM extraction step."""
    system, user = load_prompt("extract.prompt.md", {"input_document": ctx.input_document})
    raw = call_llm(client, system, user)
    ctx.raw_outputs["extract"] = raw
    raw = raw.strip()
    if raw.startswith("```"):
        lines = raw.splitlines()
        raw = "\n".join(lines[1:-1] if lines[-1] == "```" else lines[1:])
    result = json.loads(raw)
    ctx.fields = result.get("fields", {})
    confidence = result.get("confidence", 0.0)
    if ctx.fields and confidence >= 0.5:
        return "VALIDATE"
    return "ESCALATE"


def state_validate(ctx: PipelineContext) -> str:
    """VALIDATE: Code-based decision — no LLM call."""
    required_fields = ["field_one", "field_two"]  # TODO: set your required fields
    present = [f for f in required_fields if ctx.fields.get(f) is not None]
    completeness = len(present) / len(required_fields) if required_fields else 1.0
    if completeness >= 0.8:
        return "ENRICH"
    if ctx.retry_count < 1:
        ctx.retry_count += 1
        return "EXTRACT"  # Retry with same model (or swap to stronger in retry handler)
    return "ESCALATE"


def state_enrich(client: anthropic.Anthropic, ctx: PipelineContext) -> str:
    """ENRICH: Optional enrichment step — add or normalize fields."""
    # TODO: implement enrichment prompt if needed
    # For now, pass through
    return "OUTPUT"


# ---------------------------------------------------------------------------
# FSM runner
# ---------------------------------------------------------------------------

def run_fsm(client: anthropic.Anthropic, ctx: PipelineContext) -> TerminalResult:
    state = "EXTRACT"
    while True:
        ctx.state_history.append(state)

        if state == "EXTRACT":
            next_state = state_extract(client, ctx)
        elif state == "VALIDATE":
            next_state = state_validate(ctx)
        elif state == "ENRICH":
            next_state = state_enrich(client, ctx)
        elif state == "OUTPUT":
            log_transition(state, "TERMINAL:accept", "output_ready", ctx)
            return TerminalResult.ACCEPT
        elif state == "ESCALATE":
            log_transition(state, "TERMINAL:escalate", "escalated", ctx)
            return TerminalResult.ESCALATE
        else:
            raise ValueError(f"Unknown state: {state}")

        log_transition(state, next_state, "guard_passed", ctx)
        state = next_state


# ---------------------------------------------------------------------------
# Pipeline entry
# ---------------------------------------------------------------------------

def run(input_document: str) -> dict[str, Any]:
    client = anthropic.Anthropic()
    ctx = PipelineContext(input_document=input_document)
    result = run_fsm(client, ctx)
    return {
        "result": result.value,
        "fields": ctx.fields,
        "state_history": ctx.state_history,
        "retries": ctx.retry_count,
    }


if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("Usage: python pipeline.py '<input document>'")
        sys.exit(1)
    print(json.dumps(run(sys.argv[1]), indent=2))