import base64
import enum
import hashlib
import hmac
import json
import logging
import os
import pickle
import re
import time
from datetime import datetime
from json import JSONDecodeError
from typing import Any, ClassVar, TypedDict, cast
from uuid import uuid4

import sqlalchemy as sa
from sqlalchemy import DateTime, String, func, select
from sqlalchemy.orm import Mapped, Session, mapped_column

from configs import dify_config
from core.rag.entities import ParentMode, Rule
from core.rag.index_processor.constant.built_in_field import BuiltInField, MetadataDataSource
from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from core.rag.index_processor.constant.query_type import QueryType
from core.rag.retrieval.retrieval_methods import RetrievalMethod
from core.tools.signature import sign_upload_file_preview_url
from extensions.ext_storage import storage
from libs.uuid_utils import uuidv7

from .account import Account
from .base import Base, TypeBase
from .engine import db
from .enums import (
    CollectionBindingType,
    CreatorUserRole,
    DatasetMetadataType,
    DatasetQuerySource,
    DatasetRuntimeMode,
    DataSourceType,
    DocumentCreatedFrom,
    DocumentDocType,
    IndexingStatus,
    ProcessRuleMode,
    SegmentStatus,
    SegmentType,
    SummaryStatus,
    TidbAuthBindingStatus,
)
from .model import App, Tag, TagBinding, UploadFile
from .types import AdjustedJSON, BinaryData, EnumText, LongText, StringUUID, adjusted_json_index

logger = logging.getLogger(__name__)


class PreProcessingRuleItem(TypedDict):
    id: str
    enabled: bool


class SegmentationConfig(TypedDict):
    delimiter: str
    max_tokens: int
    chunk_overlap: int


class AutomaticRulesConfig(TypedDict):
    pre_processing_rules: list[PreProcessingRuleItem]
    segmentation: SegmentationConfig


class ProcessRuleDict(TypedDict):
    id: str
    dataset_id: str
    mode: str
    rules: dict[str, Any] | None


class DocMetadataDetailItem(TypedDict):
    id: str
    name: str
    type: str
    value: Any


class AttachmentItem(TypedDict):
    id: str
    name: str
    size: int
    extension: str
    mime_type: str
    source_url: str


class DatasetBindingItem(TypedDict):
    id: str
    name: str


class ExternalKnowledgeApiDict(TypedDict):
    id: str
    tenant_id: str
    name: str
    description: str
    settings: dict[str, Any] | None
    dataset_bindings: list[DatasetBindingItem]
    created_by: str
    created_at: str


class DocumentDict(TypedDict):
    id: str
    tenant_id: str
    dataset_id: str
    position: int
    data_source_type: str
    data_source_info: str | None
    dataset_process_rule_id: str | None
    batch: str
    name: str
    created_from: str
    created_by: str
    created_api_request_id: str | None
    created_at: datetime
    processing_started_at: datetime | None
    file_id: str | None
    word_count: int | None
    parsing_completed_at: datetime | None
    cleaning_completed_at: datetime | None
    splitting_completed_at: datetime | None
    tokens: int | None
    indexing_latency: float | None
    completed_at: datetime | None
    is_paused: bool | None
    paused_by: str | None
    paused_at: datetime | None
    error: str | None
    stopped_at: datetime | None
    indexing_status: str
    enabled: bool
    disabled_at: datetime | None
    disabled_by: str | None
    archived: bool
    archived_reason: str | None
    archived_by: str | None
    archived_at: datetime | None
    updated_at: datetime
    doc_type: str | None
    doc_metadata: Any
    doc_form: IndexStructureType
    doc_language: str | None
    display_status: str | None
    data_source_info_dict: dict[str, Any]
    average_segment_length: int
    dataset_process_rule: ProcessRuleDict | None
    dataset: None
    segment_count: int | None
    hit_count: int | None


class DatasetPermissionEnum(enum.StrEnum):
    ONLY_ME = "only_me"
    ALL_TEAM = "all_team_members"
    PARTIAL_TEAM = "partial_members"


