import json
from models.base import Base, BaseMixin
from sqlalchemy import (
    Boolean,
    Column,
    Integer,
    VARCHAR,
    JSON,
    String,
    or_
)
from consts.consts import ModuleType, KnowledgeType


class KnowledgeBase(Base, BaseMixin):
    __tablename__ = 'knowledge_base'

    org_id = Column(String(64), nullable=False)
    cluster_id = Column(String(64), nullable=False, default=0, index=True)
    db_name = Column(String(64), nullable=False)
    description = Column(String(4096), nullable=False, default="")
    is_default = Column(Boolean, nullable=False, default=False)

    @classmethod
    def create(cls, s, org_id, cluster_id, db_name, description):
        item = cls(org_id=org_id, cluster_id=cluster_id, db_name=db_name, description=description)
        s.add(item)
        s.flush()
        return item

    def update_description_and_default(self, s, description, is_default):
        # set all other knowledge base to non-default
        if is_default:
            s.query(KnowledgeBase).filter(
                KnowledgeBase.cluster_id == self.cluster_id,
                KnowledgeBase.db_name == self.db_name,
            ).update({KnowledgeBase.is_default: False})

        self.description = description
        self.is_default = is_default
        s.flush()

    @classmethod
    def get_all_by_cluster_id_and_db_name(cls, s, cluster_id, db_name):
        return s.query(cls).filter(cls.cluster_id == cluster_id, cls.db_name == db_name).order_by(cls.id.desc()).all()

    def to_json(self, simple=True):
        return {
            "knowledge_base_id": self.id,
            "description": self.description,
            "default": self.is_default,
        }

    def set_default(self, s):
        s.query(KnowledgeBase).filter(KnowledgeBase.cluster_id == self.cluster_id).update({KnowledgeBase.is_default: False})
        self.is_default = True
        s.flush()


class DomainKnowledge(Base, BaseMixin):
    __tablename__ = 'domain_knowledge'

    raw_data = Column(JSON(), nullable=False)
    type = Column(VARCHAR(256), nullable=False)
    app_id = Column(Integer(), nullable=False, default=0, index=True)
    module = Column(VARCHAR(256), nullable=False, default=ModuleType.ALL, index=True)
    meta_data = Column(JSON(), nullable=False)
    created_by = Column(VARCHAR(1024), nullable=True)
    knowledge_base_id = Column(Integer, nullable=False, index=True)

    @classmethod
    def create(cls, s, raw_data, type, app_id, module, meta_data, created_by, knowledge_base_id):
        item = cls(
            raw_data=raw_data, type=type, app_id=app_id, module=module, meta_data=meta_data,
            created_by=created_by, knowledge_base_id=knowledge_base_id,
        )
        s.add(item)
        s.flush()
        return item

    @classmethod
    def get_all_by_app_id(cls, s, app_id, knowledge_base_id, type=None):
        filters = [cls.knowledge_base_id == knowledge_base_id]
        if type:
            filters.append(cls.type == type)

        return s.query(cls).filter(*filters).order_by(cls.created_at.desc()).all()

    @classmethod
    def get_by_app_id_and_module(cls, s, type, app_id=0, knowledge_base_ids=[], module=ModuleType.ALL):
        # Param check
        if not type:
            return []
        # app_id default is 0
        filters = [cls.type == type]
        # compatible for old logic
        # if knowledge items are binded to app_id, it is for global use in this app. if they are binded to knowledge_base_id, it is used for this db.
        filters.append(or_(cls.app_id == app_id, cls.knowledge_base_id.in_(knowledge_base_ids)))
        # retrievalTable and text2sql use the same instruction
        if type == KnowledgeType.INSTRUCTION and (module == ModuleType.TASK_RETRIEVAL_TABLES or module == ModuleType.TASK_TEXT2SQL):
            filters.append(or_(cls.module == ModuleType.ALL, cls.module == ModuleType.TASK_RETRIEVAL_TABLES, cls.module == ModuleType.TASK_TEXT2SQL))
        else:
            filters.append(or_(cls.module == ModuleType.ALL, cls.module == module))
        # Query by conditions
        return s.query(cls.raw_data, cls.meta_data, cls.id, cls.type, cls.knowledge_base_id).filter(*filters).order_by(cls.created_at.desc()).all()
        
        
    @classmethod
    def delete_by_id(cls, s, id):
        s.query(cls).filter(cls.id == id).delete()
        s.flush()

    def update_data(self, s, module, raw_data, meta_data):
        self.raw_data = json.dumps(raw_data)
        self.meta_data = json.dumps(meta_data)
        self.module = module
        s.flush()

    def to_json(self):
        return {
            "knowledge_data_id": self.id,
            "knowledge_base_id": self.knowledge_base_id,
            'meta_data': self.meta_data,
            'raw_data': self.raw_data,
        }