#!/usr/bin/env python3
"""Script to test and categorize OpenAI Responses models by reasoning support.

This script:
1. Loads all models from openai.types.shared.responses_model
2. Tests each uncategorized model with a simple API call using reasoning parameter
3. Updates model_features.py with the results

Usage:
    python scripts/update_openai_responses_models.py

Requires OPENAI_API_KEY environment variable (loaded from .env if present).
"""

import os
from pathlib import Path
from typing import Literal, get_args

from dotenv import load_dotenv
from openai import OpenAI
from openai.types.shared import ResponsesModel

from mirascope.llm.clients.openai.responses._utils.model_features import (
    NO_RESPONSES_API_SUPPORT_MODELS,
    NON_EXISTENT_MODELS,
    NON_REASONING_MODELS,
    REASONING_MODELS,
)

load_dotenv()


def get_all_openai_models() -> set[str]:
    """Get all supported model IDs from OpenAI's ResponsesModel type."""
    all_models = set()
    for arg in get_args(ResponsesModel):
        # Skip the `str` type itself
        if arg is str:
            continue
        # Extract values from Literal types
        if hasattr(arg, "__args__"):
            all_models.update(arg.__args__)
    return all_models


ModelCategory = Literal[
    "reasoning", "non_reasoning", "does_not_exist", "no_responses_api_support"
]


def test_model_reasoning_support(client: OpenAI, model: str) -> ModelCategory:
    """Test if a model supports the reasoning parameter.

    Returns one of: "reasoning", "non_reasoning", "does_not_exist", or "no_responses_api_support"
    """
    try:
        client.responses.create(
            model=model,
            input=[{"role": "user", "content": "Say hello"}],
            reasoning={"effort": "medium"},
        )
        return "reasoning"
    except Exception as e:
        error_str = str(e).lower()
        if "is not supported with the responses api." in error_str:
            return "no_responses_api_support"
        if "does not exist" in error_str or "model not found" in error_str:
            return "does_not_exist"
        if (
            "unsupported parameter: 'reasoning.effort' is not supported with this model"
            in error_str
        ):
            return "non_reasoning"
        raise e


def write_categorization(
    file_path: Path,
    reasoning_models: set[str],
    non_reasoning_models: set[str],
    non_existent_models: set[str],
    no_responses_api_support_models: set[str],
) -> None:
    """Write updated categorization to model_features.py."""
    content = '''"""OpenAI Responses models categorized by reasoning support.

This file is auto-generated by scripts/update_openai_responses_model_features.py
Run that script to update these sets when OpenAI releases new models.
"""

REASONING_MODELS: set[str] = {
'''

    for model in sorted(reasoning_models):
        content += f'    "{model}",\n'

    content += '''}
"""Models that have been tested and confirmed to support the reasoning parameter."""

NON_REASONING_MODELS: set[str] = {
'''

    for model in sorted(non_reasoning_models):
        content += f'    "{model}",\n'

    content += '''}
"""Models that have been tested and confirmed to NOT support the reasoning parameter."""

NON_EXISTENT_MODELS: set[str] = {
'''

    for model in sorted(non_existent_models):
        content += f'    "{model}",\n'

    content += '''}
"""Models that are listed in OpenAI's types but no longer exist in their API."""

NO_RESPONSES_API_SUPPORT_MODELS: set[str] = {
'''

    for model in sorted(no_responses_api_support_models):
        content += f'    "{model}",\n'

    content += '''}
"""Models that do not support the Responses API."""
'''

    file_path.write_text(content)


def main() -> None:
    """Main script execution."""
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        print("Error: OPENAI_API_KEY environment variable not set")
        print("Please set it in your .env file or environment")
        return

    client = OpenAI(api_key=api_key)

    script_dir = Path(__file__).parent
    reasoning_models_file = (
        script_dir.parent
        / "mirascope"
        / "llm"
        / "clients"
        / "openai"
        / "responses"
        / "_utils"
        / "model_features.py"
    )

    reasoning_models = set(REASONING_MODELS)
    non_reasoning_models = set(NON_REASONING_MODELS)
    non_existent_models = set(NON_EXISTENT_MODELS)
    no_responses_api_support_models = set(NO_RESPONSES_API_SUPPORT_MODELS)

    print("Currently categorized:")
    print(f"  Reasoning models: {len(reasoning_models)}")
    print(f"  Non-reasoning models: {len(non_reasoning_models)}")
    print(f"  Non-existent models: {len(non_existent_models)}")
    print(f"  No Responses API support: {len(no_responses_api_support_models)}")

    all_models = get_all_openai_models()
    print(f"\nTotal OpenAI Responses models: {len(all_models)}")

    categorized = (
        reasoning_models
        | non_reasoning_models
        | non_existent_models
        | no_responses_api_support_models
    )
    uncategorized = all_models - categorized

    if not uncategorized:
        print("\nAll models already categorized!")
        return

    print(f"\nUncategorized models: {len(uncategorized)}")
    print("Testing models...")

    for model in sorted(uncategorized):
        print(f"  Testing {model}...", end=" ")
        try:
            category = test_model_reasoning_support(client, model)
            match category:
                case "reasoning":
                    reasoning_models.add(model)
                    print("✓ reasoning supported")
                case "non_reasoning":
                    non_reasoning_models.add(model)
                    print("✗ no reasoning support")
                case "does_not_exist":
                    non_existent_models.add(model)
                    print("⚠ does not exist")
                case "no_responses_api_support":
                    no_responses_api_support_models.add(model)
                    print("⚠ no Responses API support")
                case _:
                    raise RuntimeError(f"Unexpected category: {category}")
        except KeyboardInterrupt:
            print("\n\nInterrupted by user. Saving progress...")
            break
        except Exception as e:
            print(f"⚠ error (skipping): {e}")
            continue

    write_categorization(
        reasoning_models_file,
        reasoning_models,
        non_reasoning_models,
        non_existent_models,
        no_responses_api_support_models,
    )

    print(f"\n✓ Updated {reasoning_models_file}")
    print(f"  Reasoning models: {len(reasoning_models)}")
    print(f"  Non-reasoning models: {len(non_reasoning_models)}")
    print(f"  Non-existent models: {len(non_existent_models)}")
    print(f"  No Responses API support: {len(no_responses_api_support_models)}")


if __name__ == "__main__":
    main()
