"""Utilities for fetching, filtering, and sorting OpenAI model lists.

Provides functions to retrieve the raw model catalogue from the OpenAI
API and to filter it into a curated, priority-sorted list of supported
models. Deprecated, dated-snapshot, and prohibitively expensive models
are removed by default. Reasoning models (o-series) are expanded into
``-low``, ``-medium``, and ``-high`` reasoning-effort variants.
"""

import traceback
from typing import TYPE_CHECKING, Any, List, Optional

from arbol import aprint

if TYPE_CHECKING:
    from openai import OpenAI

deprecated_models = [
    # 2025 – latest wave
    "text-moderation-007",
    "text-moderation-stable",
    "text-moderation-latest",
    "o1-preview",
    "o1-mini",
    "gpt-4.5-preview",
    # 2024 – vision & large-context models
    "gpt-4-32k",
    "gpt-4-32k-0613",
    "gpt-4-32k-0314",
    "gpt-4-vision-preview",
    "gpt-4-1106-vision-preview",
    "gpt-3.5-turbo-0613",
    "gpt-3.5-turbo-16k-0613",
    "gpt-3.5-turbo-0301",
    # 2023-24 chat & snapshot models
    "gpt-4-0314",
    # GPT-3 / InstructGPT retirement (Jan 2024)
    "text-ada-001",
    "text-babbage-001",
    "text-curie-001",
    "text-davinci-001",
    "text-davinci-002",
    "text-davinci-003",
    "ada",
    "babbage",
    "curie",
    "davinci",
    "code-davinci-002",
    "code-davinci-001",
    "text-davinci-edit-001",
    "code-davinci-edit-001",
    # First-generation embedding & search models
    "text-similarity-ada-001",
    "text-search-ada-doc-001",
    "text-search-ada-query-001",
    "code-search-ada-code-001",
    "code-search-ada-text-001",
    "text-similarity-babbage-001",
    "text-search-babbage-doc-001",
    "text-search-babbage-query-001",
    "code-search-babbage-code-001",
    "code-search-babbage-text-001",
    "text-similarity-curie-001",
    "text-search-curie-doc-001",
    "text-search-curie-query-001",
    "text-similarity-davinci-001",
    "text-search-davinci-doc-001",
    "text-search-davinci-query-001",
    # Codex family sunset (Mar 2023)
    "code-cushman-002",
    "code-cushman-001",
]


def _get_raw_openai_model_list(client: "OpenAI"):
    """Fetch the raw model list from the OpenAI API.

    Parameters
    ----------
    client : OpenAI
        The OpenAI client instance.

    Returns
    -------
    list
        Raw list of model objects from the API.
    """
    from openai import OpenAI

    # Explicit typing:
    client: OpenAI = client

    # Raw model list:
    models = client.models.list().data

    # Return the list of models:
    return models


