import asyncio
import hashlib
import os
from pathlib import Path
from uuid import uuid4

from fastapi import UploadFile
from sqlalchemy.ext.asyncio import async_sessionmaker

from kbqa.api.errors import AppError
from kbqa.config import Settings
from kbqa.documents.models import Document
from kbqa.documents.repository import DocumentRepository, StoredUpload
from kbqa.indexing.worker import IndexWorker
from kbqa.rag.vector_store import MilvusVectorStore

ALLOWED_EXTENSIONS = {
    ".pdf": ("pdf", "application/pdf"),
    ".md": ("md", "text/markdown"),
    ".markdown": ("md", "text/markdown"),
    ".txt": ("txt", "text/plain"),
}

ALLOWED_MIME_TYPES = {
    "pdf": {"application/pdf"},
    "md": {"text/markdown", "text/plain", "text/x-markdown"},
    "txt": {"text/plain"},
}


class DocumentService:
    def __init__(
        self,
        settings: Settings,
        session_factory: async_sessionmaker,
        worker: IndexWorker,
        vector_store: MilvusVectorStore,
        vector_mutation_lock: asyncio.Lock,
    ) -> None:
        self.settings = settings
        self.session_factory = session_factory
        self.worker = worker
        self.vector_store = vector_store
        self.vector_mutation_lock = vector_mutation_lock

    async def upload(self, file: UploadFile) -> Document:
        original_name = Path(file.filename or "").name
        suffix = Path(original_name).suffix.lower()
        if suffix not in ALLOWED_EXTENSIONS:
            raise AppError("UNSUPPORTED_FILE_TYPE", "仅支持 PDF、Markdown 和 TXT 文件", 415)
        file_type, expected_mime = ALLOWED_EXTENSIONS[suffix]
        if file.content_type not in ALLOWED_MIME_TYPES[file_type]:
            raise AppError("INVALID_MIME_TYPE", "文件扩展名与内容类型不匹配", 415)
        document_id = str(uuid4())
        target = self.settings.raw_data_dir / f"{document_id}{suffix}"
        temporary = target.with_suffix(f"{suffix}.part")
        maximum = self.settings.max_upload_mib * 1024 * 1024
        digest = hashlib.sha256()
        size = 0
        prefix = b""
        try:
            with temporary.open("wb") as handle:
                while chunk := await file.read(1024 * 1024):
                    size += len(chunk)
                    if size > maximum:
                        raise AppError("FILE_TOO_LARGE", "上传文件超过大小限制", 413)
                    if len(prefix) < 5:
                        prefix = (prefix + chunk)[:5]
                    digest.update(chunk)
                    handle.write(chunk)
            if size == 0:
                raise AppError("EMPTY_DOCUMENT", "上传文件为空", 415)
            self._validate_content(file_type, prefix, temporary)
            os.replace(temporary, target)
            upload = StoredUpload(
                document_id=document_id,
                filename=original_name,
                stored_path=str(target),
                file_type=file_type,
                mime_type=expected_mime,
                file_size=size,
                sha256=digest.hexdigest(),
            )
            async with self.session_factory() as session:
                document, job = await DocumentRepository(session).create_pending(upload)
            self.worker.enqueue(job.id)
            return document
        except Exception:
            temporary.unlink(missing_ok=True)
            if target.exists():
                target.unlink(missing_ok=True)
            raise
        finally:
            await file.close()

    @staticmethod
    def _validate_content(file_type: str, prefix: bytes, path: Path) -> None:
        if file_type == "pdf":
            if prefix != b"%PDF-":
                raise AppError("INVALID_PDF", "PDF 文件格式无效", 415)
            return
        try:
            path.read_text(encoding="utf-8")
        except UnicodeDecodeError as exc:
            raise AppError("TEXT_ENCODING_ERROR", "文本文件必须使用 UTF-8 编码", 415) from exc

    async def retry(self, document_id: str) -> Document:
        async with self.session_factory() as session:
            repository = DocumentRepository(session)
            job = await repository.requeue_failed(document_id)
            document = await repository.get(document_id)
        if document is None:
            raise AppError("DOCUMENT_NOT_FOUND", "文档不存在", 404)
        if job is None:
            raise AppError("DOCUMENT_NOT_RETRYABLE", "当前文档状态不能重试", 409)
        self.worker.enqueue(job.id)
        return document

    async def delete(self, document_id: str) -> None:
        async with self.session_factory() as session:
            repository = DocumentRepository(session)
            document = await repository.mark_deleting(document_id)
        if document is None:
            raise AppError("DOCUMENT_NOT_FOUND", "文档不存在", 404)
        await self._finish_delete(document)

    async def recover_deletions(self) -> None:
        async with self.session_factory() as session:
            documents = await DocumentRepository(session).deleting_documents()
        for document in documents:
            try:
                await self._finish_delete(document)
            except Exception:
                continue

    async def _finish_delete(self, document: Document) -> None:
        async with self.vector_mutation_lock:
            await asyncio.to_thread(self.vector_store.delete_document, document.id)
            await asyncio.to_thread(Path(document.stored_path).unlink, True)
            async with self.session_factory() as session:
                await DocumentRepository(session).delete_document(document.id)
