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

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

Usage:
    python scripts/update_openai_completions_model_features.py

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

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

from dotenv import load_dotenv
from openai import OpenAI
from openai.types import ChatModel

ModelStatus = Literal["supports_audio", "no_audio_support", "unavailable"]

load_dotenv()


def load_existing_model_features() -> dict[str, ModelStatus]:
    """Dynamically load existing model_features if the file exists."""
    try:
        import importlib.util
        import sys

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

        if not model_features_path.exists():
            return {}

        spec = importlib.util.spec_from_file_location(
            "model_features_module", model_features_path
        )
        if spec and spec.loader:
            module = importlib.util.module_from_spec(spec)
            sys.modules["model_features_module"] = module
            spec.loader.exec_module(module)
            return dict(module.model_features)
    except Exception:
        pass
    return {}


def get_all_openai_models() -> set[str]:
    """Get all supported model IDs from OpenAI's ChatModel type."""
    all_models = set()
    for arg in get_args(ChatModel):
        if isinstance(arg, str):
            all_models.add(arg)
    return all_models


def test_model_audio_support(client: OpenAI, model: str) -> ModelStatus:
    """Test if a model supports audio inputs.

    Returns one of: "supports_audio", "no_audio_support", "unavailable"
    """
    script_dir = Path(__file__).parent
    audio_path = (
        script_dir.parent / "tests" / "e2e" / "assets" / "audio" / "tagline.mp3"
    )
    audio_base64 = base64.b64encode(audio_path.read_bytes()).decode("utf-8")

    try:
        client.chat.completions.create(
            model=model,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "input_audio",
                            "input_audio": {"data": audio_base64, "format": "mp3"},
                        }
                    ],
                }
            ],
        )
        return "supports_audio"
    except Exception as e:
        error_str = str(e).lower()
        error_code = getattr(e, "status_code", None)

        if error_code == 404:
            return "unavailable"

        if (
            "content blocks are expected to be either text or image_url type"
            in error_str
        ):
            return "no_audio_support"

        raise e


def write_model_features(
    file_path: Path,
    model_features: dict[str, ModelStatus],
) -> None:
    """Write updated model info to model_features.py."""
    content = '''"""OpenAI ChatCompletions models categorized by audio support.

This file is auto-generated by scripts/update_openai_completions_model_features.py
Run that script to update this map when OpenAI releases new models.
"""

from typing import Literal

ModelStatus = Literal["supports_audio", "no_audio_support", "unavailable"]

MODEL_FEATURES: dict[str, ModelStatus] = {
'''

    for model in sorted(model_features.keys()):
        status = model_features[model]
        content += f'    "{model}": "{status}",\n'

    content += '''}
"""Map of model IDs to their status.

- "supports_audio": model exists and supports audio inputs
- "no_audio_support": model exists but does not support audio inputs
- "unavailable": model has been deprecated or removed
- If model not in map: assume it exists and supports audio (optimistic default)
"""
'''

    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
    model_features_file = (
        script_dir.parent
        / "mirascope"
        / "llm"
        / "clients"
        / "openai"
        / "completions"
        / "_utils"
        / "model_features.py"
    )

    model_features = load_existing_model_features()

    models_with_audio = sum(1 for v in model_features.values() if v == "supports_audio")
    models_without_audio = sum(
        1 for v in model_features.values() if v == "no_audio_support"
    )
    unavailable_models = sum(1 for v in model_features.values() if v == "unavailable")

    print("Currently categorized:")
    print(f"  Models with audio support: {models_with_audio}")
    print(f"  Models without audio support: {models_without_audio}")
    print(f"  Unavailable models: {unavailable_models}")

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

    uncategorized = all_models - set(model_features.keys())

    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_audio_support(client, model)
            model_features[model] = category
            match category:
                case "supports_audio":
                    print("✓ audio supported")
                case "no_audio_support":
                    print("✗ no audio support")
                case "unavailable":
                    print("⚠ unavailable")
                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_model_features(model_features_file, model_features)

    models_with_audio = sum(1 for v in model_features.values() if v == "supports_audio")
    models_without_audio = sum(
        1 for v in model_features.values() if v == "no_audio_support"
    )
    unavailable_models = sum(1 for v in model_features.values() if v == "unavailable")

    print(f"\n✓ Updated {model_features_file}")
    print(f"  Models with audio support: {models_with_audio}")
    print(f"  Models without audio support: {models_without_audio}")
    print(f"  Unavailable models: {unavailable_models}")


if __name__ == "__main__":
    main()
