"""会话存储的 MongoDB 实现，方法语义与内存实现一致"""

from datetime import UTC, datetime, timedelta
from uuid import uuid4

from bson.codec_options import CodecOptions
from pymongo import ASCENDING, DESCENDING
from pymongo.asynchronous.database import AsyncDatabase

from app.config import settings
from app.schemas.conversation import Turn

# 单次清理的空会话上限，避免积压时一次请求扫描过多
PURGE_BATCH = 200


def _now() -> datetime:
    return datetime.now(UTC)


class MongoConversationStore:
    def __init__(
        self, db: AsyncDatabase, empty_ttl: int | None = None, prefix: str = ""
    ):
        ttl = empty_ttl if empty_ttl is not None else settings.conversation_empty_ttl
        self._empty_ttl = timedelta(seconds=ttl)
        # 读回的时间统一带 UTC 时区，与 Turn 的时间保持一致；集合名前缀用于测试隔离
        options = CodecOptions(tz_aware=True, tzinfo=UTC)
        self._conversations = db.get_collection(
            f"{prefix}conversations", codec_options=options
        )
        self._turns = db.get_collection(f"{prefix}turns", codec_options=options)

    async def ensure_indexes(self) -> None:
        # 创建索引是幂等的，启动时调用一次即可
        # 同一任务只有一轮，恢复后的重复写入靠唯一索引落到同一条记录
        await self._turns.create_index(
            [("conversation_id", ASCENDING), ("thread_id", ASCENDING)], unique=True
        )
        await self._turns.create_index(
            [("conversation_id", ASCENDING), ("created_at", ASCENDING)]
        )
        # 空会话清理按创建时间扫描
        await self._conversations.create_index("created_at")
        # 按用户查询会话
        await self._conversations.create_index("user_id")

    async def _purge_expired(self) -> None:
        # 只清理没有任何轮次且超过保留时长的空会话，有对话的会话不受影响
        deadline = _now() - self._empty_ttl
        expired = await (
            self._conversations.find({"created_at": {"$lt": deadline}}, {"_id": 1})
            .limit(PURGE_BATCH)
            .to_list()
        )
        ids = [item["_id"] for item in expired]
        if not ids:
            return
        used = set(
            await self._turns.distinct(
                "conversation_id", {"conversation_id": {"$in": ids}}
            )
        )
        empty = [item for item in ids if item not in used]
        if empty:
            await self._conversations.delete_many({"_id": {"$in": empty}})

    async def create(self, user_id: int) -> str:
        await self._purge_expired()
        conversation_id = str(uuid4())
        # user_id 来自用户中心，会话只对创建者可见
        await self._conversations.insert_one(
            {"_id": conversation_id, "created_at": _now(), "user_id": user_id}
        )
        return conversation_id

    async def exists(self, conversation_id: str, user_id: int) -> bool:
        # 不属于该用户的会话与不存在同样处理，不泄露会话是否存在
        document = await self._conversations.find_one(
            {"_id": conversation_id, "user_id": user_id}
        )
        if document is None:
            return False
        if document["created_at"] >= _now() - self._empty_ttl:
            return True
        # 超过空会话保留时长后，只有已产生对话的会话仍视为存在
        return (
            await self._turns.find_one({"conversation_id": conversation_id}, {"_id": 1})
            is not None
        )

    async def load_turns(self, conversation_id: str, limit: int | None) -> list[Turn]:
        if limit == 0:
            return []
        cursor = self._turns.find({"conversation_id": conversation_id})
        if limit is None:
            documents = await cursor.sort(
                [("created_at", ASCENDING), ("_id", ASCENDING)]
            ).to_list()
        else:
            # 倒序取最近 limit 轮，再反转为时间正序
            documents = await (
                cursor.sort([("created_at", DESCENDING), ("_id", DESCENDING)])
                .limit(limit)
                .to_list()
            )
            documents.reverse()
        return [
            Turn.model_validate({k: v for k, v in item.items() if k != "_id"})
            for item in documents
        ]

    async def save_turn(self, turn: Turn) -> None:
        data = turn.model_dump()
        created_at = data.pop("created_at")
        # 同一任务再次写入（澄清恢复后）保留首次创建时间，只更新内容
        await self._turns.update_one(
            {"conversation_id": turn.conversation_id, "thread_id": turn.thread_id},
            {"$set": data, "$setOnInsert": {"created_at": created_at}},
            upsert=True,
        )