"""
LangChain 1.0 文档加载器封装 + PDF图片提取
"""
import logging
from typing import List, Optional, Dict
from pathlib import Path
import fitz  # pymupdf
from langchain_community.document_loaders import (
    PyPDFLoader,
    Docx2txtLoader,
    TextLoader,
    UnstructuredMarkdownLoader
)
from langchain_core.documents import Document

logger = logging.getLogger(__name__)


class PDFImageExtractor:
    """使用pymupdf从PDF中提取内嵌图片"""

    @staticmethod
    def extract_images(file_path: str, output_dir: str) -> List[dict]:
        """提取PDF中所有图片，返回图片元数据列表"""
        Path(output_dir).mkdir(parents=True, exist_ok=True)
        images = []
        doc = fitz.open(file_path)

        for page_num in range(len(doc)):
            page = doc[page_num]
            # 获取页面文本作为图片上下文
            page_text = page.get_text("text")
            # 提取页面内嵌图片
            image_list = page.get_images(full=True)

            for img_idx, img_info in enumerate(image_list):
                xref = img_info[0]
                try:
                    base_image = doc.extract_image(xref)
                    image_bytes = base_image["image"]
                    ext = base_image["ext"]

                    # 过滤太小的图片（图标、装饰元素等）
                    if len(image_bytes) < 2048:
                        continue

                    # 图片周围文本（取该页文字前后各300字作为上下文）
                    context_start = max(0, page_text.find(
                        page_text[:len(page_text)//2]) if len(page_text) > 600
                        else 0)
                    context_text = page_text[context_start:context_start+600].strip()

                    filename = f"page{page_num+1}_img{img_idx+1}.{ext}"
                    output_path = Path(output_dir) / filename
                    output_path.write_bytes(image_bytes)

                    images.append({
                        "path": str(output_path),
                        "filename": filename,
                        "page": page_num + 1,
                        "context_text": context_text,
                        "size": len(image_bytes),
                    })
                except Exception as e:
                    logger.warning(f"提取图片失败 page={page_num+1} img={img_idx}: {e}")
                    continue

        doc.close()
        return images


class DocumentLoaderFactory:
    """文档加载器工厂"""
    
    @staticmethod
    def get_loader(file_path: str, file_type: str):
        """根据文件类型获取对应的加载器"""
        loaders = {
            ".pdf": PyPDFLoader,
            ".docx": Docx2txtLoader,
            ".txt": TextLoader,
            ".md": UnstructuredMarkdownLoader,
        }
        
        loader_class = loaders.get(file_type)
        if not loader_class:
            raise ValueError(f"Unsupported file type: {file_type}")
        
        return loader_class(file_path)
    
    @staticmethod
    def load_document(file_path: str, file_type: str, metadata: Optional[dict] = None) -> List[Document]:
        """加载文档并添加元数据"""
        loader = DocumentLoaderFactory.get_loader(file_path, file_type)
        documents = loader.load()
        
        if metadata:
            for doc in documents:
                doc.metadata.update(metadata)
        
        return documents