"""ChatSessionRepository — /chat 세션 상태 DB 영속화 (P8-A).

인메모리 대신 learning.chat_sessions 에 진행 상태를 저장 → 서버 재시작/멀티워커 안전.
company_id 스코프. ChatSession(dataclass) ↔ row 직렬화.
"""

import uuid
from typing import Optional

from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession

from common.database.db_session_manager import DB_SESSION_MNG
from common.database.model.models import ChatSessionRow
from common.enums import DBType, DBWRType, ErrorType
from common.utils.gtime import GTime
from negotiation.chat.service.chat_engine import ChatSession


class ChatSessionRepository:
    def __init__(self, company_id: str):
        self.company_id = company_id

    async def get(self, session_id: Optional[str]) -> Optional[ChatSession]:
        if not session_id:
            return None

        def _q(s: AsyncSession):
            stmt = select(ChatSessionRow).where(
                ChatSessionRow.session_id == session_id,
                ChatSessionRow.company_id == self.company_id,
            ).limit(1)
            return DB_SESSION_MNG.execute(s, stmt)

        err, rows = await DB_SESSION_MNG.execute_lambda(DBType.MAIN.value, DBWRType.DB_READ.value, _q)
        if err != ErrorType.SUCCESS or not rows:
            return None
        r = rows[0]
        return ChatSession(
            session_id=str(r.session_id), tenant_id=r.tenant_id, company_id=r.company_id,
            rq_type=r.rq_type, step=r.step, context=dict(r.context or {}),
            used_action_ids=set(r.used_action_ids or []), action_space_size=r.action_space_size,
            ended=r.ended,
        )

    async def save(self, session: ChatSession) -> ErrorType:
        values = dict(
            session_id=session.session_id, company_id=session.company_id, tenant_id=session.tenant_id,
            rq_type=session.rq_type, step=session.step, context=session.context,
            used_action_ids=sorted(session.used_action_ids), action_space_size=session.action_space_size,
            ended=session.ended, updated_at=GTime.UTC(),
        )

        def _do(s: AsyncSession):
            stmt = pg_insert(ChatSessionRow.__table__).values(**values).on_conflict_do_update(
                index_elements=[ChatSessionRow.session_id],
                set_={k: values[k] for k in ("step", "context", "used_action_ids", "ended", "updated_at")},
            )
            return DB_SESSION_MNG.add(s, stmt)

        return await DB_SESSION_MNG.execute_lambda_run([DBType.MAIN.value], [_do])