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

This script reads the openai.yaml feature test results and generates
python/mirascope/llm/clients/openai/model_info.py with:
- OpenAIModelId type alias containing all valid model IDs
- Model feature categorization for audio input support
- Model feature categorization for reasoning support

Logic:
- Skip models with neither completions_api nor responses_api support
- For models with at least one API:
  - Add "openai/MODEL_NAME" to the list
  - If completions_api is supported: add "openai/MODEL_NAME:completions"
  - If responses_api is supported: add "openai/MODEL_NAME:responses"
- Categorize models by audio_input support (for completions API)
- Categorize models by reasoning support (for responses API)

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

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_audio_support: set[str] = set()
    non_reasoning_models: set[str] = set()
    models_without_json_schema_support: set[str] = set()
    models_without_json_object_support: set[str] = set()

    for model_id, model_data in sorted(data.get("models", {}).items()):
        features = model_data.get("features", {})

        # Check API support
        completions_result = features.get("completions_api", {})
        responses_result = features.get("responses_api", {})
        audio_result = features.get("audio_input", {})
        reasoning_result = features.get("reasoning", {})
        structured_output_result = features.get("structured_output", {})
        json_object_result = features.get("json_object", {})

        completions_supported = completions_result.get("status") == "supported"
        responses_supported = responses_result.get("status") == "supported"

        # Skip if neither API is supported
        if not completions_supported and not responses_supported:
            continue

        # Add base model ID
        base_id = f"openai/{model_id}"
        model_ids.append(base_id)

        # Add API-specific variants
        if completions_supported:
            model_ids.append(f"{base_id}:completions")
        if responses_supported:
            model_ids.append(f"{base_id}:responses")

        # Categorize audio support (for completions API)
        audio_status = audio_result.get("status")
        if audio_status in ("not_supported", "unavailable"):
            models_without_audio_support.add(model_id)
        # If supported, skipped, or not tested, optimistic default (not in set)

        # Categorize reasoning support (for responses API)
        reasoning_status = reasoning_result.get("status")
        if reasoning_status == "not_supported":
            non_reasoning_models.add(model_id)
        # If supported, skipped, or error, optimistic default (not in set)

        # Categorize JSON schema support (structured outputs)
        structured_output_status = structured_output_result.get("status")
        if structured_output_status in ("not_supported", "unavailable"):
            models_without_json_schema_support.add(model_id)
        # If supported, skipped, or not tested, optimistic default (not in set)

        # Categorize JSON object support
        json_object_status = json_object_result.get("status")
        if json_object_status in ("not_supported", "unavailable"):
            models_without_json_object_support.add(model_id)
        # If supported, skipped, or not tested, optimistic default (not in set)

    # Generate the file content
    lines = [
        '"""OpenAI model information.',
        "",
        "This file is auto-generated by scripts/model_features/codegen_openai.py",
        'Do not edit manually - run the codegen script to update."""',
        "",
        "from typing import Literal",
        "",
        "OpenAIKnownModels = 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 OpenAI model IDs including API-specific variants."""')
    lines.append("")
    lines.append("")

    # Add models without audio support set
    lines.extend(
        [
            "",
            "MODELS_WITHOUT_AUDIO_SUPPORT: set[str] = {",
        ]
    )
    for model_id in sorted(models_without_audio_support):
        lines.append(f'    "{model_id}",')
    lines.append("}")
    lines.append('"""Models that do not support audio inputs.')
    lines.append("")
    lines.append(
        "Models not in this set are assumed to support audio (optimistic default)."
    )
    lines.append('"""')

    # Add non-reasoning model set
    lines.extend(
        [
            "",
            "NON_REASONING_MODELS: set[str] = {",
        ]
    )
    for model_id in sorted(non_reasoning_models):
        lines.append(f'    "{model_id}",')
    lines.append("}")
    lines.append('"""Models that do not support the reasoning parameter.')
    lines.append("")
    lines.append(
        "Models not in this set are assumed to support reasoning (optimistic default)."
    )
    lines.append('"""')

    # Add models without JSON schema support set
    lines.extend(
        [
            "",
            "MODELS_WITHOUT_JSON_SCHEMA_SUPPORT: set[str] = {",
        ]
    )
    for model_id in sorted(models_without_json_schema_support):
        lines.append(f'    "{model_id}",')
    lines.append("}")
    lines.append('"""Models that do not support JSON schema (structured outputs).')
    lines.append("")
    lines.append(
        "Models not in this set are assumed to support JSON schema (optimistic default)."
    )
    lines.append('"""')

    # Add models without JSON object support set
    lines.extend(
        [
            "",
            "MODELS_WITHOUT_JSON_OBJECT_SUPPORT: set[str] = {",
        ]
    )
    for model_id in sorted(models_without_json_object_support):
        lines.append(f'    "{model_id}",')
    lines.append("}")
    lines.append('"""Models that do not support JSON object mode.')
    lines.append("")
    lines.append(
        "Models not in this set are assumed to support JSON object mode (optimistic default)."
    )
    lines.append('"""')
    lines.append("")

    return "\n".join(lines)


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

    # Check input exists
    if not input_path.exists():
        print(f"Error: Input file not found: {input_path}")
        print("Run test_openai.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
    lines = content.split("\n")
    model_count = sum(1 for line in lines if line.strip().startswith('"openai/'))
    print(f"  Total model IDs: {model_count}")

    return 0


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