import datetime

from sqlalchemy import (
    Column,
    String,
    Text,
    DateTime,
    Integer,
    JSON,
)

from models.base import Base, BaseMixin, with_session


class AILog(Base, BaseMixin):
    __tablename__ = "ai_log"

    job_id = Column(String(36), nullable=False, index=True)
    status = Column(Text, nullable=False)
    prompt = Column(Text, nullable=False)
    result = Column(Text, nullable=False)
    model = Column(String(36), nullable=False)
    openai_api_key = Column(String(64), nullable=False)
    app_id = Column(Integer)
    org_id = Column(String(64))
    session_id = Column(String(64))
    query_id = Column(String(64))
    task_id = Column(String(64))
    module = Column(String(64))
    latency = Column(Integer)
    input_token = Column(Integer)
    output_token = Column(Integer)
    created_by = Column(String(64))

    @classmethod
    def upsert(
        cls, created_at, job_id, status, prompt, result, model, openai_api_key,
        app_id, org_id, session_id, query_id, task_id, module, latency, input_token, output_token, created_by,
    ):
        with with_session() as s:
            item = cls(
                created_at=created_at, job_id=job_id, status=status, prompt=str(prompt), result=str(result),
                model=str(model), openai_api_key=openai_api_key,
                app_id=app_id, org_id=org_id, session_id=session_id,
                query_id=query_id, task_id=task_id, module=module, latency=latency,
                input_token=input_token, output_token=output_token, created_by=created_by,
            )
            s.add(item)


class HistoryLog(Base, BaseMixin):
    __tablename__ = "history_log"

    created_at = Column(DateTime, nullable=False, default=datetime.datetime.now, index=True)
    status = Column(String(32), nullable=False)  # searchable, but we do not add index because its cardinality is low
    trace_id = Column(String(36), nullable=False, index=True)
    _type = Column(String(32), nullable=False, index=True)
    org_id = Column(String(64), nullable=False, default="", index=True)
    cluster_id = Column(String(64), nullable=False, default="", index=True)
    database = Column(String(64), nullable=False, default="", index=True)
    question = Column(String(1024), nullable=False, default="")  # searchable, but we do not add index because search it using LIKE
    start_at = Column(DateTime, nullable=False)
    session_id = Column(Integer, nullable=False, default=0)
    data_summary_id = Column(Integer, nullable=False, default=0)
    input = Column(JSON, nullable=False, default={})
    output = Column(JSON, nullable=False, default={})
    recalled_knowledge_data = Column(JSON, nullable=False, default=None)
    settings = Column(JSON, nullable=False, default={})
    token_count = Column(Integer, nullable=False, default=0)

    @classmethod
    def search(cls, s, trace_id=None, _type=None, org_id=None, cluster_id=None, database=None, question=None, start_at=None, end_at=None, status=None, limit=25, offset=0, download=False):
        query = s.query(cls)
        if trace_id:
            query = query.filter(cls.trace_id == trace_id)
        if _type:
            query = query.filter(cls._type == _type)
        if org_id:
            query = query.filter(cls.org_id == org_id)
        if cluster_id:
            query = query.filter(cls.cluster_id == cluster_id)
        if database:
            query = query.filter(cls.database == database)
        if question:
            query = query.filter(cls.question.ilike(f"%{question}%"))
        if start_at:
            query = query.filter(cls.created_at >= start_at)
        if end_at:
            query = query.filter(cls.created_at <= end_at)
        if status:
            query = query.filter(cls.status == status)

        count = query.count()
        items = query.order_by(cls.id.desc()).all() if download else query.order_by(cls.id.desc()).limit(limit).offset(offset).all()

        return items, count