class Dataset(Base):
    __tablename__ = "datasets"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="dataset_pkey"),
        sa.Index("dataset_tenant_idx", "tenant_id"),
        adjusted_json_index("retrieval_model_idx", "retrieval_model"),
    )

    INDEXING_TECHNIQUE_LIST = ["high_quality", "economy", None]
    PROVIDER_LIST = ["vendor", "external", None]
    DOC_FORM_LIST = [member.value for member in IndexStructureType]

    id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()))
    tenant_id: Mapped[str] = mapped_column(StringUUID)
    name: Mapped[str] = mapped_column(String(255))
    description = mapped_column(LongText, nullable=True)
    provider: Mapped[str] = mapped_column(String(255), server_default=sa.text("'vendor'"))
    permission: Mapped[DatasetPermissionEnum] = mapped_column(
        EnumText(DatasetPermissionEnum, length=255),
        server_default=sa.text("'only_me'"),
        default=DatasetPermissionEnum.ONLY_ME,
    )
    data_source_type = mapped_column(EnumText(DataSourceType, length=255))
    indexing_technique: Mapped[IndexTechniqueType | None] = mapped_column(EnumText(IndexTechniqueType, length=255))
    index_struct = mapped_column(LongText, nullable=True)
    created_by = mapped_column(StringUUID, nullable=False)
    created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
    updated_by = mapped_column(StringUUID, nullable=True)
    updated_at = mapped_column(
        sa.DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp()
    )
    embedding_model = mapped_column(sa.String(255), nullable=True)
    embedding_model_provider = mapped_column(sa.String(255), nullable=True)
    keyword_number = mapped_column(sa.Integer, nullable=True, server_default=sa.text("10"))
    collection_binding_id = mapped_column(StringUUID, nullable=True)
    retrieval_model = mapped_column(AdjustedJSON, nullable=True)
    summary_index_setting = mapped_column(AdjustedJSON, nullable=True)
    built_in_field_enabled = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"))
    icon_info = mapped_column(AdjustedJSON, nullable=True)
    runtime_mode = mapped_column(
        EnumText(DatasetRuntimeMode, length=255), nullable=True, server_default=sa.text("'general'")
    )
    pipeline_id = mapped_column(StringUUID, nullable=True)
    chunk_structure = mapped_column(sa.String(255), nullable=True)
    enable_api = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("true"))
    is_multimodal = mapped_column(sa.Boolean, default=False, nullable=False, server_default=db.text("false"))

    @property
    def total_documents(self):
        return db.session.scalar(select(func.count(Document.id)).where(Document.dataset_id == self.id)) or 0

    @property
    def total_available_documents(self):
        return (
            db.session.scalar(
                select(func.count(Document.id)).where(
                    Document.dataset_id == self.id,
                    Document.indexing_status == "completed",
                    Document.enabled == True,
                    Document.archived == False,
                )
            )
            or 0
        )

    @property
    def dataset_keyword_table(self):
        return db.session.scalar(select(DatasetKeywordTable).where(DatasetKeywordTable.dataset_id == self.id))

    @property
    def index_struct_dict(self):
        return json.loads(self.index_struct) if self.index_struct else None

    @property
    def external_retrieval_model(self):
        default_retrieval_model = {
            "top_k": 2,
            "score_threshold": 0.0,
        }
        return self.retrieval_model or default_retrieval_model

    @property
    def created_by_account(self):
        return db.session.get(Account, self.created_by)

    @property
    def author_name(self) -> str | None:
        account = db.session.get(Account, self.created_by)
        if account:
            return account.name
        return None

    @property
    def latest_process_rule(self):
        return db.session.scalar(
            select(DatasetProcessRule)
            .where(DatasetProcessRule.dataset_id == self.id)
            .order_by(DatasetProcessRule.created_at.desc())
            .limit(1)
        )

    @property
    def app_count(self):
        return (
            db.session.scalar(
                select(func.count(AppDatasetJoin.id)).where(
                    AppDatasetJoin.dataset_id == self.id, App.id == AppDatasetJoin.app_id
                )
            )
            or 0
        )

    @property
    def document_count(self):
        return db.session.scalar(select(func.count(Document.id)).where(Document.dataset_id == self.id)) or 0

    @property
    def available_document_count(self):
        return (
            db.session.scalar(
                select(func.count(Document.id)).where(
                    Document.dataset_id == self.id,
                    Document.indexing_status == "completed",
                    Document.enabled == True,
                    Document.archived == False,
                )
            )
            or 0
        )

    @property
    def available_segment_count(self):
        return (
            db.session.scalar(
                select(func.count(DocumentSegment.id)).where(
                    DocumentSegment.dataset_id == self.id,
                    DocumentSegment.status == "completed",
                    DocumentSegment.enabled == True,
                )
            )
            or 0
        )

    @property
    def word_count(self):
        return db.session.scalar(
            select(func.coalesce(func.sum(Document.word_count), 0)).where(Document.dataset_id == self.id)
        )

    @property
    def doc_form(self) -> str | None:
        if self.chunk_structure:
            return self.chunk_structure
        document = db.session.scalar(select(Document).where(Document.dataset_id == self.id).limit(1))
        if document:
            return document.doc_form
        return None

    @property
    def retrieval_model_dict(self):
        default_retrieval_model = {
            "search_method": RetrievalMethod.SEMANTIC_SEARCH,
            "reranking_enable": False,
            "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""},
            "top_k": 2,
            "score_threshold_enabled": False,
        }
        return self.retrieval_model or default_retrieval_model

    @property
    def tags(self):
        tags = db.session.scalars(
            select(Tag)
            .join(TagBinding, Tag.id == TagBinding.tag_id)
            .where(
                TagBinding.target_id == self.id,
                TagBinding.tenant_id == self.tenant_id,
                Tag.tenant_id == self.tenant_id,
                Tag.type == "knowledge",
            )
        ).all()

        return tags or []

    @property
    def external_knowledge_info(self):
        if self.provider != "external":
            return None
        external_knowledge_binding = db.session.scalar(
            select(ExternalKnowledgeBindings).where(
                ExternalKnowledgeBindings.dataset_id == self.id,
                ExternalKnowledgeBindings.tenant_id == self.tenant_id,
            )
        )
        if not external_knowledge_binding:
            return None
        external_knowledge_api = db.session.scalar(
            select(ExternalKnowledgeApis).where(
                ExternalKnowledgeApis.id == external_knowledge_binding.external_knowledge_api_id,
                ExternalKnowledgeApis.tenant_id == self.tenant_id,
            )
        )
        if external_knowledge_api is None or external_knowledge_api.settings is None:
            return None
        return {
            "external_knowledge_id": external_knowledge_binding.external_knowledge_id,
            "external_knowledge_api_id": external_knowledge_api.id,
            "external_knowledge_api_name": external_knowledge_api.name,
            "external_knowledge_api_endpoint": json.loads(external_knowledge_api.settings).get("endpoint", ""),
        }

    @property
    def is_published(self):
        if self.pipeline_id:
            pipeline = db.session.scalar(select(Pipeline).where(Pipeline.id == self.pipeline_id))
            if pipeline:
                return pipeline.is_published
        return False

    @property
    def doc_metadata(self):
        dataset_metadatas = db.session.scalars(
            select(DatasetMetadata).where(DatasetMetadata.dataset_id == self.id)
        ).all()

        doc_metadata = [
            {
                "id": dataset_metadata.id,
                "name": dataset_metadata.name,
                "type": dataset_metadata.type,
            }
            for dataset_metadata in dataset_metadatas
        ]
        if self.built_in_field_enabled:
            doc_metadata.append(
                {
                    "id": "built-in",
                    "name": BuiltInField.document_name,
                    "type": "string",
                }
            )
            doc_metadata.append(
                {
                    "id": "built-in",
                    "name": BuiltInField.uploader,
                    "type": "string",
                }
            )
            doc_metadata.append(
                {
                    "id": "built-in",
                    "name": BuiltInField.upload_date,
                    "type": "time",
                }
            )
            doc_metadata.append(
                {
                    "id": "built-in",
                    "name": BuiltInField.last_update_date,
                    "type": "time",
                }
            )
            doc_metadata.append(
                {
                    "id": "built-in",
                    "name": BuiltInField.source,
                    "type": "string",
                }
            )
        return doc_metadata

    @staticmethod
    def gen_collection_name_by_id(dataset_id: str) -> str:
        normalized_dataset_id = dataset_id.replace("-", "_")
        return f"{dify_config.VECTOR_INDEX_NAME_PREFIX}_{normalized_dataset_id}_Node"


class DatasetProcessRule(TypeBase):
    __tablename__ = "dataset_process_rules"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="dataset_process_rule_pkey"),
        sa.Index("dataset_process_rule_dataset_id_idx", "dataset_id"),
    )

    id: Mapped[str] = mapped_column(StringUUID, nullable=False, default_factory=lambda: str(uuid4()), init=False)
    dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    mode: Mapped[ProcessRuleMode] = mapped_column(
        EnumText(ProcessRuleMode, length=255), nullable=False, server_default=sa.text("'automatic'")
    )
    rules: Mapped[str | None] = mapped_column(LongText, nullable=True)
    created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )

    MODES = ["automatic", "custom", "hierarchical"]
    PRE_PROCESSING_RULES = ["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"]
    AUTOMATIC_RULES: ClassVar[AutomaticRulesConfig] = {
        "pre_processing_rules": [
            {"id": "remove_extra_spaces", "enabled": True},
            {"id": "remove_urls_emails", "enabled": False},
        ],
        "segmentation": {"delimiter": "\n", "max_tokens": 500, "chunk_overlap": 50},
    }

    def to_dict(self) -> ProcessRuleDict:
        return {
            "id": self.id,
            "dataset_id": self.dataset_id,
            "mode": self.mode,
            "rules": self.rules_dict,
        }

    @property
    def rules_dict(self) -> dict[str, Any] | None:
        try:
            return json.loads(self.rules) if self.rules else None
        except JSONDecodeError:
            return None


