from sqlalchemy.orm import Session

from models.db_models import MessageModel, SessionModel


class SessionStore:
    def __init__(self, db: Session):
        self.db = db

    def create_session(self, title: str = "新对话") -> SessionModel:
        session = SessionModel(title=title)
        self.db.add(session)
        self.db.commit()
        self.db.refresh(session)
        return session

    def list_sessions(self) -> list[SessionModel]:
        return (
            self.db.query(SessionModel)
            .order_by(SessionModel.updated_at.desc())
            .all()
        )

    def get_session(self, session_id: str) -> SessionModel | None:
        return self.db.get(SessionModel, session_id)

    def delete_session(self, session_id: str) -> bool:
        session = self.get_session(session_id)
        if not session:
            return False
        self.db.delete(session)
        self.db.commit()
        return True

    def update_session_title(self, session_id: str, title: str) -> None:
        session = self.get_session(session_id)
        if session:
            session.title = title
            self.db.commit()

    def get_messages(self, session_id: str) -> list[MessageModel]:
        return (
            self.db.query(MessageModel)
            .filter(MessageModel.session_id == session_id)
            .order_by(MessageModel.created_at)
            .all()
        )

    def get_recent_messages(
        self, session_id: str, max_turns: int
    ) -> list[MessageModel]:
        messages = self.get_messages(session_id)
        if max_turns <= 0:
            return messages
        return messages[-(max_turns * 2) :]

    def add_message(self, session_id: str, role: str, content: str) -> MessageModel:
        message = MessageModel(session_id=session_id, role=role, content=content)
        self.db.add(message)
        session = self.get_session(session_id)
        if session:
            if role == "user" and session.title == "新对话":
                session.title = content[:30] + ("..." if len(content) > 30 else "")
        self.db.commit()
        self.db.refresh(message)
        return message