import torch
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration
import easyocr
from typing import Dict, Any
import json

class EnhancedVLMService:
    """
    一个生产级多模态图像理解服务示例，集成OCR增强。
    架构：BLIP (VLM) + EasyOCR (文字提取) -> 信息融合 -> 回答
    """
    def __init__(self, vlm_model_name: str = "Salesforce/blip-image-captioning-large", device: str = None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        # 1. 初始化VLM（这里以BLIP为例，它是一个生成式VLM）
        self.processor = BlipProcessor.from_pretrained(vlm_model_name)
        self.vlm_model = BlipForConditionalGeneration.from_pretrained(vlm_model_name).to(self.device)
        
        # 2. 初始化OCR引擎（生产环境可考虑GPU版Tesseract或云服务）
        self.reader = easyocr.Reader(['en', 'ch_sim'], gpu=(self.device == 'cuda')) # 支持中英文
        
        print(f"服务初始化完成，运行在: {self.device}")

    def _extract_text_with_ocr(self, image: Image.Image) -> str:
        """OCR文本提取，返回结构化的文本信息"""
        # 执行OCR
        ocr_results = self.reader.readtext(image, detail=0) # detail=0 只返回文本
        full_text = ' '.join(ocr_results)
        
        # 生产环境可在此添加：文本块坐标、置信度、版面分析等
        ocr_metadata = {
            "detected_text_blocks": len(ocr_results),
            "full_text": full_text
        }
        return full_text, ocr_metadata

    def _analyze_with_vlm(self, image: Image.Image, question: str) -> str:
        """使用VLM进行视觉问答"""
        # 预处理
        inputs = self.processor(image, question, return_tensors="pt").to(self.device)
        # 生成回答
        out = self.vlm_model.generate(**inputs, max_new_tokens=50)
        answer = self.processor.decode(out[0], skip_special_tokens=True)
        return answer

    def process_image_query(self, image_path: str, user_question: str) -> Dict[str, Any]:
        """
        核心处理流水线
        """
        # 0. 加载图像
        raw_image = Image.open(image_path).convert('RGB')
        
        # 1. 并行或串行执行OCR和VLM基础分析（生产环境可考虑并行）
        ocr_text, ocr_meta = self._extract_text_with_ocr(raw_image)
        vlm_direct_answer = self._analyze_with_vlm(raw_image, user_question)
        
        # 2. 信息融合策略：如果问题明显关于文字，或VLM回答置信度低，则融合OCR文本
        # 这里演示一个简单规则：若OCR提取到显著文本，则将其作为上下文重新提问
        enhanced_answer = vlm_direct_answer
        if len(ocr_text) > 20: # 简单阈值判断存在大量文字
            # 构造增强的Prompt，将OCR文本作为已知信息提供给VLM
            enhanced_prompt = f"""基于以下图片和其中的文字信息，回答问题。
图片中的文字内容：{ocr_text[:500]}... // 截断以避免过长
问题：{user_question}
答案："""
            enhanced_answer = self._analyze_with_vlm(raw_image, enhanced_prompt)
        
        # 3. 组装响应
        response = {
            "user_question": user_question,
            "direct_vlm_answer": vlm_direct_answer,
            "enhanced_answer": enhanced_answer,
            "ocr_metadata": ocr_meta,
            "fusion_strategy_applied": len(ocr_text) > 20
        }
        return response

# 使用示例
if __name__ == "__main__":
    service = EnhancedVLMService()
    
    # 假设有一张包含错误弹窗的截图
    result = service.process_image_query(
        image_path="error_screenshot.png", # 替换为你的图片路径
        user_question="这个错误弹窗是什么意思？我应该如何解决？"
    )
    
    print(json.dumps(result, indent=2, ensure_ascii=False))