"""Frontend document upload, preview, and confirmed Milvus ingestion."""
from __future__ import annotations

import logging
import json
import uuid
from inspect import isawaitable
from pathlib import Path

from rag.document_parser import SUPPORTED_EXTENSIONS
from rag.events import publish_knowledge_update
from rag.ingestion import ingest_document_atomic
from rag.preview import preview_document
from rag.chunk_config import resolve_chunk_config
from rag.milvus_collections import KNOWLEDGE_COLLECTIONS
from rag.milvus_delete import delete_document_vectors


logger = logging.getLogger("service.knowledge_base.upload")


class UploadValidationError(ValueError):
    pass


class KnowledgeUploadService:
    def __init__(
        self,
        *,
        storage_dir: str | Path,
        milvus_client,
        embedder,
        publisher,
        max_upload_bytes: int = 10 * 1024 * 1024,
        document_exists=None,
        upload_ttl_seconds: int = 1800,
    ):
        self.storage_dir = Path(storage_dir)
        self.storage_dir.mkdir(parents=True, exist_ok=True)
        self.milvus_client = milvus_client
        self.embedder = embedder
        self.publisher = publisher
        self.max_upload_bytes = max_upload_bytes
        self.document_exists = document_exists
        self.upload_ttl_seconds = upload_ttl_seconds

    async def preview(
        self,
        filename: str,
        content: bytes,
        *,
        strategy: str,
        chunk_size: int | None = None,
        chunk_overlap: int | None = None,
    ) -> dict:
        self.cleanup_expired_uploads()
        suffix = self._validate_upload(filename, content)
        upload_id = uuid.uuid4().hex
        stored_filename = f"{upload_id}{suffix}"
        path = self.storage_dir / stored_filename
        path.write_bytes(content)
        try:
            result = preview_document(
                path,
                strategy=strategy,
                chunk_size=chunk_size,
                chunk_overlap=chunk_overlap,
            )
        except Exception:
            path.unlink(missing_ok=True)
            raise
        self._manifest_path(upload_id).write_text(
            json.dumps(
                {
                    "strategy": strategy,
                    "chunk_size": result["chunk_size"],
                    "chunk_overlap": result["chunk_overlap"],
                },
                ensure_ascii=False,
            ),
            encoding="utf-8",
        )
        return {
            "upload_id": upload_id,
            "stored_filename": stored_filename,
            "filename": Path(filename).name,
            **result,
        }

    def cleanup_expired_uploads(self, *, now: float | None = None) -> int:
        """Remove expired upload inputs and their preview manifests."""
        if self.upload_ttl_seconds <= 0:
            return 0
        if now is None:
            import time

            now = time.time()
        cutoff = now - self.upload_ttl_seconds
        removed = 0
        for path in self.storage_dir.iterdir():
            if not path.is_file() or path.suffix.lower() not in SUPPORTED_EXTENSIONS:
                continue
            if path.stat().st_mtime > cutoff:
                continue
            upload_id = path.stem
            path.unlink(missing_ok=True)
            self._manifest_path(upload_id).unlink(missing_ok=True)
            removed += 1
        return removed

    async def delete_document(self, doc_id: str) -> dict:
        if not doc_id or not doc_id.strip():
            raise UploadValidationError("doc_id不能为空")
        await delete_document_vectors(doc_id, milvus_client=self.milvus_client)
        await publish_knowledge_update(
            self.publisher,
            doc_id=doc_id,
            collection_name="",
            chunk_count=0,
            action="deleted",
        )
        return {"doc_id": doc_id, "deleted": True}

    async def list_documents(self) -> list[dict]:
        documents = {}
        for collection_name in KNOWLEDGE_COLLECTIONS:
            rows = await self.milvus_client.query(
                collection_name=collection_name,
                filter="",
                output_fields=["doc_id", "title", "strategy"],
            )
            for row in rows or []:
                doc_id = row.get("doc_id", "")
                if not doc_id:
                    continue
                document = documents.setdefault(
                    doc_id,
                    {
                        "doc_id": doc_id,
                        "title": row.get("title", ""),
                        "collection_name": collection_name,
                        "strategy": row.get("strategy", ""),
                        "chunk_count": 0,
                    },
                )
                document["chunk_count"] += 1
        return list(documents.values())

    async def get_document(self, doc_id: str) -> dict | None:
        if not doc_id or not doc_id.strip():
            raise UploadValidationError("doc_id不能为空")
        escaped = doc_id.replace("\\", "\\\\").replace('"', '\\"')
        for collection_name in KNOWLEDGE_COLLECTIONS:
            rows = await self.milvus_client.query(
                collection_name=collection_name,
                filter=f'doc_id == "{escaped}"',
                output_fields=["doc_id", "title", "section_title", "strategy", "text"],
            )
            if rows:
                first = rows[0]
                return {
                    "doc_id": first.get("doc_id", doc_id),
                    "title": first.get("title", ""),
                    "collection_name": collection_name,
                    "strategy": first.get("strategy", ""),
                    "chunk_count": len(rows),
                    "chunks": [
                        {
                            "section_title": row.get("section_title") or None,
                            "text": row.get("text", ""),
                        }
                        for row in rows
                    ],
                }
        return None

    async def confirm(
        self,
        *,
        upload_id: str,
        title: str,
        doc_id: str,
        collection_name: str,
        strategy: str,
        chunk_size: int | None = None,
        chunk_overlap: int | None = None,
    ) -> dict:
        self.cleanup_expired_uploads()
        if collection_name not in KNOWLEDGE_COLLECTIONS:
            raise UploadValidationError("知识库集合不在允许范围内")
        path = self._find_upload(upload_id)
        try:
            manifest = json.loads(self._manifest_path(upload_id).read_text(encoding="utf-8"))
            if strategy != manifest["strategy"]:
                raise UploadValidationError("确认策略必须与预览策略一致")
            if chunk_size is None:
                chunk_size = manifest["chunk_size"]
            if chunk_overlap is None:
                chunk_overlap = manifest["chunk_overlap"]
            if self.document_exists is not None:
                exists = self.document_exists(doc_id)
                if isawaitable(exists):
                    exists = await exists
                if exists:
                    raise UploadValidationError("doc_id已存在，禁止重复入库")
            from rag.document_parser import parse_document

            result = await ingest_document_atomic(
                parse_document(path),
                doc_id,
                title,
                collection_name,
                strategy,
                milvus_client=self.milvus_client,
                embedder=self.embedder,
                config=resolve_chunk_config(
                    chunk_size=chunk_size,
                    chunk_overlap=chunk_overlap,
                ),
            )
            try:
                await publish_knowledge_update(
                    self.publisher,
                    doc_id=doc_id,
                    collection_name=collection_name,
                    chunk_count=result["chunk_count"],
                )
            except Exception:
                logger.exception("knowledge update event publish failed for doc_id=%s", doc_id)
                try:
                    await delete_document_vectors(doc_id, milvus_client=self.milvus_client)
                except Exception:
                    logger.exception("failed to compensate vectors for doc_id=%s", doc_id)
                raise
            return {"event_published": True, **result}
        finally:
            path.unlink(missing_ok=True)
            self._manifest_path(upload_id).unlink(missing_ok=True)

    def _validate_upload(self, filename: str, content: bytes) -> str:
        suffix = Path(filename).suffix.lower()
        if suffix not in SUPPORTED_EXTENSIONS:
            raise UploadValidationError(f"Unsupported document type: {suffix or '<none>'}")
        if not content:
            raise UploadValidationError("上传文件不能为空")
        if len(content) > self.max_upload_bytes:
            raise UploadValidationError("上传文件超过大小限制")
        return suffix

    def _find_upload(self, upload_id: str) -> Path:
        if not upload_id or Path(upload_id).name != upload_id:
            raise UploadValidationError("upload_id无效")
        paths = [
            path for path in self.storage_dir.glob(f"{upload_id}.*")
            if path.suffix.lower() in SUPPORTED_EXTENSIONS
        ]
        if len(paths) != 1:
            raise UploadValidationError("上传文件不存在或已过期")
        return paths[0]

    def _manifest_path(self, upload_id: str) -> Path:
        return self.storage_dir / f"{upload_id}.json"