app/db/models.py

# app/db/models.py

import uuid
from datetime import datetime
import enum as py_enum # Use alias

from sqlalchemy import (
    Column, String, DateTime, ForeignKey, Text, Enum as SQLAlchemyEnum,
    Integer # Import Integer type
)
# Import PostgreSQL specific types
from sqlalchemy.dialects.postgresql import UUID as PG_UUID, JSONB
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func

# Import the Base class from database.py
from .database import Base
# Import Enums from schemas.py
from .schemas import MessageRole, FeedbackTypeEnum # Import both enums now

class ChatSession(Base):
    """SQLAlchemy model for chat sessions (using _v2 table)."""
    __tablename__ = "chat_sessions_v2"

    id = Column(PG_UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    name = Column(String, index=True, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    messages = relationship("ChatMessage", back_populates="session", cascade="all, delete-orphan", lazy="select")

    def __repr__(self):
        return f"<ChatSessionV2(id={self.id}, name='{self.name}')>"

class ChatMessage(Base):
    """SQLAlchemy model for chat messages (using _v2 table)."""
    __tablename__ = "chat_messages_v2"

    id = Column(PG_UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    session_id = Column(PG_UUID(as_uuid=True), ForeignKey("chat_sessions_v2.id"), nullable=False, index=True)
    role = Column(SQLAlchemyEnum(MessageRole, name="message_role_enum_v2", create_type=False), nullable=False)
    content = Column(Text, nullable=False)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)

    # --- Columns for Bot messages (Context Tracking) ---
    original_query = Column(Text, nullable=True, comment="Original user query that led to this bot response")
    retrieved_context_ids = Column(JSONB, nullable=True, comment="List of all chunk IDs retrieved from search for the original query")
    used_context_ids = Column(JSONB, nullable=True, comment="List of chunk IDs whose content was actually used to generate this response")
    attempt_number = Column(Integer, default=1, nullable=False, comment="Tracks which regeneration attempt this message represents (1 = initial, 2 = first regen, etc.)")
    cumulative_used_context_ids = Column(JSONB, nullable=True, comment="Set of all context IDs used in this and previous attempts for the same original_query")
    # --- END Context Tracking Columns ---

    session = relationship("ChatSession", back_populates="messages")
    feedback = relationship("FeedbackLog", back_populates="message", uselist=False, cascade="all, delete-orphan", lazy="select")

    def __repr__(self):
        attempt_str = f", attempt={self.attempt_number}" if self.role == MessageRole.BOT else ""
        return f"<ChatMessageV2(id={self.id}, role='{self.role}'{attempt_str}, session_id={self.session_id})>"


class FeedbackLog(Base):
    """SQLAlchemy model for feedback logs (using _v2 table)."""
    __tablename__ = "feedback_logs_v2"

    id = Column(PG_UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    message_id = Column(PG_UUID(as_uuid=True), ForeignKey("chat_messages_v2.id"), nullable=False, unique=True)
    feedback_type = Column(SQLAlchemyEnum(FeedbackTypeEnum, name="feedback_type_enum_v2", create_type=False), nullable=False)
    feedback_comment = Column(Text, nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    message = relationship("ChatMessage", back_populates="feedback")

    def __repr__(self):
        return f"<FeedbackLogV2(id={self.id}, message_id='{self.message_id}', type='{self.feedback_type}')>"