"""
MinerU 提取工具连接与业务管理
"""
import os
import time
import logging
import requests
import zipfile
import io
from datetime import datetime
from pathlib import Path
from urllib.parse import urlparse
from typing import Optional, List, Dict, Any, Tuple

# 导入配置与基础连接
from app.core.config import config_handler
from app.base.async_mysql_connection import get_db_connection
from app.base.minio_connection import get_minio_manager

logger = logging.getLogger("MinerU")

_mineru_manager = None

class MinerUManager:
    """MinerU 管理器"""
    
    def __init__(self):
        self.token = config_handler.get("admin_app", "MINERU_TOKEN", "")
        self.api_apply = config_handler.get("admin_app", "MINERU_API_APPLY", "https://mineru.net/api/v4/file-urls/batch")
        self.api_batch_result = config_handler.get("admin_app", "MINERU_API_BATCH_RESULT", "https://mineru.net/api/v4/extract-results/batch/{}")
        
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.token}",
        }
        
        self.supported_suffix = {".pdf", ".doc", ".docx", ".ppt", ".pptx", ".png", ".jpg", ".jpeg", ".html"}
        self.minio_manager = get_minio_manager()
        logger.info("MinerU 管理器初始化完成")

    def update_db_status(self, doc_id: str, status: Optional[int] = None, error: Optional[str] = None, 
                         md_url: Optional[str] = None, json_url: Optional[str] = None):
        """更新数据库中的转换状态和 URL"""
        conn = get_db_connection()
        if not conn:
            logger.error("数据库连接失败，无法更新状态")
            return
        try:
            with conn.cursor() as cursor:
                updates = []
                params = []
                if status is not None:
                    updates.append("conversion_status = %s")
                    params.append(status)
                if error is not None:
                    updates.append("conversion_error = %s")
                    params.append(error)
                if md_url is not None:
                    updates.append("md_url = %s")
                    params.append(md_url)
                if json_url is not None:
                    updates.append("json_url = %s")
                    params.append(json_url)
                
                if not updates:
                    return
                
                updates.append("updated_time = NOW()")
                sql = f"UPDATE t_samp_document_main SET {', '.join(updates)} WHERE id = %s"
                params.append(doc_id)
                cursor.execute(sql, params)
                
                # 同步更新子表
                if json_url is not None or md_url is not None:
                    try:
                        cursor.execute("SELECT source_type, source_id FROM t_samp_document_main WHERE id = %s", (doc_id,))
                        row = cursor.fetchone()
                        if row and row['source_type'] and row['source_id']:
                            source_type, source_id = row['source_type'], row['source_id']
                            TABLE_MAP = {
                                "basis": "t_samp_standard_base_info",
                                "work": "t_samp_construction_plan_base_info",
                                "job": "t_samp_office_regulations"
                            }
                            table_name = TABLE_MAP.get(source_type)
                            if table_name:
                                sub_updates = []
                                sub_params = []
                                if json_url is not None:
                                    sub_updates.append("json_url = %s")
                                    sub_params.append(json_url)
                                # 如果子表也有 md_url，可以在此添加
                                if sub_updates:
                                    sub_sql = f"UPDATE {table_name} SET {', '.join(sub_updates)} WHERE id = %s"
                                    sub_params.append(source_id)
                                    cursor.execute(sub_sql, sub_params)
                    except Exception as e:
                        logger.error(f"同步子表 URL 失败: {e}")
                
                conn.commit()
        except Exception as e:
            logger.error(f"更新数据库状态失败: {e}")
        finally:
            conn.close()

    def apply_upload_urls(self, files_meta: List[Dict[str, Any]], model_version: str = "vlm") -> Tuple[str, List[str]]:
        """向 MinerU 申请上传链接"""
        payload = {
            "files": files_meta,
            "model_version": model_version,
        }
        try:
            r = requests.post(self.api_apply, headers=self.headers, json=payload, timeout=60)
            r.raise_for_status()
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 401:
                logger.error("MinerU Token 已过期或无效，请在 config.ini 中更新 MINERU_TOKEN")
                raise RuntimeError("MinerU 认证失败 (401): Token 已过期或无效，请联系管理员更新配置。") from e
            raise
        
        j = r.json()
        if j.get("code") != 0:
            raise RuntimeError(f"申请上传链接失败: {j.get('msg')}")
        return j["data"]["batch_id"], j["data"]["file_urls"]

    def upload_files(self, file_data_list: List[bytes], upload_urls: List[str]):
        """上传文件到 MinerU 临时存储"""
        for data, url in zip(file_data_list, upload_urls):
            res = requests.put(url, data=data, timeout=300)
            if res.status_code != 200:
                raise RuntimeError(f"文件上传失败: {url}, status={res.status_code}")

    def poll_batch(self, doc_id: str, batch_id: str, interval_sec: int = 5, timeout_sec: int = 1800) -> List[Dict[str, Any]]:
        """轮询转换结果"""
        deadline = time.time() + timeout_sec
        while True:
            try:
                r = requests.get(self.api_batch_result.format(batch_id), headers=self.headers, timeout=60)
                r.raise_for_status()
            except requests.exceptions.HTTPError as e:
                if e.response.status_code == 401:
                    logger.error("MinerU Token 已过期或无效，请在 config.ini 中更新 MINERU_TOKEN")
                    raise RuntimeError("MinerU 认证失败 (401): Token 已过期或无效，请联系管理员更新配置。") from e
                raise

            j = r.json()
            if j.get("code") != 0:
                raise RuntimeError(f"轮询失败: {j.get('msg')}")
            results = j["data"]["extract_result"]
            states = [it.get("state") for it in results]

            if all(s in ("done", "failed") for s in states):
                return results

            if time.time() > deadline:
                raise TimeoutError(f"轮询超时: batch_id={batch_id}")
            time.sleep(interval_sec)

    def process_document(self, doc_id: str, chinese_name: str, file_url: str):
        """执行完整的文档转换流程"""
        try:
            # 1. 更新状态：开始转换
            self.update_db_status(doc_id, status=1)
            
            # 2. 下载原始文件
            logger.info(f"正在下载文件: {file_url}...")
            resp = requests.get(file_url, timeout=60)
            resp.raise_for_status()
            file_content = resp.content
            
            # 检查文件类型
            content_type = resp.headers.get("Content-Type", "").lower()
            if "text/html" in content_type:
                raise RuntimeError("不支持对网页链接进行转换，请直接查看原链接。")
            
            file_ext = Path(urlparse(file_url).path).suffix.lower()
            if not file_ext:
                file_ext = ".pdf"
                
            file_name = f"{chinese_name}{file_ext}"
            
            # 3. 提交到 MinerU
            files_meta = [{"name": file_name, "data_id": doc_id}]
            batch_id, upload_urls = self.apply_upload_urls(files_meta)
            
            self.upload_files([file_content], upload_urls)
            
            # 4. 轮询结果
            results = self.poll_batch(doc_id, batch_id)
            result = results[0]
            
            if result.get("state") == "done":
                zip_url = result.get("full_zip_url")
                if zip_url:
                    # 5. 下载并处理结果
                    zip_resp = requests.get(zip_url, timeout=300)
                    zip_resp.raise_for_status()
                    
                    with zipfile.ZipFile(io.BytesIO(zip_resp.content)) as z:
                        md_files = [f for f in z.namelist() if f.endswith(".md")]
                        json_files = [f for f in z.namelist() if f.endswith(".json")]
                        
                        md_cloud_url = None
                        json_cloud_url = None
                        
                        if md_files:
                            md_content = z.read(md_files[0])
                            md_object_name = f"{self.minio_manager.base_path}/converted/{datetime.now().strftime('%Y%m%d')}/{doc_id}.md"
                            md_cloud_url = self.minio_manager.upload_file(md_content, md_object_name, content_type="text/markdown")
                            
                        if json_files:
                            json_file = next((f for f in json_files if "content_list" in f), json_files[0])
                            json_content = z.read(json_file)
                            json_object_name = f"{self.minio_manager.base_path}/converted/{datetime.now().strftime('%Y%m%d')}/{doc_id}.json"
                            json_cloud_url = self.minio_manager.upload_file(json_content, json_object_name, content_type="application/json")
                        
                        # 6. 更新数据库
                        self.update_db_status(doc_id, status=2, md_url=md_cloud_url, json_url=json_cloud_url)
                        logger.info(f"[{doc_id}] 转换成功. MD: {md_cloud_url}, JSON: {json_cloud_url}")
                else:
                    self.update_db_status(doc_id, status=3, error="未找到 ZIP 下载链接")
            else:
                self.update_db_status(doc_id, status=3, error=result.get("err_msg", "转换失败"))
                
        except Exception as e:
            logger.exception(f"[{doc_id}] 处理文档出错: {e}")
            self.update_db_status(doc_id, status=3, error=str(e))

def get_mineru_manager() -> MinerUManager:
    """获取 MinerU 管理器单例"""
    global _mineru_manager
    if _mineru_manager is None:
        _mineru_manager = MinerUManager()
    return _mineru_manager
