from enum import StrEnum

from litestar.plugins.sqlalchemy import base
from sqlalchemy import ForeignKey, Index, text
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column


class UserRole(StrEnum):
    ADMIN = "admin"
    USER = "user"


class DataSourceType(StrEnum):
    UPLOAD = "upload"
    API = "api"


class DataSourceStatus(StrEnum):
    PROCESSING = "processing"
    READY = "ready"
    ERROR = "error"


class VectorDistance(StrEnum):
    COSINE = "Cosine"
    EUCLID = "Euclid"
    DOT = "Dot"
    MANHATTAN = "Manhattan"


class RetrievalMode(StrEnum):
    DENSE = "dense"
    SPARSE = "sparse"
    HYBRID = "hybrid"


class User(base.UUIDAuditBase):
    __tablename__ = "users"
    name: Mapped[str]
    username: Mapped[str | None]
    password: Mapped[bytes | None]
    avatar: Mapped[str | None]
    role: Mapped[UserRole] = mapped_column(default=UserRole.USER)
    extra_info: Mapped[dict | None] = mapped_column(JSONB)


class KnowledgeCollection(base.UUIDAuditBase):
    __tablename__ = "knowledge_collections"

    id: Mapped[str] = mapped_column(primary_key=True)
    name: Mapped[str]
    vector_dimensions: Mapped[int] = mapped_column(default=1024)
    vector_distance: Mapped[VectorDistance] = mapped_column(
        default=VectorDistance.COSINE
    )
    retrieval_mode: Mapped[RetrievalMode] = mapped_column(default=RetrievalMode.DENSE)


class KnowledgeDataSource(base.UUIDAuditBase):
    __tablename__ = "knowledge_data_sources"
    __table_args__ = (
        Index("idx_knowledge_data_sources_collection_id", "collection_id"),
        Index(
            "idx_knowledge_data_sources_file_hash", text("(extra_info ->> 'file_hash')")
        ),
    )

    name: Mapped[str]
    definition_id: Mapped[str]
    type: Mapped[DataSourceType]
    status: Mapped[DataSourceStatus]
    extra_info: Mapped[dict | None] = mapped_column(JSONB)
    collection_id: Mapped[str] = mapped_column(ForeignKey("knowledge_collections.id"))
