# -*- coding: utf-8 -*-
"""
图像对齐处理器 - 基于 imageAlignment.py
将PDF图像与提取的概念进行智能对齐
支持多线程并发处理，提升对齐速度
"""

import requests
import json
import os
import base64
import time
import re
from typing import List, Dict
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading


class ImageAlignmentProcessor:
    """
    图像对齐处理器：使用GPT-4o分析图像，匹配提取的概念
    """
    
    def __init__(self, subject="biology"):
        """
        初始化图像对齐处理器
        
        Args:
            subject: 学科领域（默认biology）
        """
        self.subject = subject
        
        # API配置（参考 imageAlignment.py）
        self.api_url = "https://jeniya.cn/v1/chat/completions"
        self.api_key = "sk-5w2FNpdHmV4xoi07RAe2lWiqXGgvRWCbmCk9o8T39hZBhJIc"
        self.model = "gpt-4o-2024-11-20"
        
        # 多线程配置
        self.max_workers = 3  # 图像分析较慢，使用3个线程避免API限流
        self.mapping_lock = threading.Lock()  # 保护共享的映射字典
        
        print(f"✅ 图像对齐处理器初始化完成（多线程并发）")
        print(f"   学科: {subject}")
        print(f"   模型: {self.model}")
        print(f"   🚀 最大并发线程数: {self.max_workers}")
    
    def process(self, image_paths: List[str], concepts: List[str], process_id: str) -> Dict[str, List[str]]:
        """
        处理图像对齐任务（支持多线程并发）
        
        Args:
            image_paths: 图像文件路径列表
            concepts: 提取的概念列表
            process_id: 处理ID
            
        Returns:
            dict: 概念-图像映射 {"概念1": ["image1.png"], "概念2": ["image2.png", "image3.png"]}
        """
        print(f"\n{'='*80}")
        print(f"🖼️  开始图像对齐（多线程并发）")
        print(f"{'='*80}")
        print(f"   图像数量: {len(image_paths)}")
        print(f"   概念数量: {len(concepts)}")
        print(f"   使用模型: {self.model}")
        print(f"   🚀 使用 {self.max_workers} 个线程并发处理")
        
        # 概念-图像映射（共享资源）
        concept_image_mapping = {}
        
        # 准备任务列表
        tasks = []
        for idx, image_item in enumerate(image_paths, 1):
            # 处理字典格式：提取local_path用于文件读取，path用于前端显示
            if isinstance(image_item, dict):
                local_path = image_item.get('local_path', '')  # 文件系统路径
                url_path = image_item.get('path', '')  # URL路径
                filename = image_item.get('filename', '')
            else:
                # 如果是字符串，直接使用
                local_path = image_item
                url_path = image_item
                filename = os.path.basename(image_item)
            
            if not local_path or not url_path:
                print(f"\n   ⚠️  图像 {idx}: 无效路径")
                continue
            
            # 检查文件是否存在（使用local_path）
            if not os.path.exists(local_path):
                print(f"\n   ⚠️  图像 {idx}: 文件不存在 - {local_path}")
                continue
            
            tasks.append((idx, local_path, url_path, filename, concepts))
        
        print(f"   📋 准备 {len(tasks)} 个图像分析任务")
        
        # 使用线程池并发处理
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            future_to_task = {
                executor.submit(self._process_single_image, task): task
                for task in tasks
            }
            
            completed = 0
            for future in as_completed(future_to_task):
                completed += 1
                task = future_to_task[future]
                idx, local_path, url_path, filename = task[0], task[1], task[2], task[3]
                
                try:
                    result = future.result()
                    if result:
                        matched_concepts = result
                        print(f"   ✅ [{completed}/{len(tasks)}] {filename}: 匹配 {len(matched_concepts)} 个概念")
                        
                        # 使用锁保护共享映射（线程安全）
                        with self.mapping_lock:
                            for concept in matched_concepts:
                                if concept not in concept_image_mapping:
                                    concept_image_mapping[concept] = []
                                concept_image_mapping[concept].append(url_path)
                    else:
                        print(f"   ℹ️  [{completed}/{len(tasks)}] {filename}: 无匹配概念")
                        
                except Exception as e:
                    print(f"   ❌ [{completed}/{len(tasks)}] {filename}: 分析失败 - {str(e)}")
        
        # 保存对齐结果
        self._save_alignment_result(concept_image_mapping, process_id)
        
        # 统计
        total_aligned = sum(len(images) for images in concept_image_mapping.values())
        print(f"\n   ✅ 图像对齐完成:")
        print(f"      - 对齐的概念: {len(concept_image_mapping)}/{len(concepts)}")
        print(f"      - 对齐的图像: {total_aligned}/{len(image_paths)}")
        
        return concept_image_mapping
    
    def _process_single_image(self, task):
        """
        处理单张图像的对齐（线程安全）
        
        Args:
            task: (idx, local_path, url_path, filename, concepts)
            
        Returns:
            list: 匹配的概念列表 或 None
        """
        idx, local_path, url_path, filename, concepts = task
        
        try:
            # 分析图像，获取匹配的概念
            matched_concepts = self._analyze_image_with_consensus(local_path, concepts)
            return matched_concepts if matched_concepts else None
            
        except Exception as e:
            return None
    
    def _analyze_image_with_consensus(self, image_path: str, concepts: List[str]) -> List[str]:
        """
        使用三次调用的共识机制分析图像（参考 imageAlignment.py 的三次调用）
        
        Args:
            image_path: 图像文件路径
            concepts: 候选概念列表
            
        Returns:
            list: 匹配的概念列表
        """
        results = []
        
        # 三次调用（参考 imageAlignment.py line 129-133）
        for attempt in range(3):
            try:
                matched = self._analyze_image_single(image_path, concepts)
                results.append(set(matched))
                time.sleep(2)  # 避免API限流
            except Exception as e:
                print(f"         ⚠️  第 {attempt + 1} 次调用失败: {str(e)}")
                results.append(set())
        
        # 取并集（参考 imageAlignment.py line 137）
        consensus = set()
        if results:
            consensus = results[0] | results[1] | results[2]
        
        return list(consensus)
    
    def _analyze_image_single(self, image_path: str, concepts: List[str]) -> List[str]:
        """
        单次图像分析（完全参考 imageAlignment.py 的 imageAnalysis 函数）
        
        Args:
            image_path: 图像文件路径
            concepts: 候选概念列表
            
        Returns:
            list: 匹配的概念列表
        """
        # 构建Prompt（参考 imageAlignment.py line 31-58）
        prompt = f"""
You are an AI {self.subject} image analyst. Follow this strict protocol:

1. Candidate Verification:
   Within the {self.subject} concepts {concepts},
   identify and extract the concepts that appear in the image or are highly related to it. 
   Use these extracted concepts to form the C4.

2. Multilevel Analysis:
   a) OCR Extraction: Capture ALL {self.subject} pure concept → C1
   b) Regional Analysis: Identify specific {self.subject} components → C2
   c) Holistic Analysis: Derive tissue/system-level {self.subject} concepts → C3

3. Deduplication Protocol:
   Priority: C4 > C3 > C2 > C1
   - Remove duplicates ONLY for identical terms
   - Preserve hierarchical terms

4. Output Decision:
   - If any valid terms remain: List ALL in bullet format
   - If completely empty: Output "none"

Mandatory Format:
- concept1
- concept2
...

OR
"none"

Absolute constraints:
1. No term suppression except exact duplicates
2. No explanations, categories, or formatting variations.
        """
        
        # 读取图像并转为Base64（参考 imageAlignment.py line 59-60）
        with open(image_path, "rb") as image_file:
            base64_image = base64.b64encode(image_file.read()).decode("utf-8")
        
        # 构建请求（参考 imageAlignment.py line 61-73）
        payload = json.dumps({
            "model": self.model,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_image}"}},
                    ],
                }
            ]
        })
        
        headers = {
            'Accept': 'application/json',
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }
        
        # 发起请求（参考 imageAlignment.py line 79-88）
        max_retries = 3
        for retry in range(max_retries):
            try:
                response = requests.post(self.api_url, headers=headers, data=payload, timeout=60)
                response_json = response.json()
                
                # 提取结果
                content = response_json["choices"][0]["message"]["content"]
                
                # 如果返回"none"，说明没有匹配的概念
                if "none" in content.lower():
                    return []
                
                # 提取概念列表（参考 imageAlignment.py line 134）
                matched = re.findall(r"-\s*(\S+)", content)
                
                # 过滤：只保留在候选概念列表中的概念
                matched = [c for c in matched if c in concepts]
                
                return matched
                
            except Exception as e:
                print(f"         ⚠️  API调用失败 (重试 {retry + 1}/{max_retries}): {str(e)}")
                if retry < max_retries - 1:
                    time.sleep(65)  # 参考 imageAlignment.py line 87
                    continue
                else:
                    return []
        
        return []
    
    def _save_alignment_result(self, mapping: Dict[str, List[str]], process_id: str):
        """
        保存对齐结果到JSON文件
        
        Args:
            mapping: 概念-图像映射
            process_id: 处理ID
        """
        output_dir = os.path.join('processed', 'alignment', process_id)
        os.makedirs(output_dir, exist_ok=True)
        
        output_file = os.path.join(output_dir, f'{process_id}_concept_image_alignment.json')
        
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(mapping, f, ensure_ascii=False, indent=2)
            
            print(f"\n   💾 对齐结果已保存: {output_file}")
        except Exception as e:
            print(f"\n   ⚠️  保存对齐结果失败: {str(e)}")

