from dataclasses import dataclass
from enum import StrEnum
from pathlib import Path
from typing import Any, Generic, Literal, TypeVar

from jinja2 import Environment, FileSystemLoader
from pydantic import BaseModel, Field, field_validator
from unique_toolkit._common.validators import LMI, get_LMI_default_field
from unique_toolkit.agentic.tools.config import get_configuration_dict
from unique_toolkit.agentic.tools.schemas import BaseToolConfig
from unique_toolkit.language_model.infos import LanguageModelName
from unique_web_search.config import (
    ActivatedSearchEngine,
    DefaultSearchEngine,
)
from unique_web_search.services.search_engine import GoogleConfig

# Global template environment for the deep research tool
TEMPLATE_DIR = Path(__file__).parent / "templates"
TEMPLATE_ENV = Environment(loader=FileSystemLoader(str(TEMPLATE_DIR)))


class DeepResearchEngine(StrEnum):
    """Available deep research engines."""

    OPENAI = "OpenAI"
    UNIQUE = "Unique"


# Hardcoded configuration for the unique custom engine
@dataclass
class UniqueCustomEngineConfig:
    max_parallel_researchers: int = 5
    max_research_iterations_lead_researcher: int = 6
    max_research_iterations_sub_researcher: int = 10


RESPONSES_API_TIMEOUT_SECONDS = 3600


T = TypeVar("T", bound=DeepResearchEngine)


class BaseEngine(BaseModel, Generic[T]):
    model_config = get_configuration_dict()

    engine_type: T = Field(description="The type of engine to use for deep research")

    small_model: LMI = get_LMI_default_field(
        LanguageModelName.AZURE_GPT_4o_2024_1120,
        description="A smaller fast model for less demanding tasks",
    )

    large_model: LMI = get_LMI_default_field(
        LanguageModelName.AZURE_GPT_41_2025_0414,
        description="A larger model with longer context window and more powerful capabilities",
    )

    research_model: LMI = get_LMI_default_field(
        LanguageModelName.AZURE_GPT_5_2025_0807,
        description="The main research model to be used for conducting research",
    )

    def get_type(self) -> DeepResearchEngine:
        return DeepResearchEngine(self.engine_type)


class OpenAIEngine(BaseEngine[Literal[DeepResearchEngine.OPENAI]]):
    model_config = get_configuration_dict()

    engine_type: Literal[DeepResearchEngine.OPENAI] = Field(
        default=DeepResearchEngine.OPENAI
    )
    research_model: LMI = get_LMI_default_field(
        LanguageModelName.LITELLM_OPENAI_GPT_5,
        description="The main research model to be used for conducting research. This must be an OpenAI model hosted directly on OpenAI's servers. eg. litellm:openai-gpt-5",
    )


class WebToolsConfig(BaseModel):
    model_config = get_configuration_dict()

    enable: bool = Field(
        default=True,
        description="Allow agent to use web search tools to access the web",
    )

    search_engine: ActivatedSearchEngine = Field(  # pyright: ignore[reportInvalidTypeForm]
        default_factory=DefaultSearchEngine,  # pyright: ignore[reportArgumentType]
        description="Search Engine Configuration",
        discriminator="search_engine_name",
        title="Search Engine Configuration",
    )


class Tools(BaseModel):
    model_config = get_configuration_dict()

    web_tools: WebToolsConfig = Field(
        default=WebToolsConfig(),
        description="Configuration for web search tools",
    )
    internal_tools: bool = Field(
        default=True,
        description="Allow agent to use internal search tools access information from the knowledge base and uploaded documents",
    )

    @field_validator("web_tools", mode="before")
    @classmethod
    def handle_bool_case(cls, v: Any) -> Any:
        if isinstance(v, bool):
            if v:
                # Backward compatibility with old config behaviour
                return WebToolsConfig(
                    enable=True, search_engine=GoogleConfig(fetch_size=10)
                )
            else:
                return WebToolsConfig(enable=False)
        return v


class UniqueEngine(BaseEngine[Literal[DeepResearchEngine.UNIQUE]]):
    model_config = get_configuration_dict()

    engine_type: Literal[DeepResearchEngine.UNIQUE] = Field(
        default=DeepResearchEngine.UNIQUE
    )
    tools: Tools = Field(
        default=Tools(),
    )


class DeepResearchToolConfig(BaseToolConfig):
    engine: UniqueEngine | OpenAIEngine = Field(
        description="The deep research engine to use. Please be aware that OpenAI engine requires particular models to be used as the research model and they have different tools available.",
        default=UniqueEngine(),
        discriminator="engine_type",
    )
