# -*- coding: utf-8 -*-
# 
# @author: lyg
# @date: 2025-5-12
# @version: 1
# @description: 文档数据库模型

from sqlalchemy import create_engine, Column, DateTime, Integer, Text, ForeignKey
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship
from sqlalchemy.orm import sessionmaker, scoped_session

from knowledgebase.log import Log

Base = declarative_base()
metadata = Base.metadata


class TParagraph(Base):
    """
    段落表
    """
    __tablename__ = 't_paragraphs'
    id = Column(Integer, primary_key=True)
    text = Column(Text)
    title_level = Column(Integer)
    title_num = Column(Text)
    num_level = Column(Integer)
    num = Column(Integer)
    doc_id = Column(Integer, ForeignKey('t_docs.id'))
    parent_id = Column(Integer, ForeignKey('t_paragraphs.id'))
    parent_link = relationship("TParagraphLink", foreign_keys='TParagraphLink.child_id', back_populates='parent',
                               uselist=False)
    children_links = relationship("TParagraphLink", foreign_keys='TParagraphLink.parent_id', back_populates='child')
    ref_links = relationship("TParagraphRefLink", foreign_keys='TParagraphRefLink.child_id', back_populates='parent')
    entity_links = relationship("TParagraphEntityLink")
    is_del = Column(Integer)

    @property
    def children(self):
        return [link.child for link in self.children_links]

    @property
    def parent(self):
        if self.parent_link:
            return self.parent_link.parent
        return None

    def refs(self):
        return [link.child for link in self.ref_links]


class TParagraphLink(Base):
    """
    段落link表
    """
    __tablename__ = 't_paragraph_links'
    id = Column(Integer, primary_key=True)
    parent_id = Column(Integer, ForeignKey('t_paragraphs.id'))
    parent = relationship("TParagraph", foreign_keys=[parent_id], back_populates="children_links")
    child_id = Column(Integer, ForeignKey('t_paragraphs.id'))
    child = relationship("TParagraph", foreign_keys=[child_id], back_populates="parent_link")
    is_del = Column(Integer)


class TParagraphRefLink(Base):
    """
    段落引用link表
    """
    __tablename__ = 't_paragraph_ref_links'
    id = Column(Integer, primary_key=True)
    parent_id = Column(Integer, ForeignKey('t_paragraphs.id'))
    parent = relationship("TParagraph", foreign_keys=[parent_id], back_populates="ref_links")
    child_id = Column(Integer, ForeignKey('t_paragraphs.id'))
    child = relationship("TParagraph", foreign_keys=[child_id], viewonly=True)
    is_del = Column(Integer)


class TParagraphEntityLink(Base):
    """
    段落实体link表
    """
    __tablename__ = 't_paragraph_entity_links'
    id = Column(Integer, primary_key=True)
    paragraph_id = Column(Integer, ForeignKey('t_paragraphs.id'))
    paragraph = relationship("TParagraph", foreign_keys=[paragraph_id], back_populates="entity_links")
    entity_id = Column(Integer, ForeignKey('t_entities.id'))
    entity = relationship("TEntity", foreign_keys=[entity_id])
    is_del = Column(Integer)


class TDoc(Base):
    """
    文档表
    """
    __tablename__ = 't_docs'
    id = Column(Integer, primary_key=True)
    file = Column(Text)
    file_name = Column(Text)
    paragraphs = relationship("TParagraph")
    is_del = Column(Integer)


class TEntity(Base):
    """
    实体表
    """
    __tablename__ = 't_entities'
    id = Column(Integer, primary_key=True)
    name = Column(Text)
    prompts = Column(Text)
    type = Column(Text)
    doc_type = Column(Text)


class TEntityLink(Base):
    __tablename__ = 't_entity_links'
    id = Column(Integer, primary_key=True)
    entity_id = Column(Integer, ForeignKey('t_entities.id'))
    is_del = Column(Integer)


def init_doc_db():
    """
    初始化文档数据库
    :return: 数据库session实例
    """
    # mysql
    Log.info("连接并初始化文档数据库...")
    engine = create_engine('mysql+pymysql://root:123456@192.168.3.145:3306/knowledgebase', echo=False)
    Base.metadata.create_all(engine)
    SessionFactory = sessionmaker(bind=engine)
    Session = scoped_session(SessionFactory)
    session = Session()
    return session