class Document(Base):
    __tablename__ = "documents"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="document_pkey"),
        sa.Index("document_dataset_id_idx", "dataset_id"),
        sa.Index("document_is_paused_idx", "is_paused"),
        sa.Index("document_tenant_idx", "tenant_id"),
        adjusted_json_index("document_metadata_idx", "doc_metadata"),
    )

    # initial fields
    id = mapped_column(StringUUID, nullable=False, default=lambda: str(uuid4()))
    tenant_id = mapped_column(StringUUID, nullable=False)
    dataset_id = mapped_column(StringUUID, nullable=False)
    position: Mapped[int] = mapped_column(sa.Integer, nullable=False)
    data_source_type: Mapped[str] = mapped_column(EnumText(DataSourceType, length=255), nullable=False)
    data_source_info = mapped_column(LongText, nullable=True)
    dataset_process_rule_id = mapped_column(StringUUID, nullable=True)
    batch: Mapped[str] = mapped_column(String(255), nullable=False)
    name: Mapped[str] = mapped_column(String(255), nullable=False)
    created_from: Mapped[str] = mapped_column(EnumText(DocumentCreatedFrom, length=255), nullable=False)
    created_by = mapped_column(StringUUID, nullable=False)
    created_api_request_id = mapped_column(StringUUID, nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())

    # start processing
    processing_started_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)

    # parsing
    file_id = mapped_column(LongText, nullable=True)
    word_count: Mapped[int | None] = mapped_column(sa.Integer, nullable=True)  # TODO: make this not nullable
    parsing_completed_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)

    # cleaning
    cleaning_completed_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)

    # split
    splitting_completed_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)

    # indexing
    tokens: Mapped[int | None] = mapped_column(sa.Integer, nullable=True)
    indexing_latency: Mapped[float | None] = mapped_column(sa.Float, nullable=True)
    completed_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)

    # pause
    is_paused: Mapped[bool | None] = mapped_column(sa.Boolean, nullable=True, server_default=sa.text("false"))
    paused_by = mapped_column(StringUUID, nullable=True)
    paused_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)

    # error
    error = mapped_column(LongText, nullable=True)
    stopped_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)

    # basic fields
    indexing_status = mapped_column(
        EnumText(IndexingStatus, length=255), nullable=False, server_default=sa.text("'waiting'")
    )
    enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("true"))
    disabled_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
    disabled_by = mapped_column(StringUUID, nullable=True)
    archived: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"))
    archived_reason = mapped_column(String(255), nullable=True)
    archived_by = mapped_column(StringUUID, nullable=True)
    archived_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp()
    )
    doc_type = mapped_column(EnumText(DocumentDocType, length=40), nullable=True)
    doc_metadata = mapped_column(AdjustedJSON, nullable=True)
    doc_form: Mapped[IndexStructureType] = mapped_column(
        EnumText(IndexStructureType, length=255), nullable=False, server_default=sa.text("'text_model'")
    )
    doc_language = mapped_column(String(255), nullable=True)
    need_summary: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"))

    DATA_SOURCES = ["upload_file", "notion_import", "website_crawl"]

    @property
    def display_status(self):
        status = None
        if self.indexing_status == "waiting":
            status = "queuing"
        elif self.indexing_status not in {"completed", "error", "waiting"} and self.is_paused:
            status = "paused"
        elif self.indexing_status in {"parsing", "cleaning", "splitting", "indexing"}:
            status = "indexing"
        elif self.indexing_status == "error":
            status = "error"
        elif self.indexing_status == "completed" and not self.archived and self.enabled:
            status = "available"
        elif self.indexing_status == "completed" and not self.archived and not self.enabled:
            status = "disabled"
        elif self.indexing_status == "completed" and self.archived:
            status = "archived"
        return status

    @property
    def data_source_info_dict(self) -> dict[str, Any]:
        if self.data_source_info:
            try:
                data_source_info_dict: dict[str, Any] = json.loads(self.data_source_info)
            except JSONDecodeError:
                data_source_info_dict = {}

            return data_source_info_dict
        return {}

    @property
    def data_source_detail_dict(self) -> dict[str, Any]:
        if self.data_source_info:
            if self.data_source_type == "upload_file":
                data_source_info_dict: dict[str, Any] = json.loads(self.data_source_info)
                file_detail = db.session.scalar(
                    select(UploadFile).where(UploadFile.id == data_source_info_dict["upload_file_id"])
                )
                if file_detail:
                    return {
                        "upload_file": {
                            "id": file_detail.id,
                            "name": file_detail.name,
                            "size": file_detail.size,
                            "extension": file_detail.extension,
                            "mime_type": file_detail.mime_type,
                            "created_by": file_detail.created_by,
                            "created_at": file_detail.created_at.timestamp(),
                        }
                    }
            elif self.data_source_type in {"notion_import", "website_crawl"}:
                result: dict[str, Any] = json.loads(self.data_source_info)
                return result
        return {}

    @property
    def average_segment_length(self):
        if self.word_count and self.word_count != 0 and self.segment_count and self.segment_count != 0:
            return self.word_count // self.segment_count
        return 0

    @property
    def dataset_process_rule(self):
        if self.dataset_process_rule_id:
            return db.session.get(DatasetProcessRule, self.dataset_process_rule_id)
        return None

    @property
    def dataset(self):
        return db.session.scalar(select(Dataset).where(Dataset.id == self.dataset_id))

    @property
    def segment_count(self):
        return (
            db.session.scalar(select(func.count(DocumentSegment.id)).where(DocumentSegment.document_id == self.id)) or 0
        )

    @property
    def hit_count(self):
        return db.session.scalar(
            select(func.coalesce(func.sum(DocumentSegment.hit_count), 0)).where(DocumentSegment.document_id == self.id)
        )

    @property
    def uploader(self):
        user = db.session.scalar(select(Account).where(Account.id == self.created_by))
        return user.name if user else None

    @property
    def upload_date(self):
        return self.created_at

    @property
    def last_update_date(self):
        return self.updated_at

    @property
    def doc_metadata_details(self) -> list[DocMetadataDetailItem] | None:
        if self.doc_metadata:
            document_metadatas = db.session.scalars(
                select(DatasetMetadata)
                .join(DatasetMetadataBinding, DatasetMetadataBinding.metadata_id == DatasetMetadata.id)
                .where(
                    DatasetMetadataBinding.dataset_id == self.dataset_id, DatasetMetadataBinding.document_id == self.id
                )
            ).all()
            metadata_list: list[DocMetadataDetailItem] = []
            for metadata in document_metadatas:
                metadata_dict: DocMetadataDetailItem = {
                    "id": metadata.id,
                    "name": metadata.name,
                    "type": metadata.type,
                    "value": self.doc_metadata.get(metadata.name),
                }
                metadata_list.append(metadata_dict)
            # deal built-in fields
            metadata_list.extend(self.get_built_in_fields())

            return metadata_list
        return None

    @property
    def process_rule_dict(self) -> ProcessRuleDict | None:
        if self.dataset_process_rule_id and self.dataset_process_rule:
            return self.dataset_process_rule.to_dict()
        return None

    def get_built_in_fields(self) -> list[DocMetadataDetailItem]:
        built_in_fields: list[DocMetadataDetailItem] = []
        built_in_fields.append(
            {
                "id": "built-in",
                "name": BuiltInField.document_name,
                "type": "string",
                "value": self.name,
            }
        )
        built_in_fields.append(
            {
                "id": "built-in",
                "name": BuiltInField.uploader,
                "type": "string",
                "value": self.uploader,
            }
        )
        built_in_fields.append(
            {
                "id": "built-in",
                "name": BuiltInField.upload_date,
                "type": "time",
                "value": str(self.created_at.timestamp()),
            }
        )
        built_in_fields.append(
            {
                "id": "built-in",
                "name": BuiltInField.last_update_date,
                "type": "time",
                "value": str(self.updated_at.timestamp()),
            }
        )
        built_in_fields.append(
            {
                "id": "built-in",
                "name": BuiltInField.source,
                "type": "string",
                "value": MetadataDataSource[self.data_source_type],
            }
        )
        return built_in_fields

    def to_dict(self) -> DocumentDict:
        result: DocumentDict = {
            "id": self.id,
            "tenant_id": self.tenant_id,
            "dataset_id": self.dataset_id,
            "position": self.position,
            "data_source_type": self.data_source_type,
            "data_source_info": self.data_source_info,
            "dataset_process_rule_id": self.dataset_process_rule_id,
            "batch": self.batch,
            "name": self.name,
            "created_from": self.created_from,
            "created_by": self.created_by,
            "created_api_request_id": self.created_api_request_id,
            "created_at": self.created_at,
            "processing_started_at": self.processing_started_at,
            "file_id": self.file_id,
            "word_count": self.word_count,
            "parsing_completed_at": self.parsing_completed_at,
            "cleaning_completed_at": self.cleaning_completed_at,
            "splitting_completed_at": self.splitting_completed_at,
            "tokens": self.tokens,
            "indexing_latency": self.indexing_latency,
            "completed_at": self.completed_at,
            "is_paused": self.is_paused,
            "paused_by": self.paused_by,
            "paused_at": self.paused_at,
            "error": self.error,
            "stopped_at": self.stopped_at,
            "indexing_status": self.indexing_status,
            "enabled": self.enabled,
            "disabled_at": self.disabled_at,
            "disabled_by": self.disabled_by,
            "archived": self.archived,
            "archived_reason": self.archived_reason,
            "archived_by": self.archived_by,
            "archived_at": self.archived_at,
            "updated_at": self.updated_at,
            "doc_type": self.doc_type,
            "doc_metadata": self.doc_metadata,
            "doc_form": self.doc_form,
            "doc_language": self.doc_language,
            "display_status": self.display_status,
            "data_source_info_dict": self.data_source_info_dict,
            "average_segment_length": self.average_segment_length,
            "dataset_process_rule": self.dataset_process_rule.to_dict() if self.dataset_process_rule else None,
            "dataset": None,
            "segment_count": self.segment_count,
            "hit_count": self.hit_count,
        }
        return result

    @classmethod
    def from_dict(cls, data: dict[str, Any]):
        return cls(
            id=data.get("id"),
            tenant_id=data.get("tenant_id"),
            dataset_id=data.get("dataset_id"),
            position=data.get("position"),
            data_source_type=data.get("data_source_type"),
            data_source_info=data.get("data_source_info"),
            dataset_process_rule_id=data.get("dataset_process_rule_id"),
            batch=data.get("batch"),
            name=data.get("name"),
            created_from=data.get("created_from"),
            created_by=data.get("created_by"),
            created_api_request_id=data.get("created_api_request_id"),
            created_at=data.get("created_at"),
            processing_started_at=data.get("processing_started_at"),
            file_id=data.get("file_id"),
            word_count=data.get("word_count"),
            parsing_completed_at=data.get("parsing_completed_at"),
            cleaning_completed_at=data.get("cleaning_completed_at"),
            splitting_completed_at=data.get("splitting_completed_at"),
            tokens=data.get("tokens"),
            indexing_latency=data.get("indexing_latency"),
            completed_at=data.get("completed_at"),
            is_paused=data.get("is_paused"),
            paused_by=data.get("paused_by"),
            paused_at=data.get("paused_at"),
            error=data.get("error"),
            stopped_at=data.get("stopped_at"),
            indexing_status=data.get("indexing_status"),
            enabled=data.get("enabled"),
            disabled_at=data.get("disabled_at"),
            disabled_by=data.get("disabled_by"),
            archived=data.get("archived"),
            archived_reason=data.get("archived_reason"),
            archived_by=data.get("archived_by"),
            archived_at=data.get("archived_at"),
            updated_at=data.get("updated_at"),
            doc_type=data.get("doc_type"),
            doc_metadata=data.get("doc_metadata"),
            doc_form=data.get("doc_form"),
            doc_language=data.get("doc_language"),
        )


