"""
目录层级分类模块（基于LLM API智能识别）

使用 model_setting.yaml 中的配置进行一级、二级目录分类。
"""

from __future__ import annotations

import asyncio
import json
import re
from collections import Counter
from typing import Any, Dict, List, Optional

from foundation.infrastructure.config.config import config_handler
from foundation.observability.logger.loggering import review_logger as logger
from foundation.ai.agent.generate.model_generate import generate_model_client
from foundation.observability.cachefiles.cache_manager import cache, CacheBaseDir

from ..interfaces import HierarchyClassifier as IHierarchyClassifier
from ..config.provider import default_config_provider
from ..utils.prompt_loader import PromptLoader


def _extract_json(text: str) -> Optional[Dict[str, Any]]:
    """从字符串中提取第一个有效 JSON 对象"""
    if not text or not text.strip():
        return None

    text = text.strip()

    # 首先尝试直接解析整个文本（处理干净的情况）
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # 尝试从代码块中提取
    for pattern in [r"```json\s*(\{.*?})\s*```", r"```\s*(\{.*?})\s*```", r"```json\s*(\[.*?])\s*```", r"```\s*(\[.*?])\s*```"]:
        m = re.search(pattern, text, re.DOTALL)
        if m:
            try:
                return json.loads(m.group(1))
            except json.JSONDecodeError:
                pass

    # 尝试匹配最外层的大括号（对象）
    try:
        for candidate in re.findall(r"(\{[\s\S]*?})", text):
            try:
                result = json.loads(candidate)
                if isinstance(result, dict):
                    return result
            except json.JSONDecodeError:
                continue
    except Exception:
        pass

    # 尝试匹配方括号（数组）- 如果解析为数组，包装成对象
    try:
        for candidate in re.findall(r"(\[[\s\S]*?])", text):
            try:
                result = json.loads(candidate)
                if isinstance(result, list):
                    # 如果返回的是数组，尝试找出常见的字段名
                    return {"classifications": result}
            except json.JSONDecodeError:
                continue
    except Exception:
        pass

    return None


