from typing import Any, Literal

from pydantic import BaseModel, field_validator

from core.rag.entities import Rule
from core.rag.entities.metadata_entities import MetadataFilteringCondition
from core.rag.index_processor.constant.index_type import IndexStructureType
from core.rag.retrieval.retrieval_methods import RetrievalMethod


class RerankingModel(BaseModel):
    reranking_provider_name: str | None = None
    reranking_model_name: str | None = None


class NotionIcon(BaseModel):
    type: str
    url: str | None = None
    emoji: str | None = None


class NotionPage(BaseModel):
    page_id: str
    page_name: str
    page_icon: NotionIcon | None = None
    type: str


class NotionInfo(BaseModel):
    credential_id: str
    workspace_id: str
    pages: list[NotionPage]


class WebsiteInfo(BaseModel):
    provider: str
    job_id: str
    urls: list[str]
    only_main_content: bool = True


class FileInfo(BaseModel):
    file_ids: list[str]


class InfoList(BaseModel):
    data_source_type: Literal["upload_file", "notion_import", "website_crawl"]
    notion_info_list: list[NotionInfo] | None = None
    file_info_list: FileInfo | None = None
    website_info_list: WebsiteInfo | None = None


class DataSource(BaseModel):
    info_list: InfoList


class ProcessRule(BaseModel):
    mode: Literal["automatic", "custom", "hierarchical"]
    rules: Rule | None = None


class WeightVectorSetting(BaseModel):
    vector_weight: float
    embedding_provider_name: str
    embedding_model_name: str


class WeightKeywordSetting(BaseModel):
    keyword_weight: float


class WeightModel(BaseModel):
    weight_type: Literal["semantic_first", "keyword_first", "customized"] | None = None
    vector_setting: WeightVectorSetting | None = None
    keyword_setting: WeightKeywordSetting | None = None


class RetrievalModel(BaseModel):
    search_method: RetrievalMethod
    reranking_enable: bool
    reranking_model: RerankingModel | None = None
    reranking_mode: str | None = None
    top_k: int
    score_threshold_enabled: bool
    score_threshold: float | None = None
    weights: WeightModel | None = None
    metadata_filtering_conditions: MetadataFilteringCondition | None = None


class MetaDataConfig(BaseModel):
    doc_type: str
    doc_metadata: dict[str, Any]


class KnowledgeConfig(BaseModel):
    original_document_id: str | None = None
    duplicate: bool = True
    indexing_technique: Literal["high_quality", "economy"]
    data_source: DataSource | None = None
    process_rule: ProcessRule | None = None
    retrieval_model: RetrievalModel | None = None
    summary_index_setting: dict[str, Any] | None = None
    doc_form: str = "text_model"
    doc_language: str = "English"
    embedding_model: str | None = None
    embedding_model_provider: str | None = None
    name: str | None = None
    is_multimodal: bool = False

    @field_validator("doc_form")
    @classmethod
    def validate_doc_form(cls, value: str) -> str:
        valid_forms = [
            IndexStructureType.PARAGRAPH_INDEX,
            IndexStructureType.QA_INDEX,
            IndexStructureType.PARENT_CHILD_INDEX,
        ]
        if value not in valid_forms:
            raise ValueError("Invalid doc_form.")
        return value


class SegmentCreateArgs(BaseModel):
    content: str | None = None
    answer: str | None = None
    keywords: list[str] | None = None
    attachment_ids: list[str] | None = None


class SegmentUpdateArgs(BaseModel):
    content: str | None = None
    answer: str | None = None
    keywords: list[str] | None = None
    regenerate_child_chunks: bool = False
    enabled: bool | None = None
    attachment_ids: list[str] | None = None
    summary: str | None = None  # Summary content for summary index


class ChildChunkUpdateArgs(BaseModel):
    id: str | None = None
    content: str


class MetadataArgs(BaseModel):
    type: Literal["string", "number", "time"]
    name: str


class MetadataUpdateArgs(BaseModel):
    name: str
    value: str | int | float | None = None


class MetadataDetail(BaseModel):
    id: str
    name: str
    value: str | int | float | None = None


class DocumentMetadataOperation(BaseModel):
    document_id: str
    metadata_list: list[MetadataDetail]
    partial_update: bool = False


class MetadataOperationData(BaseModel):
    """
    Metadata operation data
    """

    operation_data: list[DocumentMetadataOperation]