from logging import getLogger
from typing import Annotated, Literal

from pydantic import BaseModel, Field, field_validator, model_validator
from pydantic.json_schema import SkipJsonSchema
from unique_toolkit._common.chunk_relevancy_sorter.config import (
    ChunkRelevancySortConfig,
)
from unique_toolkit._common.feature_flags.schema import (
    FeatureExtendedSourceSerialization,
)
from unique_toolkit._common.pydantic.rjsf_tags import RJSFMetaTag
from unique_toolkit._common.validators import LMI, get_LMI_default_field
from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
from unique_toolkit.agentic.tools.config import get_configuration_dict
from unique_toolkit.agentic.tools.schemas import BaseToolConfig
from unique_toolkit.language_model.default_language_model import DEFAULT_GPT_4o
from unique_toolkit.language_model.infos import ModelCapabilities

from unique_web_search.prompts import (
    DEFAULT_TOOL_FORMAT_INFORMATION_FOR_SYSTEM_PROMPT,
)
from unique_web_search.services.content_processing.config import (
    ContentProcessorConfig,
)
from unique_web_search.services.crawlers import (
    get_crawler_config_types_from_names,
    get_default_crawler_config,
)
from unique_web_search.services.executors.configs import (
    RefineQueryMode,
    WebSearchMode,
    WebSearchModeConfig,
    WebSearchV1Config,
    WebSearchV2Config,
    WebSearchV3Config,
    get_default_web_search_mode_config,
)
from unique_web_search.services.query_elicitation import QueryElicitationConfig
from unique_web_search.services.search_engine import (
    get_default_search_engine_config,
    get_search_engine_config_types_from_names,
)
from unique_web_search.settings import env_settings

_LOGGER = getLogger(__name__)

DEFAULT_MODEL_NAME = DEFAULT_GPT_4o

ActivatedSearchEngine = get_search_engine_config_types_from_names(
    env_settings.active_search_engines
)
DefaultSearchEngine = get_default_search_engine_config(
    env_settings.active_search_engines
)

ActivatedCrawler = get_crawler_config_types_from_names(env_settings.active_crawlers)
DefaultCrawler = get_default_crawler_config(env_settings.active_crawlers)

DEFAULT_WEB_SEARCH_MODE_CONFIG = get_default_web_search_mode_config()


class AnswerGenerationConfig(BaseModel):
    model_config = get_configuration_dict()

    limit_token_sources: int = Field(
        default=10000,
        description="Token Source Limit",
    )
    max_chunks_to_consider: int = Field(
        default=20,
        description="Token Source Limit",
    )
    number_history_interactions_included: int = Field(
        default=2,
        description="Number of history interactions included",
    )


class ExperimentalFeatures(FeatureExtendedSourceSerialization):
    query_elicitation_config: QueryElicitationConfig = Field(
        default_factory=QueryElicitationConfig,
        title="Query Review",
        description="Allow users to review and modify search queries before execution.",
    )


