import os
import asyncio
import zipfile
import shutil
import requests
import random
from typing import Dict, Any

class MinerUClient:
    """通用的高可靠 MinerU 官方 v4 API 异步客户端封装"""
    def __init__(self, api_url: str, api_token: str, max_poll_time: int = 900):
        self.api_url = api_url.rstrip("/")
        self.api_token = api_token
        self.max_poll_time = max_poll_time
        self.headers = {"Authorization": f"Bearer {api_token}"}

    async def parse_pdf_to_markdown(self, pdf_path: str, output_dir: str) -> str:
        """
        端到端解析 PDF 并在指定目录中生成结构化 Markdown 及图片包
        
        :param pdf_path: 待解析的本地 PDF 物理路径
        :param output_dir: 结果输出目录
        :return: 生成的 Markdown 绝对路径
        """
        if not os.path.exists(pdf_path):
            raise FileNotFoundError(f"未找到目标 PDF 文件: {pdf_path}")
        
        filename = os.path.basename(pdf_path)
        zip_temp_path = os.path.join(output_dir, f"{os.path.splitext(filename)[0]}_temp.zip")
        os.makedirs(output_dir, exist_ok=True)

        try:
            # 1. 申请签名上传 URL
            submit_url = f"{self.api_url}/api/v4/file-urls/batch"
            payload = {
                "files": [
                    {
                        "name": filename
                    }
                ],
                "model_version": "vlm"
            }
            
            # 使用 asyncio.to_thread 防止同步 requests 库阻塞异步事件循环
            resp_data = await asyncio.to_thread(
                lambda: requests.post(submit_url, json=payload, headers=self.headers, timeout=15).json()
            )
            
            data = resp_data.get("data", {})
            batch_id = data.get("batch_id")
            upload_url = data.get("file_urls", [None])[0]
            
            if not batch_id or not upload_url:
                raise RuntimeError(f"获取 MinerU 上传签名 URL 失败: {resp_data}")

            # 2. 上传 PDF 文件流 (注意: 签名上传 PUT 必须忽略额外的 Header)
            def do_upload():
                with open(pdf_path, "rb") as f:
                    resp = requests.put(upload_url, data=f, timeout=60)
                    resp.raise_for_status()
            await asyncio.to_thread(do_upload)

            # 3. 异步非阻塞状态轮询
            poll_url = f"{self.api_url}/api/v4/extract-results/batch/{batch_id}"
            success = False
            full_zip_url = None
            elapsed_time = 0
            attempt = 0
            base_interval = 2.0
            max_interval = 30.0

            while elapsed_time < self.max_poll_time:
                # 引入指数退避 (Exponential Backoff)，加随机抖动以避开并发踩踏 (Jitter, 0 到 1 秒随机偏置)
                sleep_time = min(max_interval, base_interval * (2 ** attempt)) + random.random()
                
                try:
                    status_resp = await asyncio.to_thread(
                        lambda: requests.get(poll_url, headers=self.headers, timeout=10).json()
                    )
                    status_data = status_resp.get("data", {})
                    extract_results = status_data.get("extract_result", [])
                    
                    if extract_results:
                        file_res = extract_results[0]
                        state = file_res.get("state")
                        
                        if state == "done":
                            full_zip_url = file_res.get("full_zip_url")
                            success = True
                            break
                        elif state in ["failed", "error"]:
                            err_msg = file_res.get("err_msg", "")
                            raise RuntimeError(f"MinerU 解析任务失败。原因: {err_msg}")
                except Exception as e:
                    if "MinerU 解析任务失败" in str(e):
                        raise e
                    # 允许网络抖动时继续轮询
                    await asyncio.sleep(sleep_time)
                    elapsed_time += sleep_time
                    attempt += 1
                    continue

                await asyncio.sleep(sleep_time)
                elapsed_time += sleep_time
                attempt += 1

            if not success or not full_zip_url:
                raise TimeoutError("轮询超时，未能成功获取 MinerU 提取包下载地址。")

            # 4. 下载提取 ZIP 包并流式写入本地
            def download_zip():
                resp = requests.get(full_zip_url, stream=True, timeout=120)
                resp.raise_for_status()
                with open(zip_temp_path, "wb") as f:
                    for chunk in resp.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
            await asyncio.to_thread(download_zip)

            # 5. 解压并重命名结构化 Markdown
            def extract_and_rename():
                # 清理输出目录中的历史杂质，为新解析内容腾空
                for item in os.listdir(output_dir):
                    item_path = os.path.join(output_dir, item)
                    if item_path != zip_temp_path:
                        if os.path.isdir(item_path):
                            shutil.rmtree(item_path)
                        else:
                            os.remove(item_path)
                            
                with zipfile.ZipFile(zip_temp_path, "r") as z:
                    z.extractall(output_dir)

            await asyncio.to_thread(extract_and_rename)

        finally:
            # 安全释放临时压缩包
            if os.path.exists(zip_temp_path):
                os.remove(zip_temp_path)

        # 定位解压出的 Markdown 文件并统一重命名为以 PDF 文件名命名的 md
        target_md_name = f"{os.path.splitext(filename)[0]}.md"
        target_md_path = os.path.join(output_dir, target_md_name)
        
        full_md_path = os.path.join(output_dir, "full.md")
        if os.path.exists(full_md_path):
            os.rename(full_md_path, target_md_path)
        else:
            # 降级方案：寻找目录下的任意 markdown
            md_files = [f for f in os.listdir(output_dir) if f.endswith(".md")]
            if md_files:
                os.rename(os.path.join(output_dir, md_files[0]), target_md_path)
            else:
                raise FileNotFoundError("解压包内未发现任何 Markdown 解析结果。")

        return target_md_path