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

This script reads the google.yaml feature test results and generates
python/mirascope/llm/providers/google/model_info.py with:
- GoogleKnownModels type alias containing all valid model IDs
- MODELS_WITHOUT_STRUCTURED_OUTPUT_AND_TOOLS_SUPPORT set of model IDs

Logic:
- For each discovered model that supports chat functionality:
  - Add "google/MODEL_NAME" to the known models list
- Track models that don't support structured outputs with tools (negative tracking,
  assuming new models will support it)

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

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: list[str] = []
    models_without_structured_output_and_tools: set[str] = set()

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

    for model_id, model_info in sorted(models_data.items()):
        features = model_info.get("features", {})

        # Check if model supports chat
        chat_result = features.get("chat_model", {})
        chat_supported = chat_result.get("status") == "supported"

        # Skip models that don't support chat
        if not chat_supported:
            continue

        # Add model ID
        full_model_id = f"google/{model_id}"
        model_ids.append(full_model_id)

        # Check structured output with tools support
        structured_output_tools_result = features.get(
            "structured_output_with_tools", {}
        )
        supports_structured_output_and_tools = (
            structured_output_tools_result.get("status") == "supported"
        )

        if not supports_structured_output_and_tools:
            models_without_structured_output_and_tools.add(model_id)

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

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

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

    # Add models without structured output + tools support
    for i, model_id in enumerate(sorted(models_without_structured_output_and_tools)):
        comma = "," if i < len(models_without_structured_output_and_tools) - 1 else ""
        lines.append(f'    "{model_id}"{comma}')

    lines.append("}")
    lines.append(
        '"""Models that do not support structured outputs when tools are present."""'
    )
    lines.append("")

    return "\n".join(lines)


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

    # Check input exists
    if not input_path.exists():
        print(f"Error: Input file not found: {input_path}")
        print("Run test_google.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 chat models
    chat_model_count = sum(
        1
        for model_info in models_data.values()
        if model_info.get("features", {}).get("chat_model", {}).get("status")
        == "supported"
    )

    # Count models without structured output + tools support
    models_without_support_count = sum(
        1
        for model_info in models_data.values()
        if model_info.get("features", {})
        .get("structured_output_with_tools", {})
        .get("status")
        != "supported"
        and model_info.get("features", {}).get("chat_model", {}).get("status")
        == "supported"
    )

    print(f"  Discovered models: {discovered_model_count}")
    print(f"  Chat models: {chat_model_count}")
    print(
        f"  Models without structured output + tools support: {models_without_support_count}"
    )

    return 0


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