import re
from dataclasses import dataclass

from sqlalchemy import delete, func, select
from sqlalchemy.ext.asyncio import AsyncSession

from kbqa.chat.models import ChatSession, Message, MessageSource
from kbqa.database import utc_now
from kbqa.documents.models import Chunk, Document


@dataclass(frozen=True, slots=True)
class MessageSourceInput:
    document_id: str
    chunk_id: str
    citation_label: str
    score: float


@dataclass(frozen=True, slots=True)
class MessageSourceView:
    message_id: str
    document_id: str
    chunk_id: str
    citation_label: str
    score: float
    filename: str
    order_index: int
    page_number: int | None
    content: str


class ChatRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def create_session(self, title: str | None) -> ChatSession:
        session = ChatSession(title=title.strip() if title and title.strip() else None)
        self.session.add(session)
        await self.session.commit()
        await self.session.refresh(session)
        return session

    async def get_session(self, session_id: str) -> ChatSession | None:
        return await self.session.get(ChatSession, session_id)

    async def page_sessions(self, offset: int, limit: int) -> tuple[list[ChatSession], int]:
        total = await self.session.scalar(select(func.count(ChatSession.id)))
        rows = await self.session.scalars(
            select(ChatSession).order_by(ChatSession.updated_at.desc()).offset(offset).limit(limit)
        )
        return list(rows), int(total or 0)

    async def load_recent_messages(self, session_id: str, limit: int) -> list[Message]:
        rows = list(
            await self.session.scalars(
                select(Message)
                .where(Message.session_id == session_id, Message.status == "completed")
                .order_by(Message.order_index.desc())
                .limit(limit)
            )
        )
        rows.reverse()
        return rows

    async def page_messages(
        self,
        session_id: str,
        offset: int,
        limit: int,
    ) -> tuple[list[Message], int]:
        total = await self.session.scalar(
            select(func.count(Message.id)).where(Message.session_id == session_id)
        )
        rows = await self.session.scalars(
            select(Message)
            .where(Message.session_id == session_id)
            .order_by(Message.order_index)
            .offset(offset)
            .limit(limit)
        )
        return list(rows), int(total or 0)

    async def sources_for_messages(
        self, message_ids: list[str]
    ) -> dict[str, list[MessageSourceView]]:
        if not message_ids:
            return {}
        rows = await self.session.execute(
            select(MessageSource, Chunk, Document)
            .join(Chunk, Chunk.id == MessageSource.chunk_id)
            .join(Document, Document.id == MessageSource.document_id)
            .where(MessageSource.message_id.in_(message_ids))
            .order_by(MessageSource.citation_label)
        )
        result: dict[str, list[MessageSourceView]] = {}
        for source, chunk, document in rows:
            view = MessageSourceView(
                message_id=source.message_id,
                document_id=source.document_id,
                chunk_id=source.chunk_id,
                citation_label=source.citation_label,
                score=source.score,
                filename=document.filename,
                order_index=chunk.order_index,
                page_number=chunk.page_number,
                content=chunk.content,
            )
            result.setdefault(source.message_id, []).append(view)
        return result

    async def begin_turn(self, session_id: str, query: str) -> tuple[Message, Message]:
        chat_session = await self.session.get(ChatSession, session_id)
        if chat_session is None:
            raise LookupError("session not found")
        maximum = await self.session.scalar(
            select(func.max(Message.order_index)).where(Message.session_id == session_id)
        )
        next_order = int(maximum if maximum is not None else -1) + 1
        user = Message(
            session_id=session_id,
            role="user",
            content=query,
            status="completed",
            order_index=next_order,
        )
        assistant = Message(
            session_id=session_id,
            role="assistant",
            content="",
            status="generating",
            order_index=next_order + 1,
        )
        if chat_session.title is None:
            chat_session.title = query.strip()[:30]
        chat_session.updated_at = utc_now()
        self.session.add_all([user, assistant])
        await self.session.commit()
        await self.session.refresh(user)
        await self.session.refresh(assistant)
        return user, assistant

    async def complete_answer(
        self,
        message_id: str,
        content: str,
        sources: list[MessageSourceInput],
    ) -> None:
        message = await self.session.get(Message, message_id)
        if message is None:
            raise LookupError("message not found")
        await self.session.execute(
            delete(MessageSource).where(MessageSource.message_id == message_id)
        )
        message.content = content
        message.status = "completed"
        message.error_code = None
        self.session.add_all(
            [
                MessageSource(
                    message_id=message_id,
                    document_id=source.document_id,
                    chunk_id=source.chunk_id,
                    citation_label=source.citation_label,
                    score=source.score,
                )
                for source in sources
            ]
        )
        chat_session = await self.session.get(ChatSession, message.session_id)
        if chat_session is not None:
            chat_session.updated_at = utc_now()
        await self.session.commit()

    async def fail_answer(self, message_id: str, error_code: str) -> None:
        message = await self.session.get(Message, message_id)
        if message is None:
            return
        message.status = "failed"
        message.error_code = error_code
        await self.session.commit()

    async def delete_session(self, session_id: str) -> bool:
        chat_session = await self.session.get(ChatSession, session_id)
        if chat_session is None:
            return False
        await self.session.delete(chat_session)
        await self.session.commit()
        return True

    @staticmethod
    def citation_labels(content: str) -> set[str]:
        return set(re.findall(r"\[(S\d+)]", content))
