"""
MinerU Official 解析器

使用 MinerU 官方云服务 API 进行文档解析
"""

import os
import tempfile
import time
import zipfile
from pathlib import Path
from typing import Any

import requests

from yuxi.knowledge.parser.base import BaseDocumentProcessor, DocumentParserException
from yuxi.knowledge.parser.zip_utils import process_zip_file_sync
from yuxi.utils import hashstr, logger


class MinerUOfficialParser(BaseDocumentProcessor):
    """MinerU 官方 API 解析器"""

    def __init__(self, api_key: str | None = None):
        self.api_key = api_key or os.getenv("MINERU_API_KEY")
        if not self.api_key:
            raise DocumentParserException("MINERU_API_KEY 环境变量未设置", "mineru_official", "missing_api_key")

        self.api_base = "https://mineru.net/api/v4"
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
        }

    def get_service_name(self) -> str:
        return "mineru_official"

    def get_supported_extensions(self) -> list[str]:
        """MinerU 官方 API 支持的文件格式"""
        return [".pdf", ".doc", ".docx", ".ppt", ".pptx", ".png", ".jpg", ".jpeg"]

    def check_health(self) -> dict[str, Any]:
        """检查 API 可用性和密钥有效性"""
        try:
            # 使用一个简单的测试请求来验证 API 密钥
            # 由于没有专门的 ping 接口，我们尝试创建一个测试任务的请求
            test_data = {"url": "https://cdn-mineru.openxlab.org.cn/demo/example.pdf", "is_ocr": True}

            response = requests.post(f"{self.api_base}/extract/task", headers=self.headers, json=test_data, timeout=10)

            # 如果返回 401 或特定的 API 错误码，说明密钥有问题
            if response.status_code == 401:
                return {"status": "unhealthy", "message": "API 密钥无效或已过期", "details": {"error_code": "A0202"}}
            elif response.status_code == 403:
                return {"status": "unhealthy", "message": "API 密钥权限不足", "details": {"error_code": "A0211"}}
            elif response.status_code == 200:
                # 解析响应检查是否成功创建任务
                try:
                    result = response.json()
                    if result.get("code") == 0:
                        return {
                            "status": "healthy",
                            "message": "MinerU 官方 API 服务可用",
                            "details": {"api_base": self.api_base},
                        }
                    else:
                        return {
                            "status": "unhealthy",
                            "message": f"API 返回错误: {result.get('msg', '未知错误')}",
                            "details": {"error_code": result.get("code")},
                        }
                except Exception:
                    return {
                        "status": "healthy",
                        "message": "MinerU 官方 API 服务可用",
                        "details": {"api_base": self.api_base},
                    }
            else:
                return {
                    "status": "unhealthy",
                    "message": f"API 服务异常: HTTP {response.status_code}",
                    "details": {"status_code": response.status_code},
                }

        except requests.exceptions.Timeout:
            return {"status": "timeout", "message": "API 请求超时", "details": {"timeout": "10s"}}
        except requests.exceptions.ConnectionError:
            return {
                "status": "unavailable",
                "message": "无法连接到 MinerU 官方 API 服务",
                "details": {"api_base": self.api_base},
            }
        except Exception as e:
            return {"status": "error", "message": f"健康检查失败: {str(e)}", "details": {"error": str(e)}}

    def process_file(self, file_path: str, params: dict[str, Any] | None = None) -> str:
        """
        使用 MinerU 官方 API 处理文件

        Args:
            file_path: 本地文件路径
            params: 处理参数
                - is_ocr: 是否启用 OCR (默认: True)
                - enable_formula: 是否启用公式识别 (默认: True)
                - enable_table: 是否启用表格识别 (默认: True)
                - language: 文档语言 (默认: "ch")
                - page_ranges: 页码范围 (默认: None)
                - model_version: 模型版本 "pipeline" 或 "vlm" (默认: "pipeline")

        Returns:
            str: 提取的 Markdown 文本
        """
        if not os.path.exists(file_path):
            raise DocumentParserException(f"文件不存在: {file_path}", self.get_service_name(), "file_not_found")

        file_ext = Path(file_path).suffix.lower()
        if not self.supports_file_type(file_ext):
            raise DocumentParserException(
                f"不支持的文件类型: {file_ext}", self.get_service_name(), "unsupported_file_type"
            )

        # 先检查 API 健康状态
        health = self.check_health()
        if health["status"] != "healthy":
            raise DocumentParserException(
                f"MinerU 官方 API 不可用: {health['message']}", self.get_service_name(), health["status"]
            )

        # 处理参数
        params = params or {}

        # 由于官方 API 不支持直接文件上传，我们需要先上传文件到可访问的 URL
        # 这里使用批量文件上传接口
        try:
            start_time = time.time()
            logger.info(f"MinerU Official 开始处理: {os.path.basename(file_path)}")

            # 步骤 1: 申请文件上传链接
            batch_id = self._upload_file(file_path, params)
            logger.info(f"文件上传成功，batch_id: {batch_id}")

            # 步骤 2: 轮询任务结果
            result = self._poll_batch_result(batch_id)
            logger.info(f"任务完成，状态: {result['state']}")

            zip_url = result.get("full_zip_url")

            try:
                zip_path = self._download_zip(zip_url)
            except Exception:
                text = self._download_and_extract(zip_url)
                processing_time = time.time() - start_time
                logger.info(
                    f"MinerU Official: {os.path.basename(file_path)} - {len(text)} 字符 ({processing_time:.2f}s)"
                )
                return text

            try:
                image_bucket = params.get("image_bucket") or "public"
                image_prefix = params.get("image_prefix") or "unknown/kb-images"

                processed = process_zip_file_sync(
                    zip_path,
                    image_bucket=image_bucket,
                    image_prefix=image_prefix,
                )
                text = processed["markdown_content"]
            except Exception:
                import zipfile

                text = ""
                logger.error(f"从 zip 文件中提取 full.md 失败: {zip_path}，使用第一个 md 文件")
                with zipfile.ZipFile(zip_path, "r") as zf:
                    md_files = [n for n in zf.namelist() if n.lower().endswith(".md")]
                    if md_files:
                        md_file = next((n for n in md_files if Path(n).name == "full.md"), md_files[0])
                        with zf.open(md_file) as f:
                            text = f.read().decode("utf-8")
            finally:
                try:
                    os.unlink(zip_path)
                except Exception:
                    pass

            processing_time = time.time() - start_time
            logger.info(
                f"MinerU Official 处理成功: {os.path.basename(file_path)} - {len(text)} 字符 ({processing_time:.2f}s)"
            )

            return text

        except Exception as e:
            if isinstance(e, DocumentParserException):
                raise
            processing_time = time.time() - start_time
            error_msg = f"MinerU Official 处理失败: {str(e)}"
            logger.error(f"{error_msg} ({processing_time:.2f}s)")
            raise DocumentParserException(error_msg, self.get_service_name(), "processing_failed")

    def _upload_file(self, file_path: str, params: dict[str, Any]) -> str:
        """上传文件并返回 batch_id"""
        filename = os.path.basename(file_path)

        data_id = params.get("data_id", filename)
        if len(data_id) > 30:
            data_id = data_id[:30] + "_" + hashstr(data_id, length=8)

        upload_data = {
            "enable_formula": params.get("enable_formula", True),
            "enable_table": params.get("enable_table", True),
            "language": params.get("language", "ch"),
            "files": [
                {
                    "name": filename,
                    "is_ocr": params.get("is_ocr", True),
                    "data_id": data_id,
                    "page_ranges": params.get("page_ranges"),
                }
            ],
        }

        # 申请上传链接
        response = requests.post(f"{self.api_base}/file-urls/batch", headers=self.headers, json=upload_data, timeout=30)

        if response.status_code != 200:
            raise DocumentParserException(
                f"申请上传链接失败: HTTP {response.status_code}", self.get_service_name(), "upload_url_failed"
            )

        result = response.json()
        if result.get("code") != 0:
            error_msg = result.get("msg", "未知错误")
            raise DocumentParserException(
                f"申请上传链接失败: {error_msg}", self.get_service_name(), f"api_error_{result.get('code', 'unknown')}"
            )

        batch_id = result["data"]["batch_id"]
        upload_urls = result["data"]["file_urls"]

        if not upload_urls:
            raise DocumentParserException("未获取到文件上传链接", self.get_service_name(), "no_upload_url")

        # 上传文件
        upload_url = upload_urls[0]
        with open(file_path, "rb") as f:
            upload_response = requests.put(upload_url, data=f, timeout=60)

        if upload_response.status_code != 200:
            raise DocumentParserException(
                f"文件上传失败: HTTP {upload_response.status_code}", self.get_service_name(), "file_upload_failed"
            )

        return batch_id

    def _poll_batch_result(self, batch_id: str, max_wait_time: int = 600) -> dict[str, Any]:
        """轮询批量任务结果"""
        start_time = time.time()

        while time.time() - start_time < max_wait_time:
            response = requests.get(
                f"{self.api_base}/extract-results/batch/{batch_id}", headers=self.headers, timeout=30
            )

            if response.status_code != 200:
                raise DocumentParserException(
                    f"查询任务状态失败: HTTP {response.status_code}", self.get_service_name(), "status_query_failed"
                )

            result = response.json()
            if result.get("code") != 0:
                error_msg = result.get("msg", "未知错误")
                raise DocumentParserException(
                    f"查询任务状态失败: {error_msg}",
                    self.get_service_name(),
                    f"api_error_{result.get('code', 'unknown')}",
                )

            extract_results = result["data"].get("extract_result", [])
            if not extract_results:
                time.sleep(5)
                continue

            # 检查第一个文件的状态
            file_result = extract_results[0]
            state = file_result.get("state")

            if state == "done":
                return file_result
            elif state == "failed":
                err_msg = file_result.get("err_msg", "未知错误")
                raise DocumentParserException(f"文档解析失败: {err_msg}", self.get_service_name(), "parsing_failed")

            # 继续等待
            time.sleep(5)

        raise DocumentParserException("任务处理超时", self.get_service_name(), "timeout")

    def _download_and_extract(self, zip_url: str) -> str:
        """下载并解压结果文件"""
        if not zip_url:
            raise DocumentParserException("未获取到结果下载链接", self.get_service_name(), "no_download_url")

        # 下载文件
        response = requests.get(zip_url, timeout=60)
        if response.status_code != 200:
            raise DocumentParserException(
                f"下载结果失败: HTTP {response.status_code}", self.get_service_name(), "download_failed"
            )

        # 解压到临时目录
        with tempfile.NamedTemporaryFile(delete=False, suffix=".zip") as tmp_file:
            tmp_file.write(response.content)
            tmp_file.flush()

            try:
                with tempfile.TemporaryDirectory() as tmp_dir:
                    with zipfile.ZipFile(tmp_file.name, "r") as zip_ref:
                        zip_ref.extractall(tmp_dir)

                    # 查找 markdown 文件
                    md_files = list(Path(tmp_dir).glob("*.md"))
                    if md_files:
                        with open(md_files[0], encoding="utf-8") as f:
                            return f.read()

                    # 如果没有 markdown 文件，查找 json 文件
                    json_files = list(Path(tmp_dir).glob("*.json"))
                    if json_files:
                        import json

                        with open(json_files[0], encoding="utf-8") as f:
                            data = json.load(f)
                            # 尝试提取文本内容
                            if isinstance(data, dict) and "content" in data:
                                return str(data["content"])
                            return str(data)

                    # 如果都没有，返回第一个文本文件的内容
                    text_files = list(Path(tmp_dir).glob("*"))
                    if text_files:
                        with open(text_files[0], encoding="utf-8") as f:
                            return f.read()

                    raise DocumentParserException(
                        "无法从结果中提取文本内容", self.get_service_name(), "extract_content_failed"
                    )

            finally:
                os.unlink(tmp_file.name)

    def _download_zip(self, zip_url: str) -> str:
        """下载结果ZIP到临时文件并返回路径"""
        if not zip_url:
            raise DocumentParserException("未获取到结果下载链接", self.get_service_name(), "no_download_url")
        response = requests.get(zip_url, timeout=60)
        if response.status_code != 200:
            raise DocumentParserException(
                f"下载结果失败: HTTP {response.status_code}", self.get_service_name(), "download_failed"
            )
        import tempfile

        with tempfile.NamedTemporaryFile(delete=False, suffix=".zip") as tmp_file:
            tmp_file.write(response.content)
            tmp_file.flush()
            return tmp_file.name