class DocumentSegment(TypeBase):
    __tablename__ = "document_segments"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="document_segment_pkey"),
        sa.Index("document_segment_dataset_id_idx", "dataset_id"),
        sa.Index("document_segment_document_id_idx", "document_id"),
        sa.Index("document_segment_tenant_dataset_idx", "dataset_id", "tenant_id"),
        sa.Index("document_segment_tenant_document_idx", "document_id", "tenant_id"),
        sa.Index("document_segment_node_dataset_idx", "index_node_id", "dataset_id"),
        sa.Index("document_segment_tenant_idx", "tenant_id"),
    )

    # initial fields
    id: Mapped[str] = mapped_column(StringUUID, nullable=False, default_factory=lambda: str(uuid4()), init=False)
    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    document_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    position: Mapped[int]
    content: Mapped[str] = mapped_column(LongText, nullable=False)
    word_count: Mapped[int]
    tokens: Mapped[int]

    created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
    # basic fields
    # indexing fields
    index_node_id: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None)
    index_node_hash: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None)
    enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("true"), default=True)
    answer: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None)
    keywords: Mapped[Any] = mapped_column(sa.JSON, nullable=True, default=None)
    disabled_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None)
    disabled_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
    status: Mapped[SegmentStatus] = mapped_column(
        EnumText(SegmentStatus, length=255), server_default=sa.text("'waiting'"), default=SegmentStatus.WAITING
    )
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )
    updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )
    indexing_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None)
    completed_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None)
    error: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None)
    stopped_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None)
    hit_count: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0)

    @property
    def dataset(self):
        return db.session.scalar(select(Dataset).where(Dataset.id == self.dataset_id))

    @property
    def document(self):
        return db.session.scalar(select(Document).where(Document.id == self.document_id))

    @property
    def previous_segment(self):
        return db.session.scalar(
            select(DocumentSegment).where(
                DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position - 1
            )
        )

    @property
    def next_segment(self):
        return db.session.scalar(
            select(DocumentSegment).where(
                DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position + 1
            )
        )

    @property
    def child_chunks(self):
        if not self.document:
            return []
        process_rule = self.document.dataset_process_rule
        if process_rule and process_rule.mode == "hierarchical":
            rules_dict = process_rule.rules_dict
            if rules_dict:
                rules = Rule.model_validate(rules_dict)
                if rules.parent_mode and rules.parent_mode != ParentMode.FULL_DOC:
                    child_chunks = db.session.scalars(
                        select(ChildChunk).where(ChildChunk.segment_id == self.id).order_by(ChildChunk.position.asc())
                    ).all()
                    return child_chunks or []
        return []

    def get_child_chunks(self):
        if not self.document:
            return []
        process_rule = self.document.dataset_process_rule
        if process_rule and process_rule.mode == "hierarchical":
            rules_dict = process_rule.rules_dict
            if rules_dict:
                rules = Rule.model_validate(rules_dict)
                if rules.parent_mode:
                    child_chunks = db.session.scalars(
                        select(ChildChunk).where(ChildChunk.segment_id == self.id).order_by(ChildChunk.position.asc())
                    ).all()
                    return child_chunks or []
        return []

    @property
    def sign_content(self) -> str:
        return self.get_sign_content()

    def get_sign_content(self) -> str:
        signed_urls: list[tuple[int, int, str]] = []
        text = self.content

        # For data before v0.10.0
        pattern = r"/files/([a-f0-9\-]+)/image-preview(?:\?.*?)?"
        matches = re.finditer(pattern, text)
        for match in matches:
            upload_file_id = match.group(1)
            nonce = os.urandom(16).hex()
            timestamp = str(int(time.time()))
            data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}"
            secret_key = dify_config.SECRET_KEY.encode()
            sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
            encoded_sign = base64.urlsafe_b64encode(sign).decode()

            params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
            base_url = f"/files/{upload_file_id}/image-preview"
            signed_url = f"{base_url}?{params}"
            signed_urls.append((match.start(), match.end(), signed_url))

        # For data after v0.10.0
        pattern = r"/files/([a-f0-9\-]+)/file-preview(?:\?.*?)?"
        matches = re.finditer(pattern, text)
        for match in matches:
            upload_file_id = match.group(1)
            nonce = os.urandom(16).hex()
            timestamp = str(int(time.time()))
            data_to_sign = f"file-preview|{upload_file_id}|{timestamp}|{nonce}"
            secret_key = dify_config.SECRET_KEY.encode()
            sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
            encoded_sign = base64.urlsafe_b64encode(sign).decode()

            params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
            base_url = f"/files/{upload_file_id}/file-preview"
            signed_url = f"{base_url}?{params}"
            signed_urls.append((match.start(), match.end(), signed_url))

        # For tools directory - direct file formats (e.g., .png, .jpg, etc.)
        # Match URL including any query parameters up to common URL boundaries (space, parenthesis, quotes)
        pattern = r"/files/tools/([a-f0-9\-]+)\.([a-zA-Z0-9]+)(?:\?[^\s\)\"\']*)?"
        matches = re.finditer(pattern, text)
        for match in matches:
            upload_file_id = match.group(1)
            file_extension = match.group(2)
            nonce = os.urandom(16).hex()
            timestamp = str(int(time.time()))
            data_to_sign = f"file-preview|{upload_file_id}|{timestamp}|{nonce}"
            secret_key = dify_config.SECRET_KEY.encode()
            sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
            encoded_sign = base64.urlsafe_b64encode(sign).decode()

            params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
            base_url = f"/files/tools/{upload_file_id}.{file_extension}"
            signed_url = f"{base_url}?{params}"
            signed_urls.append((match.start(), match.end(), signed_url))

        # Reconstruct the text with signed URLs
        offset = 0
        for start, end, signed_url in signed_urls:
            text = text[: start + offset] + signed_url + text[end + offset :]
            offset += len(signed_url) - (end - start)

        return text

    @property
    def attachments(self) -> list[AttachmentItem]:
        # Use JOIN to fetch attachments in a single query instead of two separate queries
        attachments_with_bindings = db.session.execute(
            select(SegmentAttachmentBinding, UploadFile)
            .join(UploadFile, UploadFile.id == SegmentAttachmentBinding.attachment_id)
            .where(
                SegmentAttachmentBinding.tenant_id == self.tenant_id,
                SegmentAttachmentBinding.dataset_id == self.dataset_id,
                SegmentAttachmentBinding.document_id == self.document_id,
                SegmentAttachmentBinding.segment_id == self.id,
            )
        ).all()
        if not attachments_with_bindings:
            return []
        attachment_list: list[AttachmentItem] = []
        for _, attachment in attachments_with_bindings:
            upload_file_id = attachment.id
            nonce = os.urandom(16).hex()
            timestamp = str(int(time.time()))
            data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}"
            secret_key = dify_config.SECRET_KEY.encode()
            sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
            encoded_sign = base64.urlsafe_b64encode(sign).decode()

            params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
            reference_url = dify_config.FILES_URL or dify_config.CONSOLE_API_URL or ""
            base_url = f"{reference_url}/files/{upload_file_id}/image-preview"
            source_url = f"{base_url}?{params}"
            attachment_list.append(
                {
                    "id": attachment.id,
                    "name": attachment.name,
                    "size": attachment.size,
                    "extension": attachment.extension,
                    "mime_type": attachment.mime_type,
                    "source_url": source_url,
                }
            )
        return attachment_list