class WebSearchConfig(BaseToolConfig):
    language_model: LMI = get_LMI_default_field(
        DEFAULT_MODEL_NAME,
        title="Query Refinement Language Model (V1)",
        description="The AI model used to refine and improve the user's search query in V1 search mode."
        " **This setting will be moved to the V1 Search Mode Settings in a future release.**",
    )
    # TODO [UN-17641]: Remove this field in a future release.

    limit_token_sources: SkipJsonSchema[int] = Field(
        default=60_000,  # TODO: Remove SkipJsonSchema once UI (Spaces 2.0) can be configured to not include certain fields
        description="Token Source Limit",
    )
    percentage_of_input_tokens_for_sources: float = Field(
        default=0.4,
        title="Source Content Budget",
        description="Percentage of the AI model's input capacity reserved for web search results. Higher values include more search content but leave less room for conversation history.",
        ge=0.0,
        le=1.0,
    )
    language_model_max_input_tokens: SkipJsonSchema[int | None] = Field(
        default=None,
        description="Language model maximum input tokens",
    )

    web_search_active_mode: Annotated[
        WebSearchMode,
        RJSFMetaTag(
            {
                "ui:enumNames": WebSearchMode.get_enum_names(),
            }
        ),
    ] = Field(
        default=DEFAULT_WEB_SEARCH_MODE_CONFIG,
        title="Search Mode",
        description="Choose which search strategy to use. Each option is described in the selector.",
    )

    web_search_mode_config_v1: WebSearchV1Config = Field(
        default_factory=WebSearchV1Config,
        title="Search Mode V1 Settings",
        description="Settings for the basic search mode (V1), including query refinement and tool behavior.",
    )
    web_search_mode_config_v2: WebSearchV2Config = Field(
        default_factory=WebSearchV2Config,
        title="Search Mode V2 Settings",
        description="Settings for the advanced AI-planned search mode (V2), including step limits and tool behavior.",
    )
    web_search_mode_config_v3: WebSearchV3Config = Field(
        default_factory=WebSearchV3Config,
        title="Search Mode V3 Settings",
        description="Settings for V3: pre-filter search results by relevance (snippet judge) before fetching full pages.",
    )

    # Todo [UN-17655] RJSF Tags don't function properly when using union + dscriminator
    search_engine_config: ActivatedSearchEngine = Field(  # type: ignore (This type is computed at runtime so pyright is not able to infer it)
        default_factory=DefaultSearchEngine,  # type: ignore (This type is computed at runtime so pyright is not able to infer it)
        title="Search Engine",
        description="Choose and configure which search engine to use for web searches (e.g. Google, Bing, Brave).",
        discriminator="search_engine_name",
    )

    crawler_config: ActivatedCrawler = Field(  # type: ignore (This type is computed at runtime so pyright is not able to infer it)
        default_factory=DefaultCrawler,  # type: ignore (This type is computed at runtime so pyright is not able to infer it)
        title="Web Page Reader",
        description="Choose and configure how web pages are fetched and converted to readable text.",
        discriminator="crawler_type",
    )

    content_processor_config: ContentProcessorConfig = Field(
        default_factory=ContentProcessorConfig,
        title="Content Processing",
        description="Configure how fetched web page content is cleaned, shortened, and prepared before being used.",
    )

    chunk_relevancy_sort_config: ChunkRelevancySortConfig = Field(
        default_factory=ChunkRelevancySortConfig,
        title="Result Relevancy Sorting",
        description="Use AI to evaluate and sort search results by relevance to the user's question, prioritizing the most useful content.",
    )

    evaluation_check_list: list[EvaluationMetricName] = Field(
        default=[
            EvaluationMetricName.HALLUCINATION,
        ],
        title="Answer Quality Checks",
        description="Automated checks run after generating an answer to detect issues like hallucinations or irrelevant responses.",
    )

    tool_format_information_for_system_prompt: Annotated[
        str,
        RJSFMetaTag.StringWidget.textarea(
            rows=int(
                len(DEFAULT_TOOL_FORMAT_INFORMATION_FOR_SYSTEM_PROMPT.split("\n")) / 3
            )
        ),
    ] = Field(
        default=DEFAULT_TOOL_FORMAT_INFORMATION_FOR_SYSTEM_PROMPT,
        title="Source Citation Instructions",
        description="Advanced: Instructions that tell the AI how to cite web search sources in its answers.",
    )

    experimental_features: ExperimentalFeatures = Field(
        default_factory=ExperimentalFeatures,
        title="Experimental Features",
        description="Features that are still in testing and may change or be removed in future updates.",
    )

    debug: bool = Field(
        default=False,
        title="Debug Mode",
        description="When enabled, logs additional technical details for troubleshooting.",
    )

    @model_validator(mode="after")
    def disable_query_refinement_if_no_structured_output(self):
        if (
            ModelCapabilities.STRUCTURED_OUTPUT not in self.language_model.capabilities
            and self.web_search_active_mode == WebSearchMode.V1
            and self.web_search_mode_config_v1.refine_query_mode.mode
            != RefineQueryMode.DEACTIVATED
        ):
            self.web_search_mode_config_v1.refine_query_mode.mode = (
                RefineQueryMode.DEACTIVATED
            )
            _LOGGER.warning(
                "The language model does not support structured output. Query refinement is disabled."
            )
        return self

    @property
    def web_search_mode_config(self) -> WebSearchModeConfig:
        if self.web_search_active_mode == WebSearchMode.V1:
            return self.web_search_mode_config_v1
        if self.web_search_active_mode == WebSearchMode.V3:
            return self.web_search_mode_config_v3
        return self.web_search_mode_config_v2

    @field_validator("web_search_active_mode", mode="before")
    @classmethod
    def validate_web_search_active_mode(cls, v: str) -> Literal["v1", "v2", "v3"]:
        if "v3" in v.lower():
            return "v3"
        if "v2" in v.lower():
            return "v2"
        return "v1"
