from sqlalchemy import Column, String, Text, Integer, Float, ForeignKey, JSON
from sqlalchemy.orm import relationship
from app.models.base import BaseModel


class KGNode(BaseModel):
    __tablename__ = "kg_nodes"

    user_id = Column(String(36), ForeignKey("users.id"), nullable=False, index=True)
    name = Column(String(500), nullable=False)
    entity_type = Column(String(100), nullable=False)  # person, concept, task, document, chunk, tag
    description = Column(Text, nullable=True)
    properties_ = Column(JSON, nullable=True)  # 额外属性
    source_document_id = Column(String(36), ForeignKey("documents.id"), nullable=True)
    importance = Column(Float, default=0.5)  # 重要性 0-1
    last_updated_by = Column(String(36), nullable=True)  # 哪个 agent 更新过

    outgoing_edges = relationship("KGEdge", foreign_keys="KGEdge.source_id", back_populates="source_node", cascade="all, delete-orphan")
    incoming_edges = relationship("KGEdge", foreign_keys="KGEdge.target_id", back_populates="target_node", cascade="all, delete-orphan")


class KGEdge(BaseModel):
    __tablename__ = "kg_edges"

    source_id = Column(String(36), ForeignKey("kg_nodes.id"), nullable=False, index=True)
    target_id = Column(String(36), ForeignKey("kg_nodes.id"), nullable=False, index=True)
    relation_type = Column(String(100), nullable=False)  # related_to, part_of, caused_by, depends_on, etc.
    weight = Column(Float, default=0.5)  # 关系强度 0-1
    properties_ = Column(JSON, nullable=True)

    source_node = relationship("KGNode", foreign_keys=[source_id], back_populates="outgoing_edges")
    target_node = relationship("KGNode", foreign_keys=[target_id], back_populates="incoming_edges")