import torch
from transformers import AutoProcessor, AutoModelForCausalLM, MarianMTModel, MarianTokenizer
from PIL import Image, ImageDraw
import requests
from io import BytesIO
import asyncio
import re

class FlorenceService:
    def __init__(self):
        self.model = None
        self.processor = None
        self.translator_model = None
        self.translator_tokenizer = None
        self.is_loaded = False
        
        # 🚀 自动检测并使用 GPU（如果可用）
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"🎯 使用设备: {self.device}")
        
        if self.device == "cuda":
            print(f"🚀 GPU 加速已启用: {torch.cuda.get_device_name(0)}")
            print(f"💾 GPU 显存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
            print(f"📊 GPU 数量: {torch.cuda.device_count()}")
        
    async def load_model(self):
        """异步加载模型"""
        if self.is_loaded:
            return
            
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, self._load_model_sync)
        
    def _load_model_sync(self):
        """同步加载模型（在线程池中执行）"""
        print("📦 加载 Florence-2 模型...")
        self.model = AutoModelForCausalLM.from_pretrained(
            "microsoft/Florence-2-base",
            trust_remote_code=True
        ).to(self.device)
        
        self.processor = AutoProcessor.from_pretrained(
            "microsoft/Florence-2-base",
            trust_remote_code=True
        )
        
        print("✅ Florence-2 模型加载完成")
        
        # 加载英中翻译模型
        print("📦 加载英中翻译模型...")
        translator_name = "Helsinki-NLP/opus-mt-en-zh"
        self.translator_tokenizer = MarianTokenizer.from_pretrained(translator_name)
        self.translator_model = MarianMTModel.from_pretrained(translator_name).to(self.device)
        
        self.is_loaded = True
        print("✅ 翻译模型加载完成")
        
    async def detect_object_at_point(
        self, 
        image_url: str, 
        click_x: int, 
        click_y: int,
        image_index: int,
        image_width: int,
        image_height: int
    ) -> dict:
        """检测点击位置的物体（使用大区域裁剪 + 详细描述）"""
        
        # 1. 下载完整图片
        print(f"📥 下载图片: {image_url}")
        response = requests.get(image_url, timeout=10)
        image = Image.open(BytesIO(response.content)).convert("RGB")
        actual_width, actual_height = image.size
        
        print(f"📊 图片信息:")
        print(f"   - 原始尺寸: {actual_width}x{actual_height}")
        print(f"   - 传入尺寸: {image_width}x{image_height}")
        print(f"   - 点击坐标: ({click_x}, {click_y})")
        
        # 2. 在图片上绘制标注点，让模型能"看到"用户的标注
        marked_image = image.copy()
        draw = ImageDraw.Draw(marked_image)
        
        # 绘制醒目的标注标记（红色圆圈 + 十字）
        marker_size = 30
        # 外圈（红色）
        draw.ellipse(
            [click_x - marker_size, click_y - marker_size, 
             click_x + marker_size, click_y + marker_size],
            outline=(255, 0, 0),
            width=5
        )
        # 十字线（红色）
        draw.line(
            [click_x - marker_size - 10, click_y, click_x + marker_size + 10, click_y],
            fill=(255, 0, 0),
            width=5
        )
        draw.line(
            [click_x, click_y - marker_size - 10, click_x, click_y + marker_size + 10],
            fill=(255, 0, 0),
            width=5
        )
        
        print(f"✏️  已在图片上绘制标注点: ({click_x}, {click_y})")
        
        # 3. 使用标记后的完整图片进行物体检测
        print("🔍 正在检测标记后的完整图片...")
        loop = asyncio.get_event_loop()
        detected_object = await loop.run_in_executor(
            None, 
            self._detect_object_at_point_sync, 
            marked_image,
            click_x,
            click_y
        )
        
        # 3. 翻译成中文（如果是英文）
        description_cn = await self._translate_to_chinese(detected_object)
        
        # 4. 格式化返回结果（无图标，纯文本）
        formatted_text = f"图{image_index}: {description_cn}"
        
        return {
            "description": description_cn,
            "descriptionEn": detected_object,
            "confidence": 0.95,  # Florence-2 不返回置信度，固定值
            "imageIndex": image_index,
            "formattedText": formatted_text
        }
        
    def _detect_object_at_point_sync(self, image: Image.Image, click_x: int, click_y: int) -> str:
        """智能裁剪策略：聚焦在点击位置的具体物品，而不是整个人物"""
        
        # 步骤1：使用物体检测找到点击位置的物体边界框
        print("🔍 步骤1: 检测所有物体...")
        bbox = self._find_object_bbox(image, click_x, click_y)
        
        if bbox:
            x1, y1, x2, y2 = bbox
            bbox_width = x2 - x1
            bbox_height = y2 - y1
            
            print(f"📦 检测到物体边界框: ({x1:.0f}, {y1:.0f}) -> ({x2:.0f}, {y2:.0f})")
            print(f"   尺寸: {bbox_width:.0f}x{bbox_height:.0f}")
            
            # 判断：如果是大物体（人物），使用点击点周围的固定区域
            # 如果是小物体，使用完整边界框
            if bbox_width > 400 or bbox_height > 600:
                # 大物体（可能是人物）：以点击点为中心裁剪固定区域
                print("   → 检测到大物体（人物），聚焦点击位置")
                crop_size = 300
                x1 = max(0, click_x - crop_size // 2)
                y1 = max(0, click_y - crop_size // 2)
                x2 = min(image.width, click_x + crop_size // 2)
                y2 = min(image.height, click_y + crop_size // 2)
                
                print(f"✂️  聚焦裁剪: ({x1}, {y1}) -> ({x2}, {y2}), 尺寸: {x2-x1}x{y2-y1}")
            else:
                # 小物体：使用完整边界框+扩展
                print("   → 检测到小物体，使用完整边界框")
                padding = 0.15
                pad_x = int(bbox_width * padding)
                pad_y = int(bbox_height * padding)
                
                x1 = max(0, int(x1 - pad_x))
                y1 = max(0, int(y1 - pad_y))
                x2 = min(image.width, int(x2 + pad_x))
                y2 = min(image.height, int(y2 + pad_y))
                
                print(f"✂️  扩展边界框: ({x1}, {y1}) -> ({x2}, {y2})")
            
            cropped_region = image.crop((x1, y1, x2, y2))
        else:
            # 如果没有检测到物体，使用标注点周围的中等区域
            print("⚠️  未检测到物体边界框，使用标注点周围区域")
            crop_size = 300
            x1 = max(0, click_x - crop_size // 2)
            y1 = max(0, click_y - crop_size // 2)
            x2 = min(image.width, click_x + crop_size // 2)
            y2 = min(image.height, click_y + crop_size // 2)
            
            print(f"✂️  裁剪区域: ({x1}, {y1}) -> ({x2}, {y2})")
            cropped_region = image.crop((x1, y1, x2, y2))
        
        # 步骤2：对裁剪区域进行详细描述
        print("🔍 步骤2: 生成详细描述...")
        description = self._describe_region(cropped_region)
        
        print(f"✅ 最终识别: {description}")
        return description
    
    def _find_object_bbox(self, image: Image.Image, click_x: int, click_y: int):
        """使用物体检测找到点击位置的物体边界框"""
        import time
        start_time = time.time()
        
        prompt = "<OD>"
        
        inputs = self.processor(
            text=prompt, 
            images=image, 
            return_tensors="pt"
        )
        
        # 显式移动所有tensor到GPU
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        print(f"   → 推理设备: {next(self.model.parameters()).device}")
        
        with torch.no_grad():
            generated_ids = self.model.generate(
                input_ids=inputs["input_ids"],
                pixel_values=inputs["pixel_values"],
                max_new_tokens=1024,
                num_beams=3,
                do_sample=False
            )
        
        elapsed = time.time() - start_time
        print(f"   → 物体检测耗时: {elapsed:.2f}秒")
        
        generated_text = self.processor.batch_decode(
            generated_ids, 
            skip_special_tokens=False
        )[0]
        
        parsed_results = self.processor.post_process_generation(
            generated_text,
            task="<OD>",
            image_size=(image.width, image.height)
        )
        
        if parsed_results and '<OD>' in parsed_results:
            od_result = parsed_results['<OD>']
            bboxes = od_result.get('bboxes', [])
            labels = od_result.get('labels', [])
            
            print(f"   检测到 {len(bboxes)} 个物体")
            
            # 找到包含点击点的边界框
            for i, (bbox, label) in enumerate(zip(bboxes, labels)):
                x1, y1, x2, y2 = bbox
                if x1 <= click_x <= x2 and y1 <= click_y <= y2:
                    print(f"   ✓ 找到匹配物体 {i+1}: {label}")
                    return bbox
            
            # 如果没有包含的，找最近的
            print("   未找到包含点击点的物体，寻找最近的...")
            min_distance = float('inf')
            closest_bbox = None
            
            for bbox in bboxes:
                x1, y1, x2, y2 = bbox
                center_x = (x1 + x2) / 2
                center_y = (y1 + y2) / 2
                distance = ((center_x - click_x) ** 2 + (center_y - click_y) ** 2) ** 0.5
                
                if distance < min_distance:
                    min_distance = distance
                    closest_bbox = bbox
            
            if closest_bbox:
                print(f"   ✓ 最近物体距离: {min_distance:.0f}px")
                return closest_bbox
        
        return None
    
    def _describe_region(self, region_image: Image.Image) -> str:
        """对区域图像进行详细描述"""
        import time
        start_time = time.time()
        
        prompt = "<MORE_DETAILED_CAPTION>"
        
        inputs = self.processor(
            text=prompt, 
            images=region_image, 
            return_tensors="pt"
        )
        
        # 显式移动所有tensor到GPU
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        with torch.no_grad():
            generated_ids = self.model.generate(
                input_ids=inputs["input_ids"],
                pixel_values=inputs["pixel_values"],
                max_new_tokens=50,
                num_beams=3,
                do_sample=False
            )
        
        elapsed = time.time() - start_time
        print(f"   → 描述生成耗时: {elapsed:.2f}秒")
        
        generated_text = self.processor.batch_decode(
            generated_ids, 
            skip_special_tokens=True
        )[0]
        
        print(f"   原始描述: {generated_text}")
        
        # 提取核心物品名称
        description = self._extract_item_name(generated_text)
        
        return description
    
    def _extract_item_name(self, description: str) -> str:
        """从详细描述中提取核心物品名称"""
        # 移除常见的前缀短语（按长度排序，先匹配长的）
        remove_phrases = [
            "the image is a promotional poster ",
            "the image is a portrait ",
            "this image shows ",
            "the image shows ",
            "the image is a ",
            "the image is ",
            "a close-up of ",
            "a close up of ",
            "close-up of ",
            "close up of ",
            "an image of ",
            "a picture of ",
            "a photo of ",
            "there is a ",
            "there is ",
            "this is a ",
            "this is ",
            "it is a ",
            "it is ",
        ]
        
        description_lower = description.lower().strip()
        
        # 持续删除前缀，直到没有匹配的
        changed = True
        while changed:
            changed = False
            for phrase in remove_phrases:
                if description_lower.startswith(phrase):
                    description_lower = description_lower[len(phrase):].strip()
                    changed = True
                    break
        
        # 取前4-5个词作为物品名称
        words = description_lower.split()
        
        if len(words) >= 4:
            item_name = " ".join(words[:4])
        elif len(words) >= 2:
            item_name = " ".join(words[:2])
        else:
            item_name = description_lower
        
        return item_name.strip()
    
    def _detect_with_context(self, image: Image.Image) -> str:
        """使用大区域上下文识别主要物体（在线程池中执行）"""
        # 使用 <MORE_DETAILED_CAPTION> 生成详细描述，包含颜色、材质等
        prompt = "<MORE_DETAILED_CAPTION>"
        
        print(f"🎯 使用 prompt: {prompt}")
        
        inputs = self.processor(
            text=prompt, 
            images=image, 
            return_tensors="pt"
        ).to(self.device)
        
        # 生成描述
        with torch.no_grad():
            generated_ids = self.model.generate(
                input_ids=inputs["input_ids"],
                pixel_values=inputs["pixel_values"],
                max_new_tokens=40,  # 增加长度以获取更多细节
                num_beams=3,
                do_sample=False
            )
        
        # 解码结果
        generated_text = self.processor.batch_decode(
            generated_ids, 
            skip_special_tokens=False
        )[0]
        
        # 提取描述（去除提示词和特殊标记）
        description = generated_text.replace(prompt, "").strip()
        description = description.replace("</s>", "").replace("<s>", "").strip()
        description = " ".join(description.split())
        
        # 🎯 提取核心信息（颜色+物体）
        description = self._extract_core_description(description)
        
        print(f"✅ 识别结果（带上下文）: {description}")
        return description
    
    def _detect_region_sync(self, image: Image.Image, bbox: tuple) -> str:
        """使用区域标注识别完整图片中的指定区域（在线程池中执行）"""
        # 使用 <REGION_TO_CAPTION> 任务，识别完整图片中的指定区域
        # bbox 格式：(x1, y1, x2, y2)，归一化到 [0, 1000]
        x1, y1, x2, y2 = bbox
        prompt = f"<REGION_TO_CAPTION><loc_{x1}><loc_{y1}><loc_{x2}><loc_{y2}>"
        
        print(f"🎯 使用 prompt: {prompt}")
        
        inputs = self.processor(
            text=prompt, 
            images=image, 
            return_tensors="pt"
        ).to(self.device)
        
        # 生成描述
        with torch.no_grad():
            generated_ids = self.model.generate(
                input_ids=inputs["input_ids"],
                pixel_values=inputs["pixel_values"],
                max_new_tokens=30,
                num_beams=3,
                do_sample=False
            )
        
        # 解码结果
        generated_text = self.processor.batch_decode(
            generated_ids, 
            skip_special_tokens=False
        )[0]
        
        # 提取描述（去除提示词和特殊标记）
        # Florence-2 返回格式: <REGION_TO_CAPTION>description</s>
        description = generated_text.replace("<REGION_TO_CAPTION>", "").strip()
        
        # 清理特殊标记和位置标记
        description = re.sub(r'<loc_\d+>', '', description)
        description = description.replace("</s>", "").replace("<s>", "").strip()
        description = " ".join(description.split())
        
        # 🎯 提取核心信息（颜色+物体）
        description = self._extract_core_description(description)
        
        print(f"✅ 区域识别结果: {description}")
        return description
    
    def _detect_sync(self, image: Image.Image) -> str:
        """同步推理（在线程池中执行）- 已弃用，使用 _detect_region_sync"""
        # 使用 <MORE_DETAILED_CAPTION> 生成更详细的描述，包含颜色和细节
        prompt = "<MORE_DETAILED_CAPTION>"
        
        inputs = self.processor(
            text=prompt, 
            images=image, 
            return_tensors="pt"
        ).to(self.device)
        
        # 生成描述
        with torch.no_grad():
            generated_ids = self.model.generate(
                input_ids=inputs["input_ids"],
                pixel_values=inputs["pixel_values"],
                max_new_tokens=30,  # 增加长度以获取更多细节
                num_beams=3,
                do_sample=False
            )
        
        # 解码结果
        generated_text = self.processor.batch_decode(
            generated_ids, 
            skip_special_tokens=False
        )[0]
        
        # 提取描述（去除提示词和特殊标记）
        description = generated_text.replace(prompt, "").strip()
        
        # 清理特殊标记
        description = description.replace("</s>", "").replace("<s>", "").strip()
        # 清理多余的空格
        description = " ".join(description.split())
        
        # 🎯 提取核心信息（颜色+物体）
        description = self._extract_core_description(description)
        
        print(f"✅ 识别结果: {description}")
        return description
    
    def _extract_core_description(self, text: str) -> str:
        """提取核心描述（只保留主要物体名称）"""
        # 移除所有介绍性和描述性短语（包括摄影术语）
        remove_phrases = [
            "a close up of ",      # 特写镜头（优先匹配）
            "a close-up of ",
            "close up of ",
            "close-up of ",
            "the image shows ",
            "the photo shows ",
            "this is a ",
            "this is an ",
            "this is ",
            "there is a ",
            "there is an ",
            "there is ",
            "there are ",
            "a picture of ",
            "an image of ",
            "it shows ",
            "shows ",
            "depicts ",
            "featuring ",
            "with ",
            " on top of ",
            " on a ",
            " on the ",
            " in a ",
            " in the ",
            " sitting on ",
            " lying on ",
            " placed on ",
            " next to ",
            " beside ",
        ]
        
        result = text.lower()
        
        # 先移除介绍性短语（按长度降序，优先匹配长短语）
        for phrase in sorted(remove_phrases, key=len, reverse=True):
            result = result.replace(phrase, " ")
        
        # 只保留第一个逗号或句号之前的内容
        for sep in [". ", ", ", " and ", " with "]:
            if sep in result:
                result = result.split(sep)[0]
                break
        
        # 移除冠词
        result = result.replace(" a ", " ").replace(" an ", " ").replace(" the ", " ")
        
        # 清理多余空格
        result = " ".join(result.split()).strip()
        
        # 限制长度（最多3-4个词）
        words = result.split()
        if len(words) > 4:
            result = " ".join(words[:4])
        
        return result
        
    async def _translate_to_chinese(self, text: str) -> str:
        """使用关键词映射进行准确翻译（替代MarianMT）"""
        if not text or not text.strip():
            return text
        
        # 如果已经是中文，直接返回
        if any('\u4e00' <= char <= '\u9fff' for char in text):
            return text
        
        try:
            # 使用关键词映射
            translated = self._translate_with_keywords(text)
            print(f"🌐 翻译: {text} -> {translated}")
            return translated
        except Exception as e:
            print(f"⚠️  翻译失败: {e}，返回原文")
            return text
    
    def _translate_with_keywords(self, text: str) -> str:
        """使用关键词映射进行翻译"""
        keyword_map = {
            # 家具
            "chair": "椅子", "sofa": "沙发", "couch": "沙发",
            "bean bag": "懒人沙发", "beanbag": "懒人沙发",
            "table": "桌子", "desk": "书桌", "bed": "床",
            
            # 织物
            "towel": "毛巾", "pillow": "枕头", "cushion": "靠垫",
            "blanket": "毯子", "throw": "抱枕",
            "rug": "地毯", "carpet": "地毯", "mat": "垫子",
            "fabric": "织物", "knit": "针织", "knitted": "针织",
            "fluffy": "毛绒", "fuzzy": "毛茸茸",
            
            # 颜色
            "green": "绿色", "red": "红色", "blue": "蓝色",
            "yellow": "黄色", "white": "白色", "black": "黑色",
            "brown": "棕色", "pink": "粉色", "beige": "米色",
            "gray": "灰色", "grey": "灰色",
            
            # 物品
            "book": "书", "books": "书籍", "magazine": "杂志",
            "stack": "堆叠", "pile": "堆", "bottle": "瓶子",
            "plant": "植物", "vase": "花瓶",
            
            # 衣物
            "sweater": "毛衣", "shirt": "衬衫",
            "jacket": "夹克", "coat": "外套",
            "dress": "连衣裙", "scarf": "围巾",
            
            # 形状
            "heart": "心形", "round": "圆形",
            "square": "方形",
            
            # 状态
            "lying": "放置", "sitting": "放置",
            "resting": "靠着", "open": "打开",
            "with": "带有", "and": "和",
        }
        
        words = text.lower().split()
        translated_words = []
        
        for word in words:
            clean_word = word.strip('.,!?;:')
            translated_words.append(keyword_map.get(clean_word, word))
        
        return " ".join(translated_words)