from enum import StrEnum, auto
from typing import Any, Literal

from pydantic import BaseModel, Field

from core.rag.data_post_processor.data_post_processor import RerankingModelDict, WeightsDict
from core.rag.entities import MetadataFilteringCondition
from graphon.file import FileUploadConfig
from graphon.model_runtime.entities.llm_entities import LLMMode
from graphon.model_runtime.entities.message_entities import PromptMessageRole
from graphon.variables.input_entities import VariableEntity as WorkflowVariableEntity
from models.model import AppMode


class ModelConfigEntity(BaseModel):
    """
    Model Config Entity.
    """

    provider: str
    model: str
    mode: str | None = None
    parameters: dict[str, Any] = Field(default_factory=dict)
    stop: list[str] = Field(default_factory=list)


class AdvancedChatMessageEntity(BaseModel):
    """
    Advanced Chat Message Entity.
    """

    text: str
    role: PromptMessageRole


class AdvancedChatPromptTemplateEntity(BaseModel):
    """
    Advanced Chat Prompt Template Entity.
    """

    messages: list[AdvancedChatMessageEntity]


class AdvancedCompletionPromptTemplateEntity(BaseModel):
    """
    Advanced Completion Prompt Template Entity.
    """

    class RolePrefixEntity(BaseModel):
        """
        Role Prefix Entity.
        """

        user: str
        assistant: str

    prompt: str
    role_prefix: RolePrefixEntity | None = None


class PromptTemplateEntity(BaseModel):
    """
    Prompt Template Entity.
    """

    class PromptType(StrEnum):
        """
        Prompt Type.
        'simple', 'advanced'
        """

        SIMPLE = auto()
        ADVANCED = auto()

        @classmethod
        def value_of(cls, value: str):
            """
            Get value of given mode.

            :param value: mode value
            :return: mode
            """
            for mode in cls:
                if mode.value == value:
                    return mode
            raise ValueError(f"invalid prompt type value {value}")

    prompt_type: PromptType
    simple_prompt_template: str | None = None
    advanced_chat_prompt_template: AdvancedChatPromptTemplateEntity | None = None
    advanced_completion_prompt_template: AdvancedCompletionPromptTemplateEntity | None = None


class RagPipelineVariableEntity(WorkflowVariableEntity):
    """
    Rag Pipeline Variable Entity.
    """

    tooltips: str | None = None
    placeholder: str | None = None
    belong_to_node_id: str


class ExternalDataVariableEntity(BaseModel):
    """
    External Data Variable Entity.
    """

    variable: str
    type: str
    config: dict[str, Any] = Field(default_factory=dict)


class ModelConfig(BaseModel):
    provider: str
    name: str
    mode: LLMMode
    completion_params: dict[str, Any] = Field(default_factory=dict)


class DatasetRetrieveConfigEntity(BaseModel):
    """
    Dataset Retrieve Config Entity.
    """

    class RetrieveStrategy(StrEnum):
        """
        Dataset Retrieve Strategy.
        'single' or 'multiple'
        """

        SINGLE = auto()
        MULTIPLE = auto()

        @classmethod
        def value_of(cls, value: str):
            """
            Get value of given mode.

            :param value: mode value
            :return: mode
            """
            for mode in cls:
                if mode.value == value:
                    return mode
            raise ValueError(f"invalid retrieve strategy value {value}")

    query_variable: str | None = None  # Only when app mode is completion

    retrieve_strategy: RetrieveStrategy
    top_k: int | None = None
    score_threshold: float | None = 0.0
    rerank_mode: str | None = "reranking_model"
    reranking_model: RerankingModelDict | None = None
    weights: WeightsDict | None = None
    reranking_enabled: bool | None = True
    metadata_filtering_mode: Literal["disabled", "automatic", "manual"] | None = "disabled"
    metadata_model_config: ModelConfig | None = None
    metadata_filtering_conditions: MetadataFilteringCondition | None = None


class DatasetEntity(BaseModel):
    """
    Dataset Config Entity.
    """

    dataset_ids: list[str]
    retrieve_config: DatasetRetrieveConfigEntity


class SensitiveWordAvoidanceEntity(BaseModel):
    """
    Sensitive Word Avoidance Entity.
    """

    type: str
    config: dict[str, Any] = Field(default_factory=dict)


class TextToSpeechEntity(BaseModel):
    """
    Sensitive Word Avoidance Entity.
    """

    enabled: bool
    voice: str | None = None
    language: str | None = None


class TracingConfigEntity(BaseModel):
    """
    Tracing Config Entity.
    """

    enabled: bool
    tracing_provider: str


class AppAdditionalFeatures(BaseModel):
    file_upload: FileUploadConfig | None = None
    opening_statement: str | None = None
    suggested_questions: list[str] = []
    suggested_questions_after_answer: bool = False
    show_retrieve_source: bool = False
    more_like_this: bool = False
    speech_to_text: bool = False
    text_to_speech: TextToSpeechEntity | None = None
    trace_config: TracingConfigEntity | None = None


class AppConfig(BaseModel):
    """
    Application Config Entity.
    """

    tenant_id: str
    app_id: str
    app_mode: AppMode
    additional_features: AppAdditionalFeatures | None = None
    variables: list[WorkflowVariableEntity] = []
    sensitive_word_avoidance: SensitiveWordAvoidanceEntity | None = None


class EasyUIBasedAppModelConfigFrom(StrEnum):
    """
    App Model Config From.
    """

    ARGS = auto()
    APP_LATEST_CONFIG = "app-latest-config"
    CONVERSATION_SPECIFIC_CONFIG = "conversation-specific-config"


class EasyUIBasedAppConfig(AppConfig):
    """
    Easy UI Based App Config Entity.
    """

    app_model_config_from: EasyUIBasedAppModelConfigFrom
    app_model_config_id: str
    app_model_config_dict: dict[str, Any]
    model: ModelConfigEntity
    prompt_template: PromptTemplateEntity
    dataset: DatasetEntity | None = None
    external_data_variables: list[ExternalDataVariableEntity] = []


class WorkflowUIBasedAppConfig(AppConfig):
    """
    Workflow UI Based App Config Entity.
    """

    workflow_id: str