class HierarchyClassifier(IHierarchyClassifier):
    """基于层级结构的目录分类器（一级、二级目录分类）"""

    # 功能名称，对应 model_setting.yaml 中的配置
    FUNCTION_NAME_PRIMARY = "doc_classification_primary"
    FUNCTION_NAME_SECONDARY = "doc_classification_secondary"

    def __init__(self):
        self._cfg = default_config_provider
        self._concurrency = int(config_handler.get("llm_keywords", "CONCURRENT_WORKERS", "20"))
        self.category_mapping = self._cfg.get("categories.mapping", {})
        self.prompt_loader = PromptLoader()
        self.standard_categories = self.prompt_loader.get_standard_categories()
        # 二级分类模式: single | batch
        self._secondary_mode = self._cfg.get("secondary_classification.mode", "single")
        self._batch_max_titles = int(self._cfg.get("secondary_classification.batch_max_titles", "50"))

    async def _call_llm_for_classification(
        self,
        level1_title: str,
        level2_titles: str
    ) -> Optional[Dict[str, Any]]:
        """
        调用LLM进行一级分类

        使用 function_name 从 model_setting.yaml 加载模型配置
        """
        prompt = self.prompt_loader.render(
            "toc_classification",
            level1_title=level1_title,
            level2_titles=level2_titles,
        )

        try:
            content = await generate_model_client.get_model_generate_invoke(
                trace_id="hierarchy_classifier",
                system_prompt=prompt["system"],
                user_prompt=prompt["user"],
                function_name=self.FUNCTION_NAME_PRIMARY,
            )
            result = _extract_json(content)
            return result if result is not None else {"raw_content": content}
        except Exception as e:
            logger.error(f"[一级分类] LLM 调用失败: {e}")
            return None

    async def classify_async(
        self,
        toc_items: List[Dict[str, Any]],
        target_level: int = 1
    ) -> Dict[str, Any]:
        """
        一级目录分类（异步版）

        直接使用LLM对所有一级目录进行分类。

        Args:
            toc_items: 目录项列表，每项包含 title, page, level 等字段
            target_level: 目标层级，默认为1（一级目录）

        Returns:
            {
                "items": [...],           # 分类后的目录项
                "total_count": int,       # 总数
                "target_level": int,      # 目标层级
                "category_stats": {}      # 分类统计
            }
        """
        logger.info(f"[一级分类] 开始对 {target_level} 级目录进行分类...")

        # 筛选目标层级的目录项
        level1_items = [item for item in toc_items if item["level"] == target_level]
        if not level1_items:
            logger.warning(f"[一级分类] 未找到 {target_level} 级目录项")
            return {"items": [], "total_count": 0, "target_level": target_level, "category_stats": {}}

        logger.info(f"[一级分类] 找到 {len(level1_items)} 个一级目录项，准备调用LLM分类...")

        # 构建带二级子目录的层级结构
        level1_with_children = []
        for i, level1_item in enumerate(level1_items):
            level1_idx = toc_items.index(level1_item)
            next_idx = toc_items.index(level1_items[i + 1]) if i < len(level1_items) - 1 else len(toc_items)
            children = [
                item for item in toc_items[level1_idx + 1: next_idx]
                if item["level"] == target_level + 1
            ]
            level1_with_children.append({
                "level1_item": level1_item,
                "level2_children": children
            })

        # 并发调用LLM进行分类
        semaphore = asyncio.Semaphore(self._concurrency)

        async def classify_single(entry: Dict) -> Dict[str, Any]:
            """对单个一级目录进行分类"""
            async with semaphore:
                level1_item = entry["level1_item"]
                level2_children = entry["level2_children"]
                level2_titles = "\n".join(f"- {c['title']}" for c in level2_children) or "（无二级目录）"

                llm_result = await self._call_llm_for_classification(
                    level1_title=level1_item["title"],
                    level2_titles=level2_titles
                )

                # 解析LLM结果
                if llm_result and isinstance(llm_result, dict):
                    category_cn = llm_result.get("category_cn", "")
                    category_code = llm_result.get("category_code", "")
                    confidence = llm_result.get("confidence", 0.0)

                    # 处理无效返回
                    if category_code in ("non_standard_invalid", "unknown"):
                        category_cn = category_code = ""

                    # 验证类别名称
                    if not category_cn or (
                        category_cn not in self.standard_categories and category_cn != "非标准项"
                    ):
                        category_cn = "非标准项"
                        category_code = "non_standard"
                        confidence = 0.0

                    # 映射到标准代码
                    if category_cn in self.category_mapping:
                        category_code = self.category_mapping.get(category_cn, category_code)
                    elif category_cn == "非标准项":
                        category_code = "non_standard"
                else:
                    # LLM调用失败，归为"非标准项"
                    logger.warning(f"[一级分类] LLM调用失败: {level1_item['title']}")
                    category_cn = "非标准项"
                    category_code = "non_standard"
                    confidence = 0.0

                return {
                    "title": level1_item["title"],
                    "page": level1_item["page"],
                    "level": level1_item["level"],
                    "category": category_cn,
                    "category_code": category_code,
                    "original": level1_item.get("original", ""),
                    "level2_count": len(level2_children),
                    "level2_titles": [c["title"] for c in level2_children],
                    "confidence": confidence,
                }

        # 并发执行所有分类任务
        tasks = [classify_single(entry) for entry in level1_with_children]
        classified_items = await asyncio.gather(*tasks)

        # 统计分类结果
        category_stats: Counter = Counter()
        for item in classified_items:
            category_stats[item["category"]] += 1

        logger.info(f"[一级分类] 完成，共 {len(classified_items)} 个目录项，分布: {dict(category_stats)}")

        result = {
            "items": classified_items,
            "total_count": len(classified_items),
            "target_level": target_level,
            "category_stats": dict(category_stats),
        }

        # 保存一级分类结果到缓存
        try:
            cache.save(
                data=result,
                subdir="document_temp",
                filename="一级分类结果",
                base_cache_dir=CacheBaseDir.CONSTRUCTION_REVIEW
            )
            logger.info("[一级分类] 结果已保存到缓存: temp/construction_review/document_temp/一级分类结果.json")
        except Exception as e:
            logger.warning(f"[一级分类] 保存缓存失败: {e}")

        return result

    def classify(
        self,
        toc_items: List[Dict[str, Any]],
        target_level: int = 1
    ) -> Dict[str, Any]:
        """
        一级目录分类（同步包装）

        适合无事件循环的同步场景调用。
        """
        try:
            return asyncio.run(self.classify_async(toc_items, target_level))
        except RuntimeError as exc:
            raise RuntimeError(
                "HierarchyClassifier.classify 不支持在运行中的事件循环内调用，请改用 await classify_async"
            ) from exc

    async def _call_llm_for_secondary_classification(
        self,
        first_category: str,
        first_category_code: str,
        level2_titles: List[str],
        original_title: str = ""
    ) -> Optional[Dict[str, Any]]:
        """
        调用LLM进行二级分类

        支持两种模式:
        - single: 每个标题单独请求 (高准确率，但请求多)
        - batch: 批量请求 (请求少，适合标题数量多的场景)
        """
        # 获取该一级分类的二级分类标准和映射
        secondary_standards = self.prompt_loader.get_secondary_standards(first_category)
        secondary_mapping = self.prompt_loader.get_secondary_mapping(first_category)

        if self._secondary_mode == "batch":
            # 批量模式: 一次性发送所有标题
            return await self._call_llm_for_secondary_batch(
                first_category, first_category_code, level2_titles,
                secondary_standards, secondary_mapping, original_title
            )
        else:
            # 单条模式: 每个标题单独请求 (并发)
            return await self._call_llm_for_secondary_single(
                first_category, first_category_code, level2_titles,
                secondary_standards, secondary_mapping, original_title
            )

    async def _call_llm_for_secondary_batch(
        self,
        first_category: str,
        first_category_code: str,
        level2_titles: List[str],
        secondary_standards: str,
        secondary_mapping: Dict[int, Dict[str, str]],
        original_title: str = ""
    ) -> Dict[str, Any]:
        """批量模式: 一次性分类所有二级标题"""
        # 构建标题列表文本
        titles_list = "\n".join(f"{i+1}. {title}" for i, title in enumerate(level2_titles))

        # 带重试的LLM调用
        max_retries = 3
        last_error = ""
        last_raw_content = ""

        for attempt in range(max_retries):
            # 构建提示词（重试时附加错误反馈）
            extra_hint = ""
            if attempt > 0 and last_error:
                extra_hint = f"\n\n注意：上一次返回格式有误 ({last_error})，请严格按要求的JSON格式输出。"
            if attempt > 0 and last_raw_content:
                extra_hint += f"\n上次返回内容片段: {last_raw_content[:200]}"

            prompt = self.prompt_loader.render(
                "chunk_secondary_classification_batch",
                first_category=first_category,
                level2_titles_list=titles_list,
                secondary_standards=secondary_standards,
            )

            # 如果有错误提示，追加到 user prompt
            user_prompt = prompt["user"]
            if extra_hint:
                user_prompt = user_prompt + extra_hint

            try:
                content = await generate_model_client.get_model_generate_invoke(
                    trace_id="hierarchy_classifier_secondary",
                    system_prompt=prompt["system"],
                    user_prompt=user_prompt,
                    function_name=self.FUNCTION_NAME_SECONDARY,
                )
                last_raw_content = content
                result = _extract_json(content)

                # DEBUG: 记录原始返回
                logger.debug(f"[二级分类批量] 尝试{attempt+1} LLM原始返回: {content[:500]}...")
                logger.debug(f"[二级分类批量] 尝试{attempt+1} JSON解析结果: {result}")

                if result and isinstance(result, dict) and "classifications" in result:
                    classifications = result["classifications"]

                    # 验证结果数量是否匹配
                    if len(classifications) != len(level2_titles):
                        logger.warning(
                            f"[二级分类批量] 返回结果数量({len(classifications)})与标题数量({len(level2_titles)})不匹配"
                        )
                        last_error = f"返回{classifications}项，预期{len(level2_titles)}项"
                        continue

                    # 映射结果
                    mapped_results = []
                    for i, classification in enumerate(classifications):
                        category_index = classification.get("category_index", 0)
                        title = level2_titles[i] if i < len(level2_titles) else ""

                        if category_index > 0 and category_index in secondary_mapping:
                            mapped = secondary_mapping[category_index]
                            mapped_results.append({
                                "title": title,
                                "category_index": category_index,
                                "category_code": mapped.get("code", ""),
                                "category_name": mapped.get("name", ""),
                                "raw_response": content if i == 0 else "",
                            })
                        else:
                            mapped_results.append({
                                "title": title,
                                "category_index": category_index,
                                "category_code": "non_standard",
                                "category_name": "非标准项",
                                "raw_response": content if i == 0 else "",
                            })

                    return {
                        "first_category": first_category,
                        "first_category_code": first_category_code,
                        "original_title": original_title,  # 添加原始标题
                        "level2_count": len(level2_titles),
                        "classifications": mapped_results,
                    }
                else:
                    if result is None:
                        last_error = "JSON解析失败"
                    elif not isinstance(result, dict):
                        last_error = f"返回类型错误({type(result).__name__})，应为JSON对象"
                    elif "classifications" not in result:
                        available_keys = list(result.keys()) if result else "空"
                        last_error = f"缺少classifications字段，可用字段: {available_keys}"
                    logger.warning(f"[二级分类批量] 尝试{attempt+1}/{max_retries}失败: {last_error}")

            except Exception as e:
                last_error = str(e)
                logger.error(f"[二级分类批量] LLM调用失败: 尝试{attempt+1}/{max_retries}, 错误: {e}")

        # 所有重试都失败，返回全部非标准项
        logger.error(f"[二级分类批量] 所有{max_retries}次尝试均失败，最后错误: {last_error}")
        return {
            "first_category": first_category,
            "first_category_code": first_category_code,
            "original_title": original_title,  # 添加原始标题
            "level2_count": len(level2_titles),
            "classifications": [
                {
                    "title": title,
                    "category_index": 0,
                    "category_code": "non_standard",
                    "category_name": "非标准项",
                    "error": f"批量分类失败: {last_error}",
                }
                for title in level2_titles
            ],
        }

    async def _call_llm_for_secondary_single(
        self,
        first_category: str,
        first_category_code: str,
        level2_titles: List[str],
        secondary_standards: str,
        secondary_mapping: Dict[int, Dict[str, str]],
        original_title: str = ""
    ) -> Dict[str, Any]:
        """单条模式: 每个二级标题单独请求 (并发)"""
        # 构建层级路径和内容预览（简化处理）
        hierarchy_path = f"{first_category}"
        content_preview = "\n".join(f"- {title}" for title in level2_titles)

        # 并发控制
        semaphore = asyncio.Semaphore(self._concurrency)

        async def classify_single_title(chunk_title: str) -> Dict[str, Any]:
            """对单个二级标题进行分类（带重试）"""
            prompt = self.prompt_loader.render(
                "chunk_secondary_classification",
                first_category=first_category,
                chunk_title=chunk_title,
                hierarchy_path=hierarchy_path,
                content_preview=content_preview,
                secondary_standards=secondary_standards,
            )

            # 带重试的LLM调用
            max_retries = 3
            async with semaphore:
                for attempt in range(max_retries):
                    try:
                        content = await generate_model_client.get_model_generate_invoke(
                            trace_id="hierarchy_classifier_secondary",
                            system_prompt=prompt["system"],
                            user_prompt=prompt["user"],
                            function_name=self.FUNCTION_NAME_SECONDARY,
                        )
                        result = _extract_json(content)
                        if result and isinstance(result, dict) and "category_index" in result:
                            category_index = result.get("category_index", 0)
                            # 映射编号到代码和名称
                            if category_index > 0 and category_index in secondary_mapping:
                                mapped = secondary_mapping[category_index]
                                return {
                                    "title": chunk_title,
                                    "category_index": category_index,
                                    "category_code": mapped.get("code", ""),
                                    "category_name": mapped.get("name", ""),
                                    "raw_response": content,
                                }
                            else:
                                # 编号为0或未找到映射，标记为非标准项
                                return {
                                    "title": chunk_title,
                                    "category_index": category_index,
                                    "category_code": "non_standard",
                                    "category_name": "非标准项",
                                    "raw_response": content,
                                }
                        else:
                            logger.warning(f"[二级分类] JSON解析失败或缺少category_index: {chunk_title}, 尝试: {attempt + 1}/{max_retries}")
                            if attempt == max_retries - 1:
                                # 最后一次尝试失败，使用默认值
                                return {
                                    "title": chunk_title,
                                    "category_index": 0,
                                    "category_code": "non_standard",
                                    "category_name": "非标准项",
                                    "raw_response": content,
                                    "error": "JSON解析失败",
                                }
                    except Exception as e:
                        logger.error(f"[二级分类] LLM调用失败: {chunk_title}, 错误: {e}, 尝试: {attempt + 1}/{max_retries}")
                        if attempt == max_retries - 1:
                            return {
                                "title": chunk_title,
                                "category_index": 0,
                                "category_code": "non_standard",
                                "category_name": "非标准项",
                                "error": str(e),
                            }

            # 不会到达这里，但保留以防万一
            return {
                "title": chunk_title,
                "category_index": 0,
                "category_code": "non_standard",
                "category_name": "非标准项",
                "error": "未知错误",
            }

        # 并发执行所有二级标题的分类
        tasks = [classify_single_title(title) for title in level2_titles]
        results = await asyncio.gather(*tasks)

        return {
            "first_category": first_category,
            "first_category_code": first_category_code,
            "original_title": original_title,  # 添加原始标题
            "level2_count": len(level2_titles),
            "classifications": results,
        }

    async def classify_secondary_async(
        self,
        primary_result: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        二级目录分类（异步版）

        基于一级分类结果，对每个一级分类下的二级标题进行分类。
        如果未提供primary_result，则从缓存读取。

        Args:
            primary_result: 一级分类结果，如果为None则从缓存读取

        Returns:
            {
                "items": [...],           # 二级分类结果
                "total_count": int,       # 总数
                "category_stats": {}      # 分类统计
            }
        """
        # 从缓存读取一级分类结果（如果未提供）
        if primary_result is None:
            try:
                primary_result = cache.load(
                    subdir="document_temp",
                    filename="一级分类结果",
                    base_cache_dir=CacheBaseDir.CONSTRUCTION_REVIEW
                )
                logger.info("[二级分类] 从缓存加载一级分类结果")
            except Exception as e:
                logger.error(f"[二级分类] 从缓存加载一级分类结果失败: {e}")
                return {"items": [], "total_count": 0, "category_stats": {}}

        if not primary_result or "items" not in primary_result:
            logger.warning("[二级分类] 一级分类结果为空或格式错误")
            return {"items": [], "total_count": 0, "category_stats": {}}

        primary_items = primary_result["items"]
        logger.info(f"[二级分类] 开始处理 {len(primary_items)} 个一级分类...")

        # 并发控制
        semaphore = asyncio.Semaphore(self._concurrency)

        async def classify_single_primary(item: Dict) -> Optional[Dict[str, Any]]:
            """对单个一级分类下的二级标题进行分类"""
            async with semaphore:
                first_category = item.get("category", "")
                first_category_code = item.get("category_code", "")
                level2_titles = item.get("level2_titles", [])

                if not level2_titles:
                    logger.debug(f"[二级分类] 跳过无二级标题的一级分类: {first_category}")
                    return None

                logger.info(f"[二级分类] 处理 '{first_category}' 下的 {len(level2_titles)} 个二级标题...")

                result = await self._call_llm_for_secondary_classification(
                    first_category=first_category,
                    first_category_code=first_category_code,
                    level2_titles=level2_titles,
                    original_title=item.get("title", "")  # 传递清理后的标题用于匹配chunks
                )

                return result

        # 并发执行所有一级分类的二级分类
        tasks = [classify_single_primary(item) for item in primary_items]
        secondary_results = await asyncio.gather(*tasks)

        # 过滤空结果
        secondary_results = [r for r in secondary_results if r is not None]

        # 统计分类结果
        category_stats: Counter = Counter()
        for result in secondary_results:
            for classification in result.get("classifications", []):
                code = classification.get("category_code", "non_standard")
                category_stats[code] += 1

        logger.info(f"[二级分类] 完成，共 {len(secondary_results)} 个一级分类的二级标题已分类")

        final_result = {
            "items": secondary_results,
            "total_count": sum(r.get("level2_count", 0) for r in secondary_results),
            "category_stats": dict(category_stats),
        }

        # 保存二级分类结果到缓存
        try:
            cache.save(
                data=final_result,
                subdir="document_temp",
                filename="二级分类结果",
                base_cache_dir=CacheBaseDir.CONSTRUCTION_REVIEW
            )
            logger.info("[二级分类] 结果已保存到缓存: temp/construction_review/document_temp/二级分类结果.json")
        except Exception as e:
            logger.warning(f"[二级分类] 保存缓存失败: {e}")

        return final_result

    def classify_secondary(
        self,
        primary_result: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        二级目录分类（同步包装）

        适合无事件循环的同步场景调用。
        """
        try:
            return asyncio.run(self.classify_secondary_async(primary_result))
        except RuntimeError as exc:
            raise RuntimeError(
                "HierarchyClassifier.classify_secondary 不支持在运行中的事件循环内调用，请改用 await classify_secondary_async"
            ) from exc
