"""数据文件的存储：文件放本地目录，元数据启用 MongoDB 时进集合 data_files，否则放进程内存

目录结构是 DATA_FILE_DIR/<file_id>/，里面有原文件 source 和每张表的 1.parquet、2.parquet……
DuckDB 只能读本地文件，所以文件不进 MongoDB；单实例部署，本地磁盘就够了
会话内的表名 t1、t2……由每个会话一个的计数器分配，文件删除、过期后表名也不再使用
"""

from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Protocol

from bson.codec_options import CodecOptions
from pymongo import ASCENDING, ReturnDocument

from app.schemas.data_file import DataFileRecord, DataTable, FileInsights

# 原文件不带扩展名：calamine 见到 .xls 等扩展名会按扩展名打开，用户的扩展名和内容不符时会失败
SOURCE_NAME = "source"


def table_path(root: Path, file_id: str, table: DataTable) -> Path:
    return root / file_id / table.parquet


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


class DataFileStore(Protocol):
    """内存和 MongoDB 实现同一组方法；过期的记录一律当作不存在"""

    async def insert(self, record: DataFileRecord) -> None: ...

    async def get(self, file_id: str) -> DataFileRecord | None:
        """未过期且未删除的记录"""
        ...

    async def find(self, conversation_id: str) -> list[DataFileRecord]:
        """会话内未过期且未删除的记录，按上传时间排列"""
        ...

    async def reserve_table_numbers(self, conversation_id: str, count: int) -> int:
        """为会话预留 count 个连续的表号，返回第一个；表号从 1 开始，不会重复分配"""
        ...

    async def mark_deleted(self, file_id: str) -> None: ...

    async def update_insights(self, file_id: str, insights: FileInsights) -> None:
        """记录已过期或不存在时什么都不做"""
        ...

    async def live_ids(self) -> set[str]:
        """未过期且未删除的 file_id，清理目录时用"""
        ...


class InMemoryDataFileStore:
    def __init__(self) -> None:
        self._records: dict[str, DataFileRecord] = {}
        # 会话 → 已分配的最大表号
        self._table_numbers: dict[str, int] = {}

    def _purge_expired(self) -> None:
        now = _now()
        for file_id in [
            key for key, record in self._records.items() if record.expires_at <= now
        ]:
            del self._records[file_id]

    async def insert(self, record: DataFileRecord) -> None:
        self._purge_expired()
        self._records[record.file_id] = record

    async def get(self, file_id: str) -> DataFileRecord | None:
        self._purge_expired()
        record = self._records.get(file_id)
        return None if record is None or record.deleted else record

    async def find(self, conversation_id: str) -> list[DataFileRecord]:
        self._purge_expired()
        records = [
            record
            for record in self._records.values()
            if record.conversation_id == conversation_id and not record.deleted
        ]
        return sorted(records, key=lambda record: record.created_at)

    async def reserve_table_numbers(self, conversation_id: str, count: int) -> int:
        first = self._table_numbers.get(conversation_id, 0) + 1
        self._table_numbers[conversation_id] = first + count - 1
        return first

    async def mark_deleted(self, file_id: str) -> None:
        if (record := self._records.get(file_id)) is not None:
            self._records[file_id] = record.model_copy(update={"deleted": True})

    async def update_insights(self, file_id: str, insights: FileInsights) -> None:
        if (record := self._records.get(file_id)) is not None:
            self._records[file_id] = record.model_copy(update={"insights": insights})

    async def live_ids(self) -> set[str]:
        self._purge_expired()
        return {key for key, record in self._records.items() if not record.deleted}


class MongoDataFileStore:
    def __init__(self, db: Any, prefix: str = ""):
        # 读回的时间统一带 UTC 时区；集合名前缀用于测试隔离
        options = CodecOptions(tz_aware=True, tzinfo=UTC)
        self._files = db.get_collection(f"{prefix}data_files", codec_options=options)
        # 每个会话一条：_id 是会话 ID，tables 是已分配的最大表号；会话一直保留，计数器也不过期
        self._table_numbers = db.get_collection(f"{prefix}data_file_table_numbers")

    async def ensure_indexes(self) -> None:
        # 过期记录由 TTL 索引删除；TTL 每分钟才扫一次，查询时仍按 expires_at 过滤
        await self._files.create_index("expires_at", expireAfterSeconds=0)
        await self._files.create_index(
            [("conversation_id", ASCENDING), ("created_at", ASCENDING)]
        )

    async def insert(self, record: DataFileRecord) -> None:
        await self._files.insert_one({"_id": record.file_id, **record.model_dump()})

    async def get(self, file_id: str) -> DataFileRecord | None:
        document = await self._files.find_one(
            {"_id": file_id, "deleted": False, "expires_at": {"$gt": _now()}}
        )
        return None if document is None else _record(document)

    async def find(self, conversation_id: str) -> list[DataFileRecord]:
        query = {
            "conversation_id": conversation_id,
            "deleted": False,
            "expires_at": {"$gt": _now()},
        }
        documents = await (
            self._files.find(query)
            .sort([("created_at", ASCENDING), ("_id", ASCENDING)])
            .to_list()
        )
        return [_record(document) for document in documents]

    async def reserve_table_numbers(self, conversation_id: str, count: int) -> int:
        document = await self._table_numbers.find_one_and_update(
            {"_id": conversation_id},
            {"$inc": {"tables": count}},
            upsert=True,
            return_document=ReturnDocument.AFTER,
        )
        return document["tables"] - count + 1

    async def mark_deleted(self, file_id: str) -> None:
        await self._files.update_one({"_id": file_id}, {"$set": {"deleted": True}})

    async def update_insights(self, file_id: str, insights: FileInsights) -> None:
        await self._files.update_one(
            {"_id": file_id}, {"$set": {"insights": insights.model_dump()}}
        )

    async def live_ids(self) -> set[str]:
        return set(
            await self._files.distinct(
                "_id", {"deleted": False, "expires_at": {"$gt": _now()}}
            )
        )


def _record(document: dict[str, Any]) -> DataFileRecord:
    return DataFileRecord.model_validate(
        {key: value for key, value in document.items() if key != "_id"}
    )


async def open_data_file_store(db: Any) -> DataFileStore:
    """db 为空（未启用 MongoDB）时用内存实现，和会话记录一致"""
    if db is None:
        return InMemoryDataFileStore()
    store = MongoDataFileStore(db)
    await store.ensure_indexes()
    return store