"""
文档处理器
负责文档解析、内容提取和结构化处理
集成doc_worker模块的智能处理能力

重构说明:
1. 使用 UnifiedDocumentStructure 统一数据结构
2. 使用类级别共享ChunkClassifier实例，避免重复创建LLM客户端
3. 统一PDF处理流程，消除代码重复
4. 移除splits冗余数据，统一使用chunks

注意: DOCX/DOC 文件应在上传层转换为 PDF，本模块不再直接处理 DOCX
"""

import io
import json
import os
import tempfile
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Any, Optional, List
from datetime import datetime
import asyncio
import uuid

from foundation.observability.logger.loggering import review_logger as logger
from foundation.observability.cachefiles import cache, CacheBaseDir
from .constants import CategoryCode, StatusCode, StageName

# 引入doc_worker核心组件
try:
    from .doc_worker.interfaces import DocumentSource, TOCExtractor, FullTextExtractor, TextSplitter, HierarchyClassifier
    from .doc_worker.pdf_worker.toc_extractor import PdfTOCExtractor
    from .doc_worker.pdf_worker.fulltext_extractor import PdfFullTextExtractor
    from .doc_worker.pdf_worker.hybrid_extractor import HybridFullTextExtractor
    from .doc_worker.pdf_worker.text_splitter import PdfTextSplitter
    from .doc_worker.classification.hierarchy_classifier import HierarchyClassifier
    from .doc_worker.classification.chunk_classifier import ChunkClassifier
    from .doc_worker.config.provider import default_config_provider
    from .doc_worker.models import (
        UnifiedDocumentStructure,
    )
    from .minimal_pipeline import SimpleDocumentProcessor
except ImportError:
    from core.construction_review.component.doc_worker.interfaces import DocumentSource, TOCExtractor, FullTextExtractor, TextSplitter, HierarchyClassifier
    from core.construction_review.component.doc_worker.pdf_worker.toc_extractor import PdfTOCExtractor
    from core.construction_review.component.doc_worker.pdf_worker.fulltext_extractor import PdfFullTextExtractor
    from core.construction_review.component.doc_worker.pdf_worker.hybrid_extractor import HybridFullTextExtractor
    from core.construction_review.component.doc_worker.pdf_worker.text_splitter import PdfTextSplitter
    from core.construction_review.component.doc_worker.classification.hierarchy_classifier import HierarchyClassifier
    from core.construction_review.component.doc_worker.classification.chunk_classifier import ChunkClassifier
    from core.construction_review.component.doc_worker.config.provider import default_config_provider
    from core.construction_review.component.doc_worker.models import (
        UnifiedDocumentStructure,
    )
    from core.construction_review.component.minimal_pipeline import SimpleDocumentProcessor


@dataclass
class DocumentComponents:
    """文档处理组件集合，统一封装各类型文档的处理组件"""
    toc_extractor: TOCExtractor
    classifier: HierarchyClassifier
    fulltext_extractor: FullTextExtractor
    text_splitter: TextSplitter