class ChildChunk(TypeBase):
    __tablename__ = "child_chunks"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="child_chunk_pkey"),
        sa.Index("child_chunk_dataset_id_idx", "tenant_id", "dataset_id", "document_id", "segment_id", "index_node_id"),
        sa.Index("child_chunks_node_idx", "index_node_id", "dataset_id"),
        sa.Index("child_chunks_segment_idx", "segment_id"),
    )

    # initial fields
    id: Mapped[str] = mapped_column(StringUUID, nullable=False, default_factory=lambda: str(uuid4()), init=False)
    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    document_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    segment_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    position: Mapped[int] = mapped_column(sa.Integer, nullable=False)
    content: Mapped[str] = mapped_column(LongText, nullable=False)
    word_count: Mapped[int] = mapped_column(sa.Integer, nullable=False)
    # indexing fields
    created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False
    )
    updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, init=False)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime,
        nullable=False,
        server_default=sa.func.current_timestamp(),
        onupdate=func.current_timestamp(),
        init=False,
    )
    indexing_at: Mapped[datetime | None] = mapped_column(
        DateTime, nullable=True, insert_default=None, server_default=None, init=False
    )
    completed_at: Mapped[datetime | None] = mapped_column(
        DateTime, nullable=True, insert_default=None, server_default=None, init=False
    )
    index_node_id: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None)
    index_node_hash: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None)
    type: Mapped[SegmentType] = mapped_column(
        EnumText(SegmentType, length=255),
        nullable=False,
        server_default=sa.text("'automatic'"),
        default=SegmentType.AUTOMATIC,
    )
    error: Mapped[str | None] = mapped_column(LongText, nullable=True, init=False)

    @property
    def dataset(self):
        return db.session.scalar(select(Dataset).where(Dataset.id == self.dataset_id))

    @property
    def document(self):
        return db.session.scalar(select(Document).where(Document.id == self.document_id))

    @property
    def segment(self):
        return db.session.scalar(select(DocumentSegment).where(DocumentSegment.id == self.segment_id))