def get_openai_model_list(
    raw_model_list: List[Any],
    included: Optional[List[str]] = None,
    excluded: Optional[List[str]] = None,
    exclude_dated_models: bool = True,
    allow_prohibitively_expensive_models: bool = False,
    verbose: bool = False,
) -> List[str]:
    """
    Get the list of all OpenAI ChatGPT models.

    Parameters
    ----------
    raw_model_list : List[Any]
        Raw list of model objects from the OpenAI API.
    included : Optional[List[str]]
        Substrings to include. A model must contain at least one to be kept.
        If None, uses a default set of common model prefixes.
    excluded : Optional[List[str]]
        Substrings to exclude. A model containing any of these is removed.
        If None, uses a default exclusion list.
    exclude_dated_models : bool
        If True, removes models with date suffixes (e.g., ``gpt-4-0613``).
    allow_prohibitively_expensive_models : bool
        If True, includes very expensive models like o1-pro.
    verbose : bool
        If True, prints inclusion/exclusion decisions.

    Returns
    -------
    list[str]
        List of models.

    """

    try:
        if included is None:
            # Base filtering of models:
            included = [
                "dall-e",
                "audio",
                "gpt",
                "o1",
                "o3",
                "o4",
                "o5",
                "text-embedding",
                "whisper",
                "ada-002",
                "tts",
            ]

        if excluded is None:
            # Exclude models that are not supported by the API:
            excluded = ["ada-002"]

        if not allow_prohibitively_expensive_models:
            # Exclude prohibitively expensive models:
            excluded += ["o1-pro", "o3-pro", "o4-pro", "o5-pro"]

        # Convert to list of model ids:
        models = [model.id for model in raw_model_list]

        if exclude_dated_models:
            models = _remove_dated_models(models)

        # Model list to populate:
        model_list = []

        # Goes through models and populate the list:
        for model in models:

            # only keep models that match the filter:
            if not included or any(f in model for f in included):
                model_list.append(model)
                if verbose:
                    aprint(f"Included: {model}")

            # Only keep models that do not match the excluded:
            if excluded and any(e in model for e in excluded):
                model_list.remove(model)
                if verbose:
                    aprint(f"Excluded: {model}")

            # Remove deprecated models:
            if model in deprecated_models:
                model_list.remove(model)
                if verbose:
                    aprint(f"Deprecated: {model}")

        # Remove duplicates:
        model_list = list(set(model_list))

        # Actual sorting:
        sorted_model_list = sorted(model_list, key=model_key, reverse=True)

        # List of reasoning models filters:
        reasoning_models_filters = ["o1", "o3", "o4", "o5"]

        # List of reasoning models that contain the filters:
        reasoning_models = [
            model
            for model in sorted_model_list
            if any(filter_ in model for filter_ in reasoning_models_filters)
        ]

        # Replace each reasoning model 'X' with its three variants:  X-low, X-mid, X-high:
        for reasoning_model in reasoning_models:
            if reasoning_model in sorted_model_list:
                # Find index of reasoning model:
                index = sorted_model_list.index(reasoning_model)
                # Insert three variants:
                sorted_model_list.insert(index + 1, f"{reasoning_model}-high")
                sorted_model_list.insert(index + 2, f"{reasoning_model}-medium")
                sorted_model_list.insert(index + 3, f"{reasoning_model}-low")
                # Remove the original reasoning model:
                sorted_model_list.remove(reasoning_model)

        return sorted_model_list

    except Exception as e:
        # Error message:
        aprint(
            f"Error: {type(e).__name__} with message: '{str(e)}' occured while trying to get the list of OpenAI models. "
        )
        # print stacktrace:
        traceback.print_exc()

        return []


def _remove_dated_models(models):
    """Remove models with date-based suffixes from the model list.

    Parameters
    ----------
    models : list
        List of model name strings.

    Returns
    -------
    list
        Filtered list with dated model variants removed.
    """
    nodate_models = []
    for model in models:

        # Remove '-preview' from name:
        model_ = model.replace("-preview", "")

        if "-" in model_:
            parts = model_.split("-")

            if len(parts) >= 4:

                year = parts[-3]
                month = parts[-2]
                day = parts[-1]

                # Check if the date is valid:
                if len(year) == 4 and len(month) == 2 and len(day) == 2:
                    continue

            if len(parts) >= 2:
                if parts[-1].isdigit() and len(parts[-1]) == 4:
                    continue

        nodate_models.append(model)
    return nodate_models


def model_key(model):
    """Compute a priority score for sorting OpenAI models.

    Higher scores indicate more capable/newer models. Used as the
    sort key to put the best models first.

    Parameters
    ----------
    model : str
        The model name.

    Returns
    -------
    int
        Priority score for sorting (higher is better).
    """
    score = 0

    # GPT-5 series (latest)
    if "gpt-5.3" in model:
        score += 165
    elif "gpt-5.2" in model:
        score += 160
    elif "gpt-5.1" in model:
        score += 155
    elif "gpt-5" in model:
        score += 150
    # GPT-4.5 series
    elif "gpt-4.5" in model:
        score += 130
    # O-series reasoning models
    elif "o5" in model:
        score += 130
    elif "o4" in model:
        score += 125
    elif "o3" in model:
        score += 120
    elif "o1" in model:
        score += 110
    # GPT-4.1 series
    elif "gpt-4.1" in model:
        score += 105
    # GPT-4o series
    elif "gpt-4o" in model:
        score += 100
    # Legacy models
    elif "gpt-4" in model:
        score += 90
    elif "gpt-3.5" in model:
        score += 80
    elif "gpt-3" in model:
        score += 70

    if "mini" in model:
        score -= 5

    if "nano" in model:
        score -= 7

    if "pro" in model:
        score += 4

    if "preview" in model:
        score -= 2

    return score