class DocumentProcessor:
    """
    文档处理器

    改进说明:
    1. 使用 UnifiedDocumentStructure 统一数据结构
    2. 使用类级别共享 _shared_chunk_classifier，避免重复创建LLM客户端
    3. 使用 DocumentComponents 统一管理处理组件
    4. 统一处理流程 _parse_content 消除代码重复
    """

    # 类级别共享的ChunkClassifier实例，避免重复创建LLM客户端
    _shared_chunk_classifier: Optional[ChunkClassifier] = None

    def __init__(self, progress_manager=None, callback_task_id: str = None, progress_state: dict = None, use_ocr: bool = False):
        """
        初始化文档处理器

        Args:
            progress_manager: 进度管理器
            callback_task_id: 回调任务ID
            progress_state: 进度状态字典
            use_ocr: 是否启用 OCR 模式（表格页使用 OCR 识别）
        """
        self.supported_types = ['pdf']  # DOCX/DOC 应在上传层转换为 PDF
        self.config = default_config_provider
        self.use_ocr = use_ocr
        # SSE 进度推送（由 DocumentWorkflow 注入）
        self._progress_manager = progress_manager
        self._callback_task_id = callback_task_id
        # 与心跳协程共享的状态字典，更新后心跳自动反映新阶段
        self._progress_state = progress_state

        # 选择提取器
        if use_ocr:
            logger.info("DocumentProcessor 使用 OCR 模式（表格页检测 + 识别）")
            extractor = HybridFullTextExtractor()
        else:
            logger.info("DocumentProcessor 使用本地提取模式（PyMuPDF）")
            extractor = PdfFullTextExtractor()

        # 初始化PDF文档的处理组件（简化版）
        self._components: Dict[str, DocumentComponents] = {
            'pdf': DocumentComponents(
                toc_extractor=PdfTOCExtractor(),
                classifier=HierarchyClassifier(),
                fulltext_extractor=extractor,
                text_splitter=PdfTextSplitter()
            )
        }

    @classmethod
    def _get_chunk_classifier(cls) -> ChunkClassifier:
        """获取共享的ChunkClassifier实例"""
        if cls._shared_chunk_classifier is None:
            cls._shared_chunk_classifier = ChunkClassifier()
        return cls._shared_chunk_classifier

    async def process_document(self, file_content: bytes, file_type: str) -> UnifiedDocumentStructure:
        """
        处理文档

        Args:
            file_content: 文件内容（字节流）
            file_type: 文件类型（pdf/docx）

        Returns:
            UnifiedDocumentStructure: 统一文档结构

        Raises:
            ValueError: 不支持的文件类型
            RuntimeError: 文档处理完全失败
        """
        try:
            logger.info(f"开始处理文档，类型: {file_type}")

            file_type_lower = file_type.lower()
            if file_type_lower not in self.supported_types:
                raise ValueError(f"不支持的文件类型: {file_type}，支持的类型: {self.supported_types}")

            # 统一调用解析方法
            unified_doc = await self._parse_content(file_content, file_type_lower)

            # 保存到缓存
            cache.document_temp(
                unified_doc.to_dict(),
                base_cache_dir=CacheBaseDir.CONSTRUCTION_REVIEW
            )

            return unified_doc

        except Exception as e:
            logger.error(f"文档处理失败: {str(e)}", exc_info=True)
            raise

    async def _parse_content(self, file_content: bytes, file_type: str) -> UnifiedDocumentStructure:
        """
        统一的文档解析方法

        Args:
            file_content: 文件内容
            file_type: 文件类型（pdf/docx）

        Returns:
            UnifiedDocumentStructure: 统一文档结构
        """
        components = self._components.get(file_type)
        if not components:
            raise ValueError(f"未找到 {file_type} 类型的处理组件")

        try:
            logger.info(f"开始使用最简流程处理{file_type.upper()}文档")

            async def _progress_adapter(stage: str, current: int, message: str):
                await self._push_classification_progress(
                    stage=stage, current=current, message=message
                )

            simple_processor = SimpleDocumentProcessor(use_ocr=self.use_ocr)
            unified_doc = await simple_processor.process_unified(
                file_content=file_content,
                file_name=f"document_{uuid.uuid4().hex[:8]}",
                progress_callback=_progress_adapter,
            )

            # 缓存结果
            await self._cache_unified_structure(unified_doc)
            await self._cache_tertiary_results(unified_doc, [])

            return unified_doc

        except Exception as e:
            logger.error(f"{file_type.upper()}解析失败: {str(e)}", exc_info=True)
            # 如果智能处理失败，尝试基础处理
            try:
                logger.info("尝试使用基础处理模式")
                return await self._fallback_processing(file_content, file_type)
            except Exception as fallback_error:
                logger.error(f"基础处理模式也失败: {str(fallback_error)}", exc_info=True)
                raise RuntimeError(
                    f"文档处理完全失败: {file_type.upper()}智能处理({str(e)}) + 基础处理({str(fallback_error)})"
                ) from e

    def _build_unified_structure(
        self,
        primary_result: Dict[str, Any],
        secondary_result: Optional[Dict[str, Any]],
        chunks: List[Dict[str, Any]],
        pages_content: List[Dict[str, Any]],
        toc_info: Dict[str, Any],
        document_name: str,
    ) -> UnifiedDocumentStructure:
        """
        构建统一文档结构（一二级分类）

        Args:
            primary_result: 一级分类结果
            secondary_result: 二级分类结果
            chunks: 文档切分结果
            pages_content: 页面内容
            toc_info: 目录信息
            document_name: 文档名称

        Returns:
            UnifiedDocumentStructure: 统一文档结构
        """
        # 计算总行数
        total_lines = 0
        for chunk in chunks:
            content = chunk.get("review_chunk_content", "")
            total_lines += len(content.split("\n"))

        # 创建统一结构
        unified_doc = build_unified_structure(
            primary_result=primary_result,
            secondary_result=secondary_result or {"items": []},
            chunks=chunks,
            document_name=document_name,
            total_pages=len(pages_content),
        )

        # 更新总行数
        unified_doc.total_lines = total_lines
        unified_doc.document_id = str(uuid.uuid4())

        # 构建大纲（从二级分类中提取）
        outline_items = []
        for sec in unified_doc.secondary_classifications:
            outline_items.append(OutlineItem(
                first_seq=sec.first_seq,
                first_code=sec.first_code,
                first_name=sec.first_name,
                second_seq=sec.second_seq,
                second_code=sec.second_code,
                second_name=sec.second_name,
                raw_title=sec.section_label,
                page=sec.page_start,
            ))

        unified_doc.outline = Outline(items=outline_items)

        # 保存原始元数据
        unified_doc.raw_metadata = {
            "toc_info": toc_info,
            "processing_info": {
                "chunks_count": len(chunks),
                "pages_count": len(pages_content),
            }
        }

        return unified_doc

    def _merge_tertiary_results(
        self,
        unified_doc: UnifiedDocumentStructure,
        tertiary_results: List[Dict[str, Any]],
        chunks: List[Dict[str, Any]],
    ) -> UnifiedDocumentStructure:
        """
        将三级分类结果合并到统一文档结构

        Args:
            unified_doc: 统一文档结构
            tertiary_results: 三级分类结果列表
            chunks: 原始chunks（已完成三级分类）

        Returns:
            UnifiedDocumentStructure: 更新后的结构
        """
        from collections import defaultdict

        # 按二级分类代码分组收集三级分类项
        secondary_groups: Dict[str, Dict] = defaultdict(lambda: {
            "first_code": "",
            "first_name": "",
            "second_name": "",
            "section_label": "",
            "second_content": "",
            "third_items": []
        })

        # 遍历所有chunks，按二级分类分组
        for chunk in chunks:
            first_code = chunk.get("chapter_classification", "")
            second_code = chunk.get("secondary_category_code", "")

            if not second_code or second_code == "none":
                continue

            # 获取或创建分组
            group = secondary_groups[second_code]
            group["first_code"] = first_code
            group["first_name"] = chunk.get("first_name", "")
            group["second_name"] = chunk.get("second_name", "")
            group["section_label"] = chunk.get("section_label", "")

            # 合并内容
            content = chunk.get("review_chunk_content", "") or chunk.get("content", "")
            if content:
                if group["second_content"]:
                    group["second_content"] += "\n\n" + content
                else:
                    group["second_content"] = content

            # 收集三级分类详情
            details = chunk.get("tertiary_classification_details", [])
            for idx, detail in enumerate(details, 1):
                group["third_items"].append(TertiaryItem(
                    third_seq=len(group["third_items"]) + 1,
                    third_code=detail.get("third_category_code", ""),
                    third_name=detail.get("third_category_name", ""),
                    line_start=detail.get("start_line", 0),
                    line_end=detail.get("end_line", 0),
                    content=detail.get("content", ""),
                    confidence=1.0
                ))

        # 构建tertiary_classifications列表
        tertiary_list = []
        second_seq = 0

        # 调试日志
        logger.info(f"[_merge_tertiary_results] 共有 {len(secondary_groups)} 个二级分类组")

        for second_code, group in secondary_groups.items():
            logger.info(f"[_merge_tertiary_results] 处理二级分类: {second_code}, 三级项数: {len(group['third_items'])}")
            if not group["third_items"]:
                continue

            second_seq += 1

            # 计算行数统计
            total_lines = len(group["second_content"].split("\n")) if group["second_content"] else 0
            classified_lines = sum(
                item.line_end - item.line_start + 1
                for item in group["third_items"]
            )

            # 查找对应的一级seq
            first_seq = self._get_first_seq(group["first_code"], group["first_name"])

            tertiary_list.append(TertiaryClassification(
                first_seq=first_seq,
                first_code=group["first_code"],
                first_name=group["first_name"],
                second_seq=second_seq,
                second_code=second_code,
                second_name=group["second_name"],
                third_items=group["third_items"],
                total_lines=total_lines,
                classified_lines=classified_lines
            ))

        unified_doc.tertiary_classifications = tertiary_list

        # 同时更新secondary_classifications（从chunks重新构建以确保一致性）
        self._rebuild_secondary_from_chunks(unified_doc, chunks)

        return unified_doc

    def _get_first_seq(self, first_code: str, first_name: str) -> int:
        """根据一级代码或名称获取序号"""
        order_map = {
            "basis": 1,
            "overview": 2,
            "plan": 3,
            "technology": 4,
            "safety": 5,
            "quality": 6,
            "environment": 7,
            "management": 8,
            "acceptance": 9,
            "other": 10,
        }

        if first_code in order_map:
            return order_map[first_code]

        name_map = {
            "编制依据": 1,
            "工程概况": 2,
            "施工计划": 3,
            "施工工艺技术": 4,
            "安全保证措施": 5,
            "质量保证措施": 6,
            "环境保证措施": 7,
            "施工管理及作业人员配备与分工": 8,
            "验收要求": 9,
            "其它资料": 10,
            "其他资料": 10,
        }

        return name_map.get(first_name, 99)

    def _rebuild_secondary_from_chunks(
        self,
        unified_doc: UnifiedDocumentStructure,
        chunks: List[Dict[str, Any]]
    ) -> None:
        """从chunks重新构建secondary_classifications以确保数据一致性"""
        from collections import defaultdict

        # 按二级分类分组
        groups: Dict[str, Dict] = defaultdict(lambda: {
            "first_code": "",
            "first_name": "",
            "second_name": "",
            "section_label": "",
            "content": "",
            "page_start": 0,
            "page_end": 0,
            "metadata": {}
        })

        for chunk in chunks:
            second_code = chunk.get("secondary_category_code", "")
            if not second_code or second_code == "none":
                continue

            group = groups[second_code]
            group["first_code"] = chunk.get("chapter_classification", "")
            group["first_name"] = chunk.get("first_name", "")
            group["second_name"] = chunk.get("second_name", "")
            group["section_label"] = chunk.get("section_label", "")
            group["metadata"] = chunk.get("metadata", {})

            # 合并内容
            content = chunk.get("review_chunk_content", "") or chunk.get("content", "")
            if content:
                if group["content"]:
                    group["content"] += "\n\n" + content
                else:
                    group["content"] = content

            # 更新页码
            page = chunk.get("page", 0)
            if page:
                if group["page_start"] == 0 or page < group["page_start"]:
                    group["page_start"] = page
                if page > group["page_end"]:
                    group["page_end"] = page

        # 重建secondary_classifications
        secondary_list = []
        second_seq = 0

        for second_code, group in groups.items():
            second_seq += 1
            first_seq = self._get_first_seq(group["first_code"], group["first_name"])

            # 计算行号
            lines = group["content"].split("\n") if group["content"] else []
            line_start = 0
            line_end = len(lines) - 1 if lines else 0

            secondary_list.append(SecondaryClassification(
                first_seq=first_seq,
                first_code=group["first_code"],
                first_name=group["first_name"],
                second_seq=second_seq,
                second_code=second_code,
                second_name=group["second_name"],
                second_content=group["content"],
                section_label=group["section_label"],
                page_start=group["page_start"],
                page_end=group["page_end"],
                line_start=line_start,
                line_end=line_end,
                metadata=group["metadata"]
            ))

        # 按一级分类和二级分类排序
        secondary_list.sort(key=lambda x: (x.first_seq, x.second_seq))

        # 重新编号
        for idx, sec in enumerate(secondary_list, 1):
            sec.second_seq = idx

        unified_doc.secondary_classifications = secondary_list

    async def _classify_chunks_secondary(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """对chunks进行二级分类"""
        logger.info(f"{StageName.SECONDARY_CLASSIFICATION.value}: 对内容块进行二级分类")
        await self._push_classification_progress(
            stage="文档分类",
            current=20,
            message=f"正在进行二级分类，共 {len(chunks)} 个内容块..."
        )
        try:
            chunk_classifier = self._get_chunk_classifier()
            chunks = await chunk_classifier.classify_chunks_secondary_async(chunks)
            logger.info("二级分类完成")
        except Exception as e:
            logger.warning(f"二级分类失败: {str(e)}，跳过二级分类", exc_info=True)
        return chunks

    async def _classify_chunks_tertiary(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """对chunks进行三级分类，返回处理后的chunks"""
        logger.info(f"{StageName.TERTIARY_CLASSIFICATION.value}: 对内容块进行三级分类")
        await self._push_classification_progress(
            stage="文档分类",
            current=60,
            message=f"正在进行三级分类，共 {len(chunks)} 个内容块..."
        )

        try:
            cache.save(
                data=chunks,
                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}")
        try:
            chunk_classifier = self._get_chunk_classifier()

            async def _tertiary_progress(completed: int, total: int, section_name: str, success: bool):
                """将三级分类的 section 级进度映射到 60%→90%"""
                if total > 0:
                    current = 60 + int(completed / total * 30)
                    status = "完成" if success else "失败"
                    await self._push_classification_progress(
                        stage="文档分类",
                        current=current,
                        message=f"三级分类中：{section_name} {status} [{completed}/{total}]"
                    )

            chunks = await chunk_classifier.classify_chunks_tertiary_async(
                chunks, progress_callback=_tertiary_progress
            )
            logger.info("三级分类完成")
        except Exception as e:
            logger.warning(f"三级分类失败: {str(e)}，跳过三级分类", exc_info=True)
        return chunks

    async def _cache_unified_structure(self, unified_doc: UnifiedDocumentStructure) -> None:
        """
        缓存统一文档结构（二级分类后、三级分类前）

        保存路径：temp/construction_review/document_temp/统一文档结构.json
        """
        try:
            cache_path = cache.save(
                data=unified_doc.to_dict(),
                subdir='document_temp',
                filename='统一文档结构',
                base_cache_dir=CacheBaseDir.CONSTRUCTION_REVIEW
            )

            logger.info(f"[缓存] 统一文档结构已保存: {cache_path}")
            logger.debug(f"[缓存] 包含 {unified_doc.secondary_count} 个二级分类")

        except Exception as e:
            logger.warning(f"[缓存] 保存统一文档结构失败: {e}", exc_info=True)

    async def _cache_tertiary_results(
        self,
        unified_doc: UnifiedDocumentStructure,
        chunks: List[Dict[str, Any]]
    ) -> None:
        """
        缓存三级分类结果

        保存路径：
        - temp/construction_review/document_temp/三级分类结果.json
        - temp/construction_review/document_temp/三级分类_chunks.json
        """
        try:
            # 缓存统一文档结构
            cache_path = cache.save(
                data=unified_doc.to_dict(),
                subdir='document_temp',
                filename='三级分类结果',
                base_cache_dir=CacheBaseDir.CONSTRUCTION_REVIEW
            )

            logger.info(f"[缓存] 三级分类结果已保存: {cache_path}")
            logger.info(f"[缓存] 包含 {unified_doc.secondary_count} 个二级分类, {unified_doc.tertiary_count} 个三级分类")

            # 详细统计
            for t in unified_doc.tertiary_classifications:
                logger.info(f"[缓存] 三级分类 {t.second_code}: {len(t.third_items)} 个细项")

            # 缓存chunks（简化版，只保留关键字段）
            # 如果外部未传入 chunks，从 legacy_dict 中提取
            source_chunks = chunks if chunks else unified_doc.to_legacy_dict().get("chunks", [])
            chunks_summary = []
            for chunk in source_chunks:
                summary = {
                    "chunk_id": chunk.get("chunk_id"),
                    "chapter_classification": chunk.get("chapter_classification"),
                    "secondary_category_code": chunk.get("secondary_category_code"),
                    "section_label": chunk.get("section_label"),
                    "content_length": len(chunk.get("review_chunk_content", "") or chunk.get("content", "")),
                    "tertiary_classification_details": chunk.get("tertiary_classification_details", []),
                }
                chunks_summary.append(summary)

            chunks_cache_path = cache.save(
                data={
                    "total_chunks": len(source_chunks),
                    "chunks": chunks_summary
                },
                subdir='document_temp',
                filename='三级分类_chunks',
                base_cache_dir=CacheBaseDir.CONSTRUCTION_REVIEW
            )

            logger.info(f"[缓存] 三级分类chunks已保存: {chunks_cache_path}")

        except Exception as e:
            logger.warning(f"[缓存] 保存三级分类结果失败: {e}", exc_info=True)

    async def _push_classification_progress(self, stage: str, current: int, message: str) -> None:
        """推送分类阶段进度，并同步更新心跳共享状态"""
        if self._progress_state is not None:
            self._progress_state['current'] = current
            self._progress_state['message'] = message
            self._progress_state['stage'] = stage
        if self._progress_manager and self._callback_task_id:
            try:
                await self._progress_manager.update_stage_progress(
                    callback_task_id=self._callback_task_id,
                    stage_name=stage,
                    current=current,
                    status="processing",
                    message=message,
                    event_type="processing"
                )
            except Exception as e:
                logger.warning(f"分类进度推送失败: {e}")

    async def _fallback_processing(self, file_content: bytes, file_type: str) -> UnifiedDocumentStructure:
        """
        统一的基础处理模式（当智能处理失败时使用）

        Args:
            file_content: 文件内容
            file_type: 文件类型（仅支持 pdf）

        Returns:
            UnifiedDocumentStructure: 基础处理结果
        """
        return await self._fallback_pdf_processing(file_content)

    async def _fallback_pdf_processing(self, file_content: bytes) -> UnifiedDocumentStructure:
        """PDF基础处理模式（当智能处理失败时使用）"""
        try:
            from langchain_community.document_loaders import PyPDFLoader
            from langchain_text_splitters import RecursiveCharacterTextSplitter

            logger.info("使用基础PDF处理模式")

            # PyPDFLoader需要文件路径，创建临时文件
            with tempfile.NamedTemporaryFile(delete=True, suffix='.pdf') as temp_file:
                temp_file.write(file_content)
                temp_file.flush()
                temp_file_path = temp_file.name

                loader = PyPDFLoader(temp_file_path)
                documents = loader.load()

                # 文本分块
                text_splitter = RecursiveCharacterTextSplitter(
                    chunk_size=1000,
                    chunk_overlap=200,
                    separators=["\n\n", "\n", " ", ""]
                )
                splits = text_splitter.split_documents(documents)

                # 过滤空内容切块
                valid_splits = []
                for split in splits:
                    content = split.page_content.strip()
                    if content:
                        valid_splits.append(split)

                logger.info(f"基础处理完成，有效分块数量: {len(valid_splits)}")

                # 构建基础版统一文档结构
                secondary_list = []
                for i, split in enumerate(valid_splits, 1):
                    secondary_list.append(SecondaryClassification(
                        first_seq=1,
                        first_code="unknown",
                        first_name="未分类",
                        second_seq=i,
                        second_code=f"chunk_{i}",
                        second_name=f"内容块{i}",
                        second_content=split.page_content,
                        page_start=split.metadata.get("page", 0),
                        page_end=split.metadata.get("page", 0),
                    ))

                unified_doc = UnifiedDocumentStructure(
                    document_id=str(uuid.uuid4()),
                    document_name="基础处理文档.pdf",
                    total_pages=len(documents),
                    secondary_classifications=secondary_list,
                )

                return unified_doc

        except Exception as e:
            logger.error(f"基础PDF处理失败: {str(e)}", exc_info=True)
            raise