class AppDatasetJoin(TypeBase):
    __tablename__ = "app_dataset_joins"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="app_dataset_join_pkey"),
        sa.Index("app_dataset_join_app_dataset_idx", "dataset_id", "app_id"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID,
        primary_key=True,
        nullable=False,
        insert_default=lambda: str(uuid4()),
        default_factory=lambda: str(uuid4()),
        init=False,
    )
    app_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False
    )

    @property
    def app(self):
        return db.session.get(App, self.app_id)


class DatasetQuery(TypeBase):
    __tablename__ = "dataset_queries"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="dataset_query_pkey"),
        sa.Index("dataset_query_dataset_id_idx", "dataset_id"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID,
        primary_key=True,
        nullable=False,
        insert_default=lambda: str(uuid4()),
        default_factory=lambda: str(uuid4()),
        init=False,
    )
    dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    content: Mapped[str] = mapped_column(LongText, nullable=False)
    source: Mapped[str] = mapped_column(EnumText(DatasetQuerySource, length=255), nullable=False)
    source_app_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
    created_by_role: Mapped[CreatorUserRole] = mapped_column(EnumText(CreatorUserRole, length=255), nullable=False)
    created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False
    )

    @property
    def queries(self) -> list[dict[str, Any]]:
        try:
            queries = json.loads(self.content)
            if isinstance(queries, list):
                for query in queries:
                    if query["content_type"] == QueryType.IMAGE_QUERY:
                        file_info = db.session.scalar(select(UploadFile).where(UploadFile.id == query["content"]))
                        if file_info:
                            query["file_info"] = {
                                "id": file_info.id,
                                "name": file_info.name,
                                "size": file_info.size,
                                "extension": file_info.extension,
                                "mime_type": file_info.mime_type,
                                "source_url": sign_upload_file_preview_url(file_info.id, file_info.extension),
                            }
                    else:
                        query["file_info"] = None

                return queries
            else:
                return [queries]
        except JSONDecodeError:
            return [
                {
                    "content_type": QueryType.TEXT_QUERY,
                    "content": self.content,
                    "file_info": None,
                }
            ]


class DatasetKeywordTable(TypeBase):
    __tablename__ = "dataset_keyword_tables"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="dataset_keyword_table_pkey"),
        sa.Index("dataset_keyword_table_dataset_id_idx", "dataset_id"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID,
        primary_key=True,
        insert_default=lambda: str(uuid4()),
        default_factory=lambda: str(uuid4()),
        init=False,
    )
    dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False, unique=True)
    keyword_table: Mapped[str] = mapped_column(LongText, nullable=False)
    data_source_type: Mapped[str] = mapped_column(
        String(255), nullable=False, server_default=sa.text("'database'"), default="database"
    )

    @property
    def keyword_table_dict(self) -> dict[str, set[Any]] | None:
        class SetDecoder(json.JSONDecoder):
            def __init__(self, *args: Any, **kwargs: Any) -> None:
                def object_hook(dct: Any) -> Any:
                    if isinstance(dct, dict):
                        result: dict[str, Any] = {}
                        items = cast(dict[str, Any], dct).items()
                        for keyword, node_idxs in items:
                            if isinstance(node_idxs, list):
                                result[keyword] = set(cast(list[Any], node_idxs))
                            else:
                                result[keyword] = node_idxs
                        return result
                    return dct

                super().__init__(object_hook=object_hook, *args, **kwargs)

        # get dataset
        dataset = db.session.scalar(select(Dataset).where(Dataset.id == self.dataset_id))
        if not dataset:
            return None
        if self.data_source_type == "database":
            return json.loads(self.keyword_table, cls=SetDecoder) if self.keyword_table else None
        else:
            file_key = "keyword_files/" + dataset.tenant_id + "/" + self.dataset_id + ".txt"
            try:
                keyword_table_text = storage.load_once(file_key)
                if keyword_table_text:
                    return json.loads(keyword_table_text.decode("utf-8"), cls=SetDecoder)
                return None
            except Exception:
                logger.exception("Failed to load keyword table from file: %s", file_key)
                return None


class Embedding(TypeBase):
    __tablename__ = "embeddings"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="embedding_pkey"),
        sa.UniqueConstraint("model_name", "hash", "provider_name", name="embedding_hash_idx"),
        sa.Index("created_at_idx", "created_at"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID,
        primary_key=True,
        insert_default=lambda: str(uuid4()),
        default_factory=lambda: str(uuid4()),
        init=False,
    )
    model_name: Mapped[str] = mapped_column(
        String(255), nullable=False, server_default=sa.text("'text-embedding-ada-002'")
    )
    hash: Mapped[str] = mapped_column(String(64), nullable=False)
    embedding: Mapped[bytes] = mapped_column(BinaryData, nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )
    provider_name: Mapped[str] = mapped_column(String(255), nullable=False, server_default=sa.text("''"))

    def set_embedding(self, embedding_data: list[float]):
        self.embedding = pickle.dumps(embedding_data, protocol=pickle.HIGHEST_PROTOCOL)

    def get_embedding(self) -> list[float]:
        return cast(list[float], pickle.loads(self.embedding))  # noqa: S301


class DatasetCollectionBinding(TypeBase):
    __tablename__ = "dataset_collection_bindings"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="dataset_collection_bindings_pkey"),
        sa.Index("provider_model_name_idx", "provider_name", "model_name"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID,
        primary_key=True,
        insert_default=lambda: str(uuid4()),
        default_factory=lambda: str(uuid4()),
        init=False,
    )
    provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
    model_name: Mapped[str] = mapped_column(String(255), nullable=False)
    type: Mapped[str] = mapped_column(
        EnumText(CollectionBindingType, length=40), server_default=sa.text("'dataset'"), nullable=False
    )
    collection_name: Mapped[str] = mapped_column(String(64), nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )


class TidbAuthBinding(TypeBase):
    __tablename__ = "tidb_auth_bindings"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="tidb_auth_bindings_pkey"),
        sa.Index("tidb_auth_bindings_tenant_idx", "tenant_id"),
        sa.Index("tidb_auth_bindings_active_idx", "active"),
        sa.Index("tidb_auth_bindings_created_at_idx", "created_at"),
        sa.Index("tidb_auth_bindings_status_idx", "status"),
    )
    id: Mapped[str] = mapped_column(
        StringUUID,
        primary_key=True,
        insert_default=lambda: str(uuid4()),
        default_factory=lambda: str(uuid4()),
        init=False,
    )
    tenant_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
    cluster_id: Mapped[str] = mapped_column(String(255), nullable=False)
    cluster_name: Mapped[str] = mapped_column(String(255), nullable=False)
    active: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"))
    status: Mapped[TidbAuthBindingStatus] = mapped_column(
        EnumText(TidbAuthBindingStatus, length=255), nullable=False, server_default=sa.text("'CREATING'")
    )
    account: Mapped[str] = mapped_column(String(255), nullable=False)
    password: Mapped[str] = mapped_column(String(255), nullable=False)
    qdrant_endpoint: Mapped[str | None] = mapped_column(String(512), nullable=True, default=None)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )


