ForcePilot/backend/package/yuxi/channels/session_mapper.py

from __future__ import annotations

import uuid as uuid_lib

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

from yuxi.channels.models import ChannelMessage
from yuxi.storage.postgres.models_business import User
from yuxi.storage.postgres.models_channels import ChannelThreadMapping, ChannelUserMapping
from yuxi.utils.datetime_utils import utc_now_naive
from yuxi.utils.logging_config import logger

VIRTUAL_DEPARTMENT_ID = -1
USER_SOURCE_PREFIX = "channel:"


class SessionMapper:
    def __init__(self, db: AsyncSession, department_id: int = VIRTUAL_DEPARTMENT_ID):
        self.db = db
        self.department_id = department_id

    async def resolve_user(self, message: ChannelMessage) -> str:
        identity = message.identity
        mapping = await self._get_user_mapping(identity.channel_id, identity.channel_user_id)
        if mapping:
            return mapping.internal_user_id

        internal_user_id = f"ch_{identity.channel_id}_{uuid_lib.uuid4().hex[:8]}"
        username = f"{identity.channel_id}_{identity.channel_user_id}"
        source = f"{USER_SOURCE_PREFIX}{identity.channel_id}"

        stmt = (
            pg_insert(User)
            .values(
                username=username,
                user_id=internal_user_id,
                password_hash="",
                role="user",
                department_id=self.department_id,
                source=source,
            )
            .on_conflict_do_update(
                index_elements=["username"],
                set_={"username": username},
            )
            .returning(User.user_id)
        )
        result = await self.db.execute(stmt)
        actual_user_id = result.scalar_one()

        try:
            mapping = ChannelUserMapping(
                channel_id=identity.channel_id,
                channel_user_id=identity.channel_user_id,
                internal_user_id=actual_user_id,
            )
            self.db.add(mapping)
            await self.db.commit()

            logger.info(
                f"Created channel user mapping: {identity.channel_id}/{identity.channel_user_id} -> {actual_user_id}"
            )
            return actual_user_id

        except IntegrityError:
            await self.db.rollback()
            mapping = await self._get_user_mapping(identity.channel_id, identity.channel_user_id)
            if mapping:
                return mapping.internal_user_id
            raise

    async def resolve_thread(self, message: ChannelMessage, internal_user_id: str) -> str:
        identity = message.identity
        mapping = await self._get_thread_mapping(identity.channel_id, identity.channel_chat_id, internal_user_id)
        if mapping:
            mapping.last_active_at = utc_now_naive()
            await self.db.flush()
            return mapping.thread_id

        thread_id = str(uuid_lib.uuid4())

        try:
            mapping = ChannelThreadMapping(
                channel_id=identity.channel_id,
                channel_chat_id=identity.channel_chat_id,
                internal_user_id=internal_user_id,
                thread_id=thread_id,
                last_active_at=utc_now_naive(),
            )
            self.db.add(mapping)
            await self.db.commit()

            logger.info(
                f"Created channel thread mapping: "
                f"{identity.channel_id}/{identity.channel_chat_id}/{internal_user_id} -> {thread_id}"
            )
            return thread_id

        except IntegrityError:
            await self.db.rollback()
            mapping = await self._get_thread_mapping(identity.channel_id, identity.channel_chat_id, internal_user_id)
            if mapping:
                return mapping.thread_id
            raise

    async def reset_thread(self, message: ChannelMessage, internal_user_id: str) -> str:
        identity = message.identity
        mapping = await self._get_thread_mapping(identity.channel_id, identity.channel_chat_id, internal_user_id)
        if mapping:
            new_thread_id = str(uuid_lib.uuid4())
            mapping.thread_id = new_thread_id
            mapping.last_active_at = utc_now_naive()
            await self.db.commit()
            logger.info(f"Reset thread: {identity.channel_id}/{identity.channel_chat_id} -> {new_thread_id}")
            return new_thread_id

        return await self.resolve_thread(message, internal_user_id)

    async def _get_user_mapping(self, channel_id: str, channel_user_id: str) -> ChannelUserMapping | None:
        result = await self.db.execute(
            select(ChannelUserMapping).where(
                ChannelUserMapping.channel_id == channel_id,
                ChannelUserMapping.channel_user_id == channel_user_id,
            )
        )
        return result.scalar_one_or_none()

    async def _get_thread_mapping(
        self, channel_id: str, channel_chat_id: str, internal_user_id: str
    ) -> ChannelThreadMapping | None:
        result = await self.db.execute(
            select(ChannelThreadMapping).where(
                ChannelThreadMapping.channel_id == channel_id,
                ChannelThreadMapping.channel_chat_id == channel_chat_id,
                ChannelThreadMapping.internal_user_id == internal_user_id,
            )
        )
        return result.scalar_one_or_none()