from __future__ import annotations

from uuid import UUID

from fastapi import HTTPException
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from models.agent_chat_session import AgentChatSession


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

    async def get_session_owner(self, *, session_id: str) -> str:
        try:
            session_uuid = UUID(session_id)
        except ValueError as exc:
            raise HTTPException(status_code=422, detail="Invalid session_id") from exc

        stmt = select(AgentChatSession.user_id).where(
            AgentChatSession.id == session_uuid
        )
        owner_id = (await self._session.execute(stmt)).scalar_one_or_none()
        if owner_id is None:
            raise HTTPException(status_code=404, detail="Session not found")
        return str(owner_id)

    async def create_session_for_user(self, *, user_id: str) -> str:
        try:
            user_uuid = UUID(user_id)
        except ValueError as exc:
            raise HTTPException(status_code=422, detail="Invalid user_id") from exc

        session = AgentChatSession(
            user_id=user_uuid,
        )
        self._session.add(session)
        await self._session.flush()
        await self._session.refresh(session)
        return str(session.id)

    async def commit(self) -> None:
        await self._session.commit()

    async def rollback(self) -> None:
        await self._session.rollback()

    async def delete_session(self, *, session_id: str) -> None:
        try:
            session_uuid = UUID(session_id)
        except ValueError as exc:
            raise HTTPException(status_code=422, detail="Invalid session_id") from exc
        session = await self._session.get(AgentChatSession, session_uuid)
        if session is not None:
            await self._session.delete(session)
            await self._session.flush()