class Whitelist(TypeBase):
    __tablename__ = "whitelists"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="whitelists_pkey"),
        sa.Index("whitelists_tenant_idx", "tenant_id"),
    )
    id: Mapped[str] = mapped_column(
        StringUUID,
        primary_key=True,
        insert_default=lambda: str(uuid4()),
        default_factory=lambda: str(uuid4()),
        init=False,
    )
    tenant_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
    category: Mapped[str] = mapped_column(String(255), nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )


class DatasetPermission(TypeBase):
    __tablename__ = "dataset_permissions"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="dataset_permission_pkey"),
        sa.Index("idx_dataset_permissions_dataset_id", "dataset_id"),
        sa.Index("idx_dataset_permissions_account_id", "account_id"),
        sa.Index("idx_dataset_permissions_tenant_id", "tenant_id"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID,
        insert_default=lambda: str(uuid4()),
        default_factory=lambda: str(uuid4()),
        primary_key=True,
        init=False,
    )
    dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    account_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    has_permission: Mapped[bool] = mapped_column(
        sa.Boolean, nullable=False, server_default=sa.text("true"), default=True
    )
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )


class ExternalKnowledgeApis(TypeBase):
    __tablename__ = "external_knowledge_apis"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="external_knowledge_apis_pkey"),
        sa.Index("external_knowledge_apis_tenant_idx", "tenant_id"),
        sa.Index("external_knowledge_apis_name_idx", "name"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID,
        nullable=False,
        insert_default=lambda: str(uuid4()),
        default_factory=lambda: str(uuid4()),
        init=False,
    )
    name: Mapped[str] = mapped_column(String(255), nullable=False)
    description: Mapped[str] = mapped_column(String(255), nullable=False)
    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    settings: Mapped[str | None] = mapped_column(LongText, nullable=True)
    created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )
    updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
    )

    def to_dict(self) -> ExternalKnowledgeApiDict:
        return {
            "id": self.id,
            "tenant_id": self.tenant_id,
            "name": self.name,
            "description": self.description,
            "settings": self.settings_dict,
            "dataset_bindings": self.dataset_bindings,
            "created_by": self.created_by,
            "created_at": self.created_at.isoformat(),
        }

    @property
    def settings_dict(self) -> dict[str, Any] | None:
        try:
            return json.loads(self.settings) if self.settings else None
        except JSONDecodeError:
            return None

    @property
    def dataset_bindings(self) -> list[DatasetBindingItem]:
        external_knowledge_bindings = db.session.scalars(
            select(ExternalKnowledgeBindings).where(ExternalKnowledgeBindings.external_knowledge_api_id == self.id)
        ).all()
        dataset_ids = [binding.dataset_id for binding in external_knowledge_bindings]
        datasets = db.session.scalars(select(Dataset).where(Dataset.id.in_(dataset_ids))).all()
        dataset_bindings: list[DatasetBindingItem] = []
        for dataset in datasets:
            dataset_bindings.append({"id": dataset.id, "name": dataset.name})

        return dataset_bindings


class ExternalKnowledgeBindings(TypeBase):
    __tablename__ = "external_knowledge_bindings"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="external_knowledge_bindings_pkey"),
        sa.Index("external_knowledge_bindings_tenant_idx", "tenant_id"),
        sa.Index("external_knowledge_bindings_dataset_idx", "dataset_id"),
        sa.Index("external_knowledge_bindings_external_knowledge_idx", "external_knowledge_id"),
        sa.Index("external_knowledge_bindings_external_knowledge_api_idx", "external_knowledge_api_id"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID,
        nullable=False,
        insert_default=lambda: str(uuid4()),
        default_factory=lambda: str(uuid4()),
        init=False,
    )
    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    external_knowledge_api_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    external_knowledge_id: Mapped[str] = mapped_column(String(512), nullable=False)
    created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )
    updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None, init=False)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False
    )


class DatasetAutoDisableLog(TypeBase):
    __tablename__ = "dataset_auto_disable_logs"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="dataset_auto_disable_log_pkey"),
        sa.Index("dataset_auto_disable_log_tenant_idx", "tenant_id"),
        sa.Index("dataset_auto_disable_log_dataset_idx", "dataset_id"),
        sa.Index("dataset_auto_disable_log_created_atx", "created_at"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
    )
    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    document_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    notified: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"), default=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False
    )


class RateLimitLog(TypeBase):
    __tablename__ = "rate_limit_logs"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="rate_limit_log_pkey"),
        sa.Index("rate_limit_log_tenant_idx", "tenant_id"),
        sa.Index("rate_limit_log_operation_idx", "operation"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
    )
    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    subscription_plan: Mapped[str] = mapped_column(String(255), nullable=False)
    operation: Mapped[str] = mapped_column(String(255), nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )


class DatasetMetadata(TypeBase):
    __tablename__ = "dataset_metadatas"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="dataset_metadata_pkey"),
        sa.Index("dataset_metadata_tenant_idx", "tenant_id"),
        sa.Index("dataset_metadata_dataset_idx", "dataset_id"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
    )
    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    type: Mapped[str] = mapped_column(EnumText(DatasetMetadataType, length=255), nullable=False)
    name: Mapped[str] = mapped_column(String(255), nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime,
        nullable=False,
        server_default=sa.func.current_timestamp(),
        onupdate=func.current_timestamp(),
        init=False,
    )
    created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
    updated_by: Mapped[str] = mapped_column(StringUUID, nullable=True, default=None)


class DatasetMetadataBinding(TypeBase):
    __tablename__ = "dataset_metadata_bindings"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="dataset_metadata_binding_pkey"),
        sa.Index("dataset_metadata_binding_tenant_idx", "tenant_id"),
        sa.Index("dataset_metadata_binding_dataset_idx", "dataset_id"),
        sa.Index("dataset_metadata_binding_metadata_idx", "metadata_id"),
        sa.Index("dataset_metadata_binding_document_idx", "document_id"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
    )
    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    metadata_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    document_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )
    created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)


