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

import logging
import json
import re
import time
import uuid
from inspect import isawaitable
from pathlib import Path
from uuid import NAMESPACE_URL, uuid5

from rag.document_parser import SUPPORTED_EXTENSIONS
from rag.events import publish_knowledge_update
from rag.embedding import EMBEDDING_DIMENSION
from rag.milvus_collections import KNOWLEDGE_COLLECTIONS
from rag.milvus_delete import _escape_filter_value
from rag.preview import preview_document
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 = {}
        page_size = 4096
        for collection_name in KNOWLEDGE_COLLECTIONS:
            offset = 0
            while True:
                rows = await self.milvus_client.query(
                    collection_name=collection_name,
                    filter="",
                    output_fields=["doc_id", "title", "strategy"],
                    limit=page_size,
                    offset=offset,
                )
                if not rows:
                    break
                for row in rows:
                    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
                if len(rows) < page_size:
                    break
                offset += page_size
        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,
        *,
        filename: str,
        content: bytes,
        title: str | None = None,
        doc_id: str | None = None,
        collection_name: str,
        strategy: str,
        chunk_size: int | None = None,
        chunk_overlap: int | None = None,
    ) -> dict:
        suffix = self._validate_upload(filename, content)
        if collection_name not in KNOWLEDGE_COLLECTIONS:
            raise UploadValidationError("知识库集合不在允许范围内")

        stem = Path(filename).stem
        title = (title or "").strip() or stem
        if not doc_id or not doc_id.strip():
            doc_id = self._slugify(stem) or f"doc_{int(time.time())}_{uuid.uuid4().hex[:8]}"
        else:
            doc_id = doc_id.strip()

        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已存在，禁止重复入库")

        upload_id = uuid.uuid4().hex
        temp_path = self.storage_dir / f"{upload_id}{suffix}"
        temp_path.write_bytes(content)

        try:
            preview = preview_document(
                temp_path,
                strategy=strategy,
                chunk_size=chunk_size,
                chunk_overlap=chunk_overlap,
            )
            if not preview["chunks"]:
                raise UploadValidationError("文档切片为空")

            vectors = await self.embedder(
                [chunk["text"] for chunk in preview["chunks"]]
            )
            if len(vectors) != len(preview["chunks"]) or any(
                len(vector) != EMBEDDING_DIMENSION for vector in vectors
            ):
                raise UploadValidationError(
                    f"Embedding维度必须为{EMBEDDING_DIMENSION}"
                )

            rows = [
                {
                    "chunk_id": str(uuid5(NAMESPACE_URL, f"{doc_id}:{index}")),
                    "doc_id": doc_id,
                    "title": title,
                    "section_title": chunk.get("section_title") or "",
                    "text": chunk["text"],
                    "strategy": preview["actual_strategy"],
                    "vector": vector,
                }
                for index, (chunk, vector) in enumerate(
                    zip(preview["chunks"], vectors)
                )
            ]
            try:
                await self.milvus_client.insert(
                    collection_name=collection_name, data=rows
                )
            except Exception:
                logger.exception(
                    "atomic Milvus ingestion failed for doc_id=%s", doc_id
                )
                try:
                    await self.milvus_client.delete(
                        collection_name=collection_name,
                        filter=f'doc_id == "{_escape_filter_value(doc_id)}"',
                    )
                except Exception:
                    logger.exception(
                        "failed to clean up document rows for doc_id=%s", doc_id
                    )
                raise

            result = {
                "doc_id": doc_id,
                "title": title,
                "collection_name": collection_name,
                "filename": Path(filename).name,
                "strategy": preview["strategy"],
                "actual_strategy": preview["actual_strategy"],
                "degraded": preview["degraded"],
                "warning": preview["warning"],
                "cleaning_changed": preview["cleaning_changed"],
                "cleaning_warnings": preview["cleaning_warnings"],
                "chunk_count": len(rows),
                "chunk_size": preview["chunk_size"],
                "chunk_overlap": preview["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:
            temp_path.unlink(missing_ok=True)

    @staticmethod
    def _slugify(text: str) -> str:
        s = text.strip().lower()
        s = re.sub(r"[^\w\u4e00-\u9fff]+", "_", s)
        s = re.sub(r"_+", "_", s).strip("_")
        return s

    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"