import enum
import numpy as np
from typing import Optional, Any, List, Dict
from pydantic import Field as PydanticField
from datetime import datetime

from sqlmodel import (
    SQLModel,
    Field,
    Column,
    JSON,
    Text,
    BLOB,
    Relationship as SQLModelRelationship,
    DateTime,
)
from sqlalchemy import func
from tidb_vector.sqlalchemy import VectorType


def cosine_distance(v1, v2):
    return 1 - np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))


class ImportStatus(str, enum.Enum):
    PENDING = "pending"
    EXTRACTING_GRAPH = "extracting_graph"
    EXTRACTED_GRAPH = "extracted_graph"
    SAVING_TO_DB = "saving_to_db"
    DONE = "done"
    FAILED = "failed"


class EntityType(str, enum.Enum):
    original = "original"
    synopsis = "synopsis"

    def __str__(self):
        return self.value


class Chunk(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    doc_id: str = Field(max_length=256)
    text: str = Field(sa_column=Column(Text))
    embedding: Any = Field(sa_column=Column(VectorType(1536)))
    meta: List | Dict = Field(default={}, sa_column=Column(JSON))
    import_status: ImportStatus = ImportStatus.PENDING
    import_result: Optional[str] = Field(sa_column=Column(Text))
    entities_df: bytes | None = Field(sa_column=Column(BLOB))
    relationships_df: bytes | None = Field(sa_column=Column(BLOB))
    llm_usage: List | Dict | None = Field(sa_column=Column(JSON))
    last_modified_at: Optional[datetime] = Field(sa_column=Column(DateTime))
    created_at: datetime = Field(
        sa_column=Column(DateTime, server_default=func.now(), nullable=True)
    )
    updated_at: datetime = Field(
        sa_column=Column(
            DateTime, server_default=func.now(), onupdate=func.now(), nullable=True
        )
    )

    __tablename__ = "chunks"

    def __hash__(self):
        return hash(self.id)


class EntityBase(SQLModel):
    name: str = Field(max_length=512)
    description: str = Field(sa_column=Column(Text))
    meta: List | Dict = Field(default={}, sa_column=Column(JSON))
    entity_type: EntityType = EntityType.original
    synopsis_info: List | Dict | None = Field(default=None, sa_column=Column(JSON))


class Entity(EntityBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    description_vec: Any = Field(sa_column=Column(VectorType(1536)))
    meta_vec: Any = Field(sa_column=Column(VectorType(1536)))

    __tablename__ = "entities"

    def __hash__(self):
        return hash(self.id)

    def screenshot(self):
        return self.model_dump(exclude={"description_vec", "meta_vec"})


# Public Entity model will be used in API response
class EntityPublic(EntityBase):
    id: int


class RelationshipBase(SQLModel):
    relationship_desc: str = Field(sa_column=Column(Text))
    meta: List | Dict = Field(default={}, sa_column=Column(JSON))
    weight: int = 0
    source_entity_id: int = Field(foreign_key="entities.id")
    target_entity_id: int = Field(foreign_key="entities.id")
    last_modified_at: Optional[datetime] = Field(sa_column=Column(DateTime))


class Relationship(RelationshipBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    relationship_desc_vec: Any = Field(sa_column=Column(VectorType(1536)))
    source_entity: Entity = SQLModelRelationship(
        sa_relationship_kwargs={
            "primaryjoin": "Relationship.source_entity_id == Entity.id",
            "lazy": "joined",
        },
    )
    target_entity: Entity = SQLModelRelationship(
        sa_relationship_kwargs={
            "primaryjoin": "Relationship.target_entity_id == Entity.id",
            "lazy": "joined",
        },
    )

    __tablename__ = "relationships"

    def __hash__(self):
        return hash(self.id)

    def screenshot(self):
        obj_dict = self.model_dump(
            exclude={
                "relationship_desc_vec",
                "source_entity",
                "target_entity",
                "last_modified_at",
            }
        )
        obj_dict["description"] = obj_dict.pop("relationship_desc")
        return obj_dict


class RelationshipPublic(RelationshipBase):
    id: int
    relationship_desc: str = PydanticField(serialization_alias="description")


class FeedbackType(str, enum.Enum):
    like = "like"
    dislike = "dislike"

    @classmethod
    def adjust_relationship_weight(cls, feedback_type):
        weights = {cls.like: 10, cls.dislike: -10}
        return weights.get(feedback_type, 0)


class Feedback(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    feedback_type: FeedbackType = FeedbackType.like
    query: str
    langfuse_link: str
    relationships: List[int] = Field(default=[], sa_column=Column(JSON))
    created_at: datetime = Field(
        default=None,
        sa_column=Column(DateTime, server_default=func.now(), nullable=True),
    )

    __tablename__ = "feedback"


class FeedbackCreate(SQLModel):
    feedback_type: FeedbackType
    query: str
    langfuse_link: str
    relationships: List[int]