class PipelineBuiltInTemplate(TypeBase):
    __tablename__ = "pipeline_built_in_templates"
    __table_args__ = (sa.PrimaryKeyConstraint("id", name="pipeline_built_in_template_pkey"),)

    id: Mapped[str] = mapped_column(
        StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False
    )
    name: Mapped[str] = mapped_column(sa.String(255), nullable=False)
    description: Mapped[str] = mapped_column(LongText, nullable=False)
    chunk_structure: Mapped[str] = mapped_column(sa.String(255), nullable=False)
    icon: Mapped[dict[str, Any]] = mapped_column(sa.JSON, nullable=False)
    yaml_content: Mapped[str] = mapped_column(LongText, nullable=False)
    copyright: Mapped[str] = mapped_column(sa.String(255), nullable=False)
    privacy_policy: Mapped[str] = mapped_column(sa.String(255), nullable=False)
    position: Mapped[int] = mapped_column(sa.Integer, nullable=False)
    install_count: Mapped[int] = mapped_column(sa.Integer, nullable=False)
    language: Mapped[str] = mapped_column(sa.String(255), nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )
    updated_at: Mapped[datetime] = mapped_column(
        sa.DateTime,
        nullable=False,
        server_default=func.current_timestamp(),
        onupdate=func.current_timestamp(),
        init=False,
    )


class PipelineCustomizedTemplate(TypeBase):
    __tablename__ = "pipeline_customized_templates"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="pipeline_customized_template_pkey"),
        sa.Index("pipeline_customized_template_tenant_idx", "tenant_id"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False
    )
    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    name: Mapped[str] = mapped_column(sa.String(255), nullable=False)
    description: Mapped[str] = mapped_column(LongText, nullable=False)
    chunk_structure: Mapped[str] = mapped_column(sa.String(255), nullable=False)
    icon: Mapped[dict[str, Any]] = mapped_column(sa.JSON, nullable=False)
    position: Mapped[int] = mapped_column(sa.Integer, nullable=False)
    yaml_content: Mapped[str] = mapped_column(LongText, nullable=False)
    install_count: Mapped[int] = mapped_column(sa.Integer, nullable=False)
    language: Mapped[str] = mapped_column(sa.String(255), nullable=False)
    created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
    updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None, init=False)
    created_at: Mapped[datetime] = mapped_column(
        sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )
    updated_at: Mapped[datetime] = mapped_column(
        sa.DateTime,
        nullable=False,
        server_default=func.current_timestamp(),
        onupdate=func.current_timestamp(),
        init=False,
    )

    @property
    def created_user_name(self):
        account = db.session.scalar(select(Account).where(Account.id == self.created_by))
        if account:
            return account.name
        return ""


class Pipeline(TypeBase):
    __tablename__ = "pipelines"
    __table_args__ = (sa.PrimaryKeyConstraint("id", name="pipeline_pkey"),)

    id: Mapped[str] = mapped_column(
        StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False
    )
    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    name: Mapped[str] = mapped_column(sa.String(255), nullable=False)
    description: Mapped[str] = mapped_column(LongText, nullable=False, default=sa.text("''"))
    workflow_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
    is_public: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"), default=False)
    is_published: Mapped[bool] = mapped_column(
        sa.Boolean, nullable=False, server_default=sa.text("false"), default=False
    )
    created_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
    created_at: Mapped[datetime] = mapped_column(
        sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )
    updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
    updated_at: Mapped[datetime] = mapped_column(
        sa.DateTime,
        nullable=False,
        server_default=func.current_timestamp(),
        onupdate=func.current_timestamp(),
        init=False,
    )

    def retrieve_dataset(self, session: Session):
        return session.scalar(select(Dataset).where(Dataset.pipeline_id == self.id))


class DocumentPipelineExecutionLog(TypeBase):
    __tablename__ = "document_pipeline_execution_logs"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="document_pipeline_execution_log_pkey"),
        sa.Index("document_pipeline_execution_logs_document_id_idx", "document_id"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False
    )
    pipeline_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    document_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    datasource_type: Mapped[str] = mapped_column(sa.String(255), nullable=False)
    datasource_info: Mapped[str] = mapped_column(LongText, nullable=False)
    datasource_node_id: Mapped[str] = mapped_column(sa.String(255), nullable=False)
    input_data: Mapped[dict[str, Any]] = mapped_column(sa.JSON, nullable=False)
    created_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
    created_at: Mapped[datetime] = mapped_column(
        sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )


class PipelineRecommendedPlugin(TypeBase):
    __tablename__ = "pipeline_recommended_plugins"
    __table_args__ = (sa.PrimaryKeyConstraint("id", name="pipeline_recommended_plugin_pkey"),)

    id: Mapped[str] = mapped_column(
        StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False
    )
    plugin_id: Mapped[str] = mapped_column(LongText, nullable=False)
    provider_name: Mapped[str] = mapped_column(LongText, nullable=False)
    type: Mapped[str] = mapped_column(sa.String(50), nullable=False, server_default=sa.text("'tool'"))
    position: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0)
    active: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=True)
    created_at: Mapped[datetime] = mapped_column(
        sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )
    updated_at: Mapped[datetime] = mapped_column(
        sa.DateTime,
        nullable=False,
        server_default=func.current_timestamp(),
        onupdate=func.current_timestamp(),
        init=False,
    )


class SegmentAttachmentBinding(TypeBase):
    __tablename__ = "segment_attachment_bindings"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="segment_attachment_binding_pkey"),
        sa.Index(
            "segment_attachment_binding_tenant_dataset_document_segment_idx",
            "tenant_id",
            "dataset_id",
            "document_id",
            "segment_id",
        ),
        sa.Index("segment_attachment_binding_attachment_idx", "attachment_id"),
    )
    id: Mapped[str] = mapped_column(
        StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False
    )
    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    document_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    segment_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    attachment_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )


class DocumentSegmentSummary(TypeBase):
    __tablename__ = "document_segment_summaries"
    __table_args__ = (
        sa.PrimaryKeyConstraint("id", name="document_segment_summaries_pkey"),
        sa.Index("document_segment_summaries_dataset_id_idx", "dataset_id"),
        sa.Index("document_segment_summaries_document_id_idx", "document_id"),
        sa.Index("document_segment_summaries_chunk_id_idx", "chunk_id"),
        sa.Index("document_segment_summaries_status_idx", "status"),
    )

    id: Mapped[str] = mapped_column(
        StringUUID,
        nullable=False,
        insert_default=lambda: str(uuid4()),
        default_factory=lambda: str(uuid4()),
        init=False,
    )
    dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    document_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    # corresponds to DocumentSegment.id or parent chunk id
    chunk_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
    summary_content: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None)
    summary_index_node_id: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None)
    summary_index_node_hash: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None)
    tokens: Mapped[int | None] = mapped_column(sa.Integer, nullable=True, default=None)
    status: Mapped[SummaryStatus] = mapped_column(
        EnumText(SummaryStatus, length=32),
        nullable=False,
        server_default=sa.text("'generating'"),
        default=SummaryStatus.GENERATING,
    )
    error: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None)
    enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("true"), default=True)
    disabled_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None)
    disabled_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, server_default=func.current_timestamp(), init=False
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime,
        nullable=False,
        server_default=func.current_timestamp(),
        onupdate=func.current_timestamp(),
        init=False,
    )

    def __repr__(self):
        return f"<DocumentSegmentSummary id={self.id} chunk_id={self.chunk_id} status={self.status}>"
