from sqlalchemy import Column, Integer, String, DateTime, JSON, Text, ForeignKey
from sqlalchemy.ext.declarative import declarative_base
from datetime import datetime

Base = declarative_base()

class BnkAIReqRes(Base):
    """Stores AI request, response, and callback details"""

    __tablename__ = "bnk_ai_req_res"

    id = Column(Integer, primary_key=True, index=True)
    request_id = Column(String(100), index=True, nullable=False)
    request = Column(JSON, nullable=False)
    response = Column(JSON, nullable=True)
    callback_res = Column(JSON, nullable=True)
    callback_url = Column(Text, nullable=True)
    callback_res_header = Column(JSON, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)


class BnkParsedDocument(Base):
    """Stores LlamaParse extracted text from documents"""
    __tablename__ = "bnk_parsed_documents"

    id = Column(Integer, primary_key=True, index=True)
    ai_req_res_id = Column(Integer, ForeignKey("bnk_ai_req_res.id", ondelete="CASCADE"), nullable=False, index=True)
    source_url = Column(Text, nullable=False)
    extracted_text = Column(Text, nullable=False)
    token_count = Column(Integer, nullable=False)
    chunked = Column(Integer, default=0)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)


class BnkDocumentChunk(Base):
    """Stores chunks and their LLM output"""
    __tablename__ = "bnk_document_chunks"

    id = Column(Integer, primary_key=True, index=True)
    parsed_document_id = Column(Integer, ForeignKey("bnk_parsed_documents.id", ondelete="CASCADE"), nullable=False, index=True)
    chunk_index = Column(Integer, nullable=False)
    chunk_text = Column(Text, nullable=False)
    token_count = Column(Integer, nullable=False)
    llm_output = Column(JSON, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)


class BnkConsolidatedLLMOutput(Base):
    """Stores final consolidated LLM response from all chunks"""
    __tablename__ = "bnk_consolidated_llm_outputs"

    id = Column(Integer, primary_key=True, index=True)
    ai_req_res_id = Column(Integer, ForeignKey("bnk_ai_req_res.id", ondelete="CASCADE"), nullable=False, unique=True)
    consolidated_response = Column(JSON, nullable=False)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)