from datetime import date

from core.domain.models._displayed_provider import DisplayedProvider
from core.domain.models.model_data import (
    DeprecatedModel,
    LatestModel,
    MaxTokensData,
    ModelData,
    ModelFallback,
    ModelReasoningBudget,
    QualityData,
    SpeedData,
    SpeedIndex,
)
from core.domain.models.models import Model


def xai_models() -> dict[Model, ModelData | LatestModel | DeprecatedModel]:
    return {
        Model.GROK_CODE_FAST_1: ModelData(
            display_name="Grok Code Fast (1)",
            supports_json_mode=True,
            supports_input_image=False,
            supports_input_pdf=False,
            supports_input_audio=False,
            max_tokens_data=MaxTokensData(
                max_tokens=256_000,
                source="https://docs.x.ai/docs/models/grok-code-fast-1",
            ),
            icon_url="https://workflowai.blob.core.windows.net/workflowai-public/xai.svg",
            release_date=date(2025, 9, 24),
            quality_data=QualityData(
                mmlu_pro=79.3,
                gpqa_diamond=72.7,
                source="https://artificialanalysis.ai/models/grok-code-fast-1",
            ),
            provider_name=DisplayedProvider.X_AI.value,
            supports_tool_calling=True,
            supports_structured_output=True,
            fallback=ModelFallback.default("cheap"),
            speed_data=SpeedData.from_tps(value=268),
            reasoning=ModelReasoningBudget(disabled=None, medium=None),
        ),
        Model.GROK_4_FAST: ModelData(
            display_name="Grok 4 Fast",
            supports_json_mode=True,
            supports_input_image=True,
            supports_input_pdf=True,
            supports_input_audio=False,
            max_tokens_data=MaxTokensData(
                max_tokens=2_000_000,
                source="https://docs.x.ai/docs/models/grok-4-fast-reasoning",
            ),
            icon_url="https://workflowai.blob.core.windows.net/workflowai-public/xai.svg",
            release_date=date(2025, 9, 24),
            quality_data=QualityData(
                mmlu_pro=72.8,
                gpqa_diamond=85.7,
                gpqa=85.1,
                source="https://www.vals.ai/models/grok_grok-4-fast-reasoning",
            ),
            provider_name=DisplayedProvider.X_AI.value,
            supports_tool_calling=True,
            supports_structured_output=True,
            fallback=ModelFallback.default("cheap", context_exceeded="no"),
            speed_data=SpeedData.from_tps(value=191),
            reasoning=ModelReasoningBudget(disabled=0, medium=None),
        ),
        Model.GROK_4_0709: ModelData(
            display_name="Grok 4 (0709)",
            supports_json_mode=True,
            supports_input_image=True,
            supports_input_pdf=True,
            supports_input_audio=False,
            max_tokens_data=MaxTokensData(
                max_tokens=256_000,
                source="https://docs.x.ai/docs/models/grok-4-0709",
            ),
            icon_url="https://workflowai.blob.core.windows.net/workflowai-public/xai.svg",
            release_date=date(2025, 7, 9),
            quality_data=QualityData(
                mmlu_pro=87.0,
                gpqa_diamond=88.0,
                source="https://x.com/ArtificialAnlys/status/1943166841150644622",
            ),
            provider_name=DisplayedProvider.X_AI.value,
            supports_tool_calling=True,
            supports_structured_output=True,
            fallback=ModelFallback.default("medium"),
            speed_data=SpeedData.from_tps(value=42),
        ),
        Model.GROK_3_BETA: ModelData(
            display_name="Grok 3 (beta)",
            supports_json_mode=True,
            supports_input_image=True,
            supports_input_pdf=True,
            supports_input_audio=False,
            max_tokens_data=MaxTokensData(
                max_tokens=131_072,
                source="https://docs.x.ai/docs/models#models-and-pricing",
            ),
            icon_url="https://workflowai.blob.core.windows.net/workflowai-public/xai.svg",
            release_date=date(2025, 4, 10),
            quality_data=QualityData(
                mmlu_pro=79.9,
                gpqa=73.7,
                source="https://www.vals.ai/models/grok_grok-3-fast-beta",
            ),
            speed_data=SpeedData(
                index=SpeedIndex.from_experiment(output_tokens=2300, duration_seconds=42),
            ),
            provider_name=DisplayedProvider.X_AI.value,
            supports_tool_calling=True,
            supports_structured_output=True,
            fallback=ModelFallback.default("medium"),
        ),
        Model.GROK_3_FAST_BETA: ModelData(
            display_name="Grok 3 Fast (beta)",
            supports_json_mode=True,
            supports_input_image=True,
            supports_input_pdf=True,
            supports_input_audio=False,
            max_tokens_data=MaxTokensData(
                max_tokens=131_072,
                source="https://docs.x.ai/docs/models#models-and-pricing",
            ),
            icon_url="https://workflowai.blob.core.windows.net/workflowai-public/xai.svg",
            release_date=date(2025, 4, 10),
            quality_data=QualityData(
                mmlu_pro=79.9,
                gpqa=73.7,
                source="https://www.vals.ai/models/grok_grok-3-fast-beta",
            ),
            speed_data=SpeedData(
                index=SpeedIndex.from_experiment(output_tokens=2300, duration_seconds=35),
            ),
            provider_name=DisplayedProvider.X_AI.value,
            supports_tool_calling=True,
            supports_structured_output=True,
            fallback=ModelFallback.default("medium"),
        ),
        Model.GROK_3_MINI_BETA: ModelData(
            display_name="Grok 3 Mini (beta)",
            supports_json_mode=True,
            supports_input_image=True,
            supports_input_pdf=True,
            supports_input_audio=False,
            max_tokens_data=MaxTokensData(
                max_tokens=131_072,
                source="https://docs.x.ai/docs/models#models-and-pricing",
            ),
            icon_url="https://workflowai.blob.core.windows.net/workflowai-public/xai.svg",
            release_date=date(2025, 4, 10),
            # TODO: Update the quality index
            quality_data=QualityData(
                mmlu_pro=81.4,
                gpqa=79,
                source="https://www.vals.ai/models/grok_grok-3-mini-fast-beta-high-reasoning",
            ),
            speed_data=SpeedData.from_tps(value=100),
            provider_name=DisplayedProvider.X_AI.value,
            supports_tool_calling=True,
            supports_structured_output=True,
            fallback=ModelFallback.default("cheap"),
            reasoning=ModelReasoningBudget(disabled=None, medium=None),
        ),
        Model.GROK_3_MINI_FAST_BETA: ModelData(
            display_name="Grok 3 Mini Fast (beta)",
            supports_json_mode=True,
            supports_input_image=True,
            supports_input_pdf=True,
            supports_input_audio=False,
            max_tokens_data=MaxTokensData(
                max_tokens=131_072,
                source="https://docs.x.ai/docs/models#models-and-pricing",
            ),
            icon_url="https://workflowai.blob.core.windows.net/workflowai-public/xai.svg",
            release_date=date(2025, 4, 10),
            quality_data=QualityData(
                mmlu_pro=81.4,
                gpqa=79,
                source="https://www.vals.ai/models/grok_grok-3-mini-fast-beta-high-reasoning",
            ),
            speed_data=SpeedData.from_tps(value=155),
            provider_name=DisplayedProvider.X_AI.value,
            supports_tool_calling=True,
            supports_structured_output=True,
            fallback=ModelFallback.default("cheap"),
            reasoning=ModelReasoningBudget(disabled=None, medium=None),
        ),
    }
