"""Code generation for Anthropic model info from test results.

This script reads the anthropic.yaml feature test results and generates
python/mirascope/llm/providers/anthropic/model_info.py with:
- AnthropicKnownModels type alias containing all valid model IDs
- MODELS_WITHOUT_STRICT_STRUCTURED_OUTPUTS set of model IDs that don't support strict mode

Logic:
- For each discovered model with a date suffix (e.g., claude-3-5-haiku-20241022),
  generates three variants:
  1. The dated version: "anthropic/claude-3-5-haiku-20241022"
  2. The -latest alias: "anthropic/claude-3-5-haiku-latest"
  3. The base name without suffix: "anthropic/claude-3-5-haiku"
- For models with a major version suffix (e.g., claude-sonnet-4), generates:
  1. The base version: "anthropic/claude-sonnet-4"
  2. The minor version: "anthropic/claude-sonnet-4-0"
  3. The -latest alias: "anthropic/claude-sonnet-4-0-latest"
- For models without date/version suffixes, just adds the model as-is
- Models are sorted alphabetically for consistency
- Tracks which models don't support strict structured outputs (negative tracking,
  assuming new models will support it)

Usage:
    uv run python -m scripts.model_features.codegen_anthropic
"""

import re
import sys
from pathlib import Path
from typing import Any

import yaml


def load_yaml(path: Path) -> dict[str, Any]:
    """Load YAML file."""
    with open(path) as f:
        return yaml.safe_load(f)


def generate_model_info(data: dict[str, Any]) -> str:
    """Generate the model_info.py file content."""
    model_ids: set[str] = set()
    models_without_strict: set[str] = set()

    # Pattern to match date suffixes like -20241022
    date_pattern = re.compile(r"-(\d{8})$")
    # Pattern to match major version suffixes like -4 (but not -3-5 or -4-0)
    major_version_pattern = re.compile(r"-(\d+)$")

    models_data = data.get("models", {})

    for model_id, model_info in models_data.items():
        # Check if model supports strict structured outputs
        features = model_info.get("features", {})
        strict_result = features.get("strict_structured_output", {})
        supports_strict = strict_result.get("status") == "supported"

        # Check if model has a date suffix
        date_match = date_pattern.search(model_id)

        if date_match:
            # Model has a date suffix - generate variants based on base name
            base_name = model_id[: date_match.start()]  # Remove date suffix

            # Add dated version
            model_ids.add(f"anthropic/{model_id}")
            if not supports_strict:
                models_without_strict.add(f"anthropic/{model_id}")

            # Add -latest alias
            model_ids.add(f"anthropic/{base_name}-latest")
            if not supports_strict:
                models_without_strict.add(f"anthropic/{base_name}-latest")

            # Add base name without suffix
            model_ids.add(f"anthropic/{base_name}")
            if not supports_strict:
                models_without_strict.add(f"anthropic/{base_name}")

            # Check if base name has a major version suffix
            major_version_match = major_version_pattern.search(base_name)
            if major_version_match:
                # Also add minor version variant (e.g., claude-sonnet-4-0)
                minor_version = f"{base_name}-0"
                model_ids.add(f"anthropic/{minor_version}")
                if not supports_strict:
                    models_without_strict.add(f"anthropic/{minor_version}")

                # And its -latest alias (e.g., claude-sonnet-4-0-latest)
                model_ids.add(f"anthropic/{minor_version}-latest")
                if not supports_strict:
                    models_without_strict.add(f"anthropic/{minor_version}-latest")

                # And its dated version (e.g., claude-sonnet-4-0-20250929)
                dated_minor_version = f"{minor_version}-{date_match.group(1)}"
                model_ids.add(f"anthropic/{dated_minor_version}")
                if not supports_strict:
                    models_without_strict.add(f"anthropic/{dated_minor_version}")
        else:
            # Model without date suffix - add as-is
            model_ids.add(f"anthropic/{model_id}")
            if not supports_strict:
                models_without_strict.add(f"anthropic/{model_id}")

    # Sort for consistent output
    sorted_model_ids = sorted(model_ids)
    sorted_models_without_strict = sorted(models_without_strict)

    # Generate the file content
    lines = [
        '"""Anthropic model information.',
        "",
        "This file is auto-generated by scripts/model_features/codegen_anthropic.py",
        'Do not edit manually - run the codegen script to update."""',
        "",
        "from typing import Literal",
        "",
        "AnthropicKnownModels = Literal[",
    ]

    # Add model IDs as literal strings
    for i, model_id in enumerate(sorted_model_ids):
        comma = "," if i < len(sorted_model_ids) - 1 else ""
        lines.append(f'    "{model_id}"{comma}')

    lines.append("]")
    lines.append('"""Valid Anthropic model IDs."""')
    lines.append("")
    lines.append("")
    lines.append("MODELS_WITHOUT_STRICT_STRUCTURED_OUTPUTS: set[str] = {")

    # Add models without strict structured outputs
    for i, model_id in enumerate(sorted_models_without_strict):
        model_name = model_id.removeprefix("anthropic/")
        comma = "," if i < len(sorted_models_without_strict) - 1 else ""
        lines.append(f'    "{model_name}"{comma}')

    lines.append("}")
    lines.append(
        '"""Models that do not support strict structured outputs (strict mode tools)."""'
    )
    lines.append("")

    return "\n".join(lines)


def main() -> int:
    # Setup paths
    script_dir = Path(__file__).parent
    input_path = script_dir / "data" / "anthropic.yaml"
    output_path = (
        script_dir.parent.parent
        / "mirascope"
        / "llm"
        / "providers"
        / "anthropic"
        / "model_info.py"
    )

    # Check input exists
    if not input_path.exists():
        print(f"Error: Input file not found: {input_path}")
        print("Run test_anthropic.py first to generate test results")
        return 1

    # Load YAML
    print(f"Reading test results from: {input_path}")
    data = load_yaml(input_path)

    # Generate code
    print("Generating model info...")
    content = generate_model_info(data)

    # Write file
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(content)
    print(f"Generated: {output_path}")

    # Show stats
    models_data = data.get("models", {})
    discovered_model_count = len(models_data)

    # Count models without strict support (just base models, not variants)
    models_without_strict_count = sum(
        1
        for model_info in models_data.values()
        if model_info.get("features", {})
        .get("strict_structured_output", {})
        .get("status")
        != "supported"
    )

    print(f"  Discovered models: {discovered_model_count}")
    print(f"  Models without strict structured outputs: {models_without_strict_count}")

    return 0


if __name__ == "__main__":
    sys.exit(main())
