"""
文档处理服务（LangChain 1.0 + 多模态图片处理）
"""
import os
import asyncio
from pathlib import Path
from typing import List, Dict, Any, Optional
from sqlalchemy.orm import Session
from langchain_core.documents import Document as LangChainDocument

from ..models.document import Document, DocumentChunk
from ..rag.vector_store import get_vector_store
from ..rag.document_loaders import DocumentLoaderFactory, PDFImageExtractor
from ..rag.text_splitters import get_text_splitter
from ..llm.siliconflow import get_llm_client

IMAGES_DIR = Path(__file__).parent.parent.parent.parent / "data" / "images"


class DocumentService:
    """文档处理服务"""

    def __init__(self, db: Session):
        self.db = db
        self.vector_store = get_vector_store()

    async def process_document(self, document_id: int) -> bool:
        """处理文档（使用LangChain 1.0 + 多模态图片处理）"""
        try:
            document = self.db.query(Document).filter(Document.id == document_id).first()
            if not document:
                return False

            # 1. 使用LangChain加载文档（文本）
            documents = DocumentLoaderFactory.load_document(
                file_path=document.file_path,
                file_type=document.file_type,
                metadata={
                    "document_id": document.id,
                    "knowledge_base_id": document.knowledge_base_id,
                    "title": document.title,
                    "filename": document.filename
                }
            )

            # 2. 使用中文优化的文本分割器
            text_splitter = get_text_splitter(chunk_size=1000, chunk_overlap=200)
            splits = text_splitter.split_documents(documents)

            # 3. PDF图片提取和描述（仅PDF文件）
            image_chunks = []
            if document.file_type == ".pdf":
                image_chunks = await self._process_pdf_images(
                    file_path=document.file_path,
                    document_id=document.id,
                    knowledge_base_id=document.knowledge_base_id,
                    title=document.title,
                    filename=document.filename
                )

            # 4. 将文本块和图片描述合并添加到向量存储
            all_splits = splits + image_chunks
            success = self.vector_store.add_documents(all_splits)

            if success:
                document.is_processed = True
                self.db.commit()
                print(f"[DocumentService] 文档 {document.filename} 处理完成: "
                      f"{len(splits)} 个文本块, {len(image_chunks)} 个图片描述块")
                return True

            return False

        except Exception as e:
            print(f"处理文档失败: {str(e)}")
            self.db.rollback()
            return False

    async def _process_pdf_images(
        self,
        file_path: str,
        document_id: int,
        knowledge_base_id: int,
        title: str,
        filename: str
    ) -> List[LangChainDocument]:
        """提取PDF图片并用VLM生成描述"""
        image_chunks = []
        try:
            # 创建图片输出目录
            img_output_dir = IMAGES_DIR / str(document_id)

            # 提取图片
            images = PDFImageExtractor.extract_images(str(file_path), str(img_output_dir))

            if not images:
                print(f"[DocumentService] 未发现可提取的图片: {filename}")
                return []

            print(f"[DocumentService] 提取到 {len(images)} 张图片, 开始VLM描述生成...")

            # 批量调用VLM生成描述
            llm_client = get_llm_client()
            for idx, img in enumerate(images):
                try:
                    description = await llm_client.describe_image(
                        img["path"],
                        img.get("context_text", "")
                    )

                    if description:
                        # 相对路径用于URL访问
                        rel_path = f"{document_id}/{img['filename']}"
                        image_url = f"/images/{rel_path}"

                        # 构建图片描述文本块（URL写入内容，LLM可直接引用）
                        chunk_content = (
                            f"[图片描述 - 第{img['page']}页]\n"
                            f"图片URL: {image_url}\n"
                            f"图片内容：{description}"
                        )

                        chunk = LangChainDocument(
                            page_content=chunk_content,
                            metadata={
                                "document_id": document_id,
                                "knowledge_base_id": knowledge_base_id,
                                "title": title,
                                "filename": filename,
                                "source_type": "image",
                                "image_path": str(img["path"]),
                                "image_url": image_url,
                                "page": img["page"],
                            }
                        )
                        image_chunks.append(chunk)
                        print(f"[DocumentService] 图片描述成功 {idx+1}/{len(images)}: {img['filename']}")
                    else:
                        print(f"[DocumentService] 图片描述为空 {idx+1}/{len(images)}: {img['filename']}")

                except Exception as e:
                    print(f"[DocumentService] 图片处理失败 {img['filename']}: {e}")
                    continue

        except Exception as e:
            print(f"[DocumentService] PDF图片处理失败: {e}")

        return image_chunks
    
    def search_documents(self, query: str, knowledge_base_ids: Optional[List[int]] = None, limit: int = 5) -> List[Dict[str, Any]]:
        """搜索文档（保留原有接口兼容性）"""
        try:
            # 构建过滤条件
            filter_dict = None
            if knowledge_base_ids:
                filter_dict = {"knowledge_base_id": {"$in": knowledge_base_ids}}
            
            # 使用LangChain Chroma进行搜索
            results = self.vector_store.similarity_search_with_score(
                query=query,
                k=limit,
                filter=filter_dict
            )
            
            # 格式化结果
            search_results = []
            for doc, distance in results:
                metadata = doc.metadata if hasattr(doc, 'metadata') else {}
                score = self._convert_distance_to_score(distance)
                
                search_results.append({
                    "content": doc.page_content,
                    "metadata": metadata,
                    "score": score,
                    "distance": distance
                })
            
            return search_results
            
        except Exception as e:
            print(f"搜索文档失败: {str(e)}")
            return []
    
    def _convert_distance_to_score(self, distance: float) -> float:
        """将ChromaDB的distance转换为0-1范围的相似度分数"""
        import math
        
        # 内积距离（负值）
        if distance < 0:
            return (1 + distance) / 2
        
        # 大距离使用对数缩放
        if distance > 100:
            return 1 / (1 + math.log(distance))
        
        # 标准距离转换
        return 1 / (1 + distance)
    
    def get_document_chunks(self, document_id: int) -> List[DocumentChunk]:
        """获取文档的所有块"""
        return self.db.query(DocumentChunk).filter(
            DocumentChunk.document_id == document_id
        ).order_by(DocumentChunk.chunk_index).all()
    
    def delete_document_chunks(self, document_id: int) -> bool:
        """删除文档的所有块（数据库 + 向量存储）"""
        try:
            # 删除向量存储中的文档数据
            self.vector_store.delete_by_document_id(document_id)
            # 删除数据库中的chunk记录
            self.db.query(DocumentChunk).filter(
                DocumentChunk.document_id == document_id
            ).delete()
            self.db.commit()
            return True
        except Exception as e:
            print(f"删除文档块失败: {str(e)}")
            self.db.rollback()
            return False