"""
Simple Chain Pipeline — generated by aiwg nlp new
Pattern: simple-chain
Dependencies: anthropic

Install: pip install anthropic
"""

from __future__ import annotations

import json
import time
from pathlib import Path
from typing import Any

import anthropic

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

PROMPTS_DIR = Path(__file__).parent.parent / "prompts"
MODEL = "claude-haiku-4-5"
MAX_TOKENS = 512
TIMEOUT_SECONDS = 30
MAX_RETRIES = 3
RETRY_ON_STATUS = {429, 502, 503}
BACKOFF_SECONDS = 1.0


def load_prompt(filename: str, variables: dict[str, str]) -> tuple[str, str]:
    """Load a prompt file and substitute {{variable}} slots.

    Returns (system_prompt, user_prompt) tuple.
    """
    path = PROMPTS_DIR / filename
    content = path.read_text(encoding="utf-8")

    # Strip YAML frontmatter
    if content.startswith("---"):
        _, _, content = content.split("---", 2)

    # Parse ## System and ## User sections
    system = ""
    user = ""
    current_section = None
    for line in content.splitlines():
        if line.strip() == "## System":
            current_section = "system"
        elif line.strip() == "## User":
            current_section = "user"
        elif current_section == "system":
            system += line + "\n"
        elif current_section == "user":
            user += line + "\n"

    # Substitute variables
    for key, value in variables.items():
        system = system.replace("{{" + key + "}}", value)
        user = user.replace("{{" + key + "}}", value)

    return system.strip(), user.strip()


# ---------------------------------------------------------------------------
# LLM call with retry
# ---------------------------------------------------------------------------

def call_llm(
    client: anthropic.Anthropic,
    system: str,
    user: str,
    model: str = MODEL,
    max_tokens: int = MAX_TOKENS,
) -> str:
    """Call the LLM with retry logic for rate limits and transient errors."""
    last_error = None
    for attempt in range(MAX_RETRIES):
        try:
            response = client.messages.create(
                model=model,
                max_tokens=max_tokens,
                system=system,
                messages=[{"role": "user", "content": user}],
                timeout=TIMEOUT_SECONDS,
            )
            return response.content[0].text
        except anthropic.RateLimitError as e:
            last_error = e
            wait = BACKOFF_SECONDS * (2 ** attempt)
            time.sleep(wait)
        except anthropic.APIStatusError as e:
            if e.status_code in RETRY_ON_STATUS:
                last_error = e
                wait = BACKOFF_SECONDS * (2 ** attempt)
                time.sleep(wait)
            else:
                raise
        except anthropic.APITimeoutError as e:
            last_error = e
            # Don't retry on timeout by default — surface immediately
            raise

    raise RuntimeError(f"LLM call failed after {MAX_RETRIES} attempts") from last_error


# ---------------------------------------------------------------------------
# Output validation
# ---------------------------------------------------------------------------

def validate_output(raw: str) -> dict[str, Any]:
    """Parse and validate JSON output. Raises ValueError on invalid output."""
    raw = raw.strip()
    # Strip markdown code fences if present
    if raw.startswith("```"):
        lines = raw.splitlines()
        raw = "\n".join(lines[1:-1] if lines[-1] == "```" else lines[1:])
    try:
        return json.loads(raw)
    except json.JSONDecodeError as e:
        raise ValueError(f"Output is not valid JSON: {e}\nRaw: {raw[:200]}") from e


# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------

def run(input_text: str) -> dict[str, Any]:
    """Run the pipeline on a single input.

    Args:
        input_text: The unstructured input to process.

    Returns:
        Parsed pipeline output as a dict.
    """
    client = anthropic.Anthropic()

    # Step: extract
    system, user = load_prompt("generator.prompt.md", {"input_text": input_text})
    raw = call_llm(client, system, user)
    output = validate_output(raw)

    # Add additional steps here following the same pattern:
    # system, user = load_prompt("step2.prompt.md", {"field": output["field"]})
    # raw = call_llm(client, system, user)
    # output = validate_output(raw)

    return output


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python pipeline.py '<input text>'")
        sys.exit(1)

    result = run(sys.argv[1])
    print(json.dumps(result, indent=2))