import re
import uuid

import modal

from ..video import downloader_image, app, config

with downloader_image.imports():
    import os, httpx, json, time, requests
    from typing import List
    from loguru import logger
    from modal import current_function_call_id
    from ffmpy import FFmpeg, FFprobe
    from fastapi import HTTPException
    from httpx import Timeout
    from starlette import status
    import asyncio
    import random
    import subprocess
    from BowongModalFunctions.utils.SentryUtils import SentryUtils
    from BowongModalFunctions.models.media_model import MediaSource
    from BowongModalFunctions.models.web_model import SentryTransactionInfo, WebhookNotify
    from BowongModalFunctions.models.ffmpeg_worker_model import FFMpegSliceSegment
    from BowongModalFunctions.utils.TimeUtils import TimeDelta, merge_product_data
    from BowongModalFunctions.utils.VideoUtils import FFMPEGSliceOptions, VideoUtils


    @app.function(cpu=(0.5, 64), timeout=1800,
                  max_containers=config.video_downloader_concurrency,
                  secrets=[modal.Secret.from_name("gemini-prompt")],
                  volumes={
                      config.S3_mount_dir: modal.CloudBucketMount(
                          bucket_name=config.S3_bucket_name,
                          secret=modal.Secret.from_name("aws-s3-secret", environment_name=config.modal_environment),
                      ),
                  },
                  region="us",
                  cloud='gcp'
                  )
    @modal.concurrent(max_inputs=1)
    async def video_hls_slice_inference(media: MediaSource,
                                        google_api_key: str,
                                        product_grid_list: List[str],
                                        product_list: List[any],
                                        start_time: str,
                                        end_time: str,
                                        options: FFMPEGSliceOptions,
                                        sentry_trace: SentryTransactionInfo,
                                        webhook: WebhookNotify = None,
                                        retry_time: int = 3,
                                        scale:float = 0.9
                                        ):
        logger.info(
            f"region {os.environ.get('MODAL_REGION', 'unknown')}, provider {os.environ.get('MODAL_CLOUD_PROVIDER', 'unknown')}")

        # Gemini可用区
        gemini_region = [
            #"us-central1",
            "us-east1",
            "us-east5",
            "us-west1",
            "us-south1",
            "europe-central2",
            "europe-north1",
            "europe-west1",
            "europe-west4",
            "europe-west8",
        ]

        logger.info(f"\nmedia || {media.urn},"
                    f"\nstart_time || {start_time}, "
                    f"\nend_time || {end_time}, "
                    f"\nproduct_grid_list || {json.dumps(product_grid_list, ensure_ascii=False, indent=2)}, "
                    f"\nproduct_list || {json.dumps(product_list, ensure_ascii=False, indent=2)}"
                    f"\noptions || {options.model_dump_json()}, "
                    f"\nscale || {scale}")

        # 动态Prompt
        if "IMAGE_PRODUCT_IDENTIFICATION_PROMPT" in os.environ:
            IMAGE_PRODUCT_IDENTIFICATION_PROMPT = os.environ["IMAGE_PRODUCT_IDENTIFICATION_PROMPT"]
        else:
            IMAGE_PRODUCT_IDENTIFICATION_PROMPT = """<prompt>
            <instruction>
            你是专业的商品识别专家。我上传了商品图片网格，需要你识别图片中的商品并与商品列表进行匹配。
    
            **输入材料**：
            - 🖼️ **商品图片网格**：包含多个黑色边框区域，每个区域内有商品图片+商品名称文字
            - 📋 **商品列表**：标准商品名称参考清单
    
            **核心任务**：
            1. **扫描黑色边框区域**：从左上角开始，按行扫描每个黑色边框区域
            2. **提取文字信息**：精确提取每个区域内的所有文字信息
            3. **与商品列表匹配**：将图片文字与商品列表进行高相似度匹配
            4. **提取商品图片特征**：从商品图片提取详细可识别特征，包括颜色、图案、纹理、材质、版型、款式等
    
            **严格约束**：
            - 🚫 只识别有黑色边框包围的商品区域
            - 🚫 每个商品必须有清晰可见的文字标注
            - 🚫 不得推测或添加图片中不存在的商品
            - ✅ 输出商品数量不得超过图片中的黑色边框区域数量
    
            **商品列表**：
            {PRODUCT_LIST}
            </instruction>
    
            <output_format>
            请按以下JSON格式输出：
            [
              {{
                "image_order": 1,
                "image_name": "图片上显示的原始文字",
                "matched_product": "匹配到的标准商品名称或null",
                "match_confidence": 95,
                "visual_features": {{
                  "color": "商品详细颜色配色等有辨识度的色彩特征",
                  "pattern": "商品详细图案纹理等有辨识度的材质特征",
                  "style": "商品详细款式版型等有辨识度的风格特征"
                }}
              }}
            ]
            </output_format>
            </prompt>"""

        if "VIDEO_TIMELINE_ANALYSIS_PROMPT" in os.environ:
            VIDEO_TIMELINE_ANALYSIS_PROMPT = os.environ["VIDEO_TIMELINE_ANALYSIS_PROMPT"]
        else:
            VIDEO_TIMELINE_ANALYSIS_PROMPT = """<prompt>
            <instruction>
            基于已识别的商品清单，分析视频中每个商品的出现时间段。
    
            **输入材料**：
            - 📹 **视频**：直播带货片段
            - 📋 **已识别商品清单**：第一阶段确认的商品列表
    
            **分析任务**：
            1. **时间段识别**：找出每个商品在视频中的展示时间
            2. **内容类型分类**：判断每个时间段的具体内容类型
            3. **时间段合并**：合并连续或相近的时间段
    
            **内容类型**：
            - (穿着本品+介绍本品) - 主播穿着该商品且正在介绍该商品本身
            - (穿着本品+他品) - 主播穿着该商品但正在介绍其他商品  
            - (穿着本品+无关) - 主播穿着该商品但在做无关商品的事情
    
            **时间格式标准**：
            - 必须严格使用格式：HH:MM:SS.mmm - HH:MM:SS.mmm
            - HH是小时(00-23)，MM是分钟(00-59)，SS是秒(00-59)，mmm是毫秒(000-999)
            - 示例：00:05:23.500 表示0小时5分钟23.5秒，不是5小时23分钟
            - 常见错误：不要将05:23:500解释为5小时23分钟，应该是0小时5分钟23.5秒
    
            **已识别商品清单(product-商品名 feature-商品特征 feature.color-商品详细颜色配色等色彩特征 feature.pattern-商品详细图案纹理等材质特征 feature.style-商品详细款式版型等风格特征)**：
            {IDENTIFIED_PRODUCTS}
            </instruction>
    
            <output_format>
            请按以下JSON格式输出：
            [
              {{
                "product": "商品名称",
                "timeline": [
                  "00:01:15.200 - 00:02:30.800 (穿着本品+介绍本品)",
                  "00:05:10.100 - 00:06:25.600 (穿着本品+他品)"
                ]
              }}
            ]
            </output_format>
            </prompt>"""


        def upload(file_path):
            with open(file_path, "rb") as video:
                video = video.read()
            content_length = len(video)
            content_type = f"video/{file_path.split('.')[-1]}"
            filename = ".".join([str(uuid.uuid4()), file_path.split(".")[-1]])
            if content_type not in ['video/mp4', 'video/mpeg', 'video/mov', 'video/avi', 'video/x-flv', 'video/mpg',
                                    'video/webm', 'video/wmv', 'video/3gpp']:
                raise HTTPException(status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE)
            logger.info(
                f"Uploading name = {filename}, size = {content_length}, type = {content_type} to google file")
            with httpx.Client(timeout=1800) as client:
                upload_response = client.post(
                    url=f"https://storage.googleapis.com/upload/storage/v1/b/dy-media-storage/o?uploadType=media&name=videos%2F{filename}",
                    content=video,
                    headers={
                        "Authorization": f"Bearer {google_api_key}",
                        "Content-Type": content_type
                    })
                upload_response.raise_for_status()

                upload_url = f"gs://dy-media-storage/videos/{filename}"

            return upload_url, upload_response.status_code

        def parse_stage1_result(result_text):
            """解析第一阶段结果，提取已识别的商品列表"""
            try:
                # 清理结果文本，提取JSON部分
                clean_result = result_text.strip()
                if clean_result.startswith('```') and clean_result.endswith('```'):
                    clean_result = clean_result[3:-3].strip()
                    if clean_result.startswith('json'):
                        clean_result = clean_result[4:].strip()

                # 解析JSON
                stage1_data = json.loads(clean_result)

                identified_products = []
                if isinstance(stage1_data, list):
                    for item in stage1_data:
                        if isinstance(item, dict) and 'matched_product' in item:
                            matched_product = item['matched_product']
                            if matched_product and matched_product != "null":
                                identified_products.append({"product": matched_product, "feature": item["visual_features"]})

                logger.info(f"一阶段解析结果: {len(identified_products)} 个有效商品")
                return identified_products

            except Exception as e:
                logger.exception(f"一阶段结果解析失败 {str(e)}")
                raise e


        @SentryUtils.webhook_handler(webhook, current_function_call_id())
        @SentryUtils.sentry_tracker(sentry_trace.x_trace_id, sentry_trace.x_baggage, op="inference_gemini",
                                    name="Gemini推理", fn_id=current_function_call_id())
        async def _handler(media: MediaSource,
                           google_api_key: str,
                           product_grid_list: List[str],
                           product_list: List[any],
                           start_time: str,
                           end_time: str,
                           options: FFMPEGSliceOptions,
                           sentry_trace: SentryTransactionInfo,
                           retry_time: int = 3,
                           scale:float = 0.9):

            video_gemini_uri = None

            try:
                # 1、首先获取全量商品列表
                logger.info("1、获取直播间商品列表")
                if len(product_list) == 0:
                    logger.error("商品列表为空, 退出推理")
                    raise Exception("商品列表为空, 退出推理")
                product_title_list = []
                if isinstance(product_list[0], dict):
                    for product in product_list:
                        product_title_list.append(product["title"])
                else:
                    product_title_list = product_list

                # 2、切20分钟的条
                logger.info("2、开始截取指定视频")
                # 计算缩放尺寸
                metadata = VideoUtils.ffprobe_media_metadata(media.path)
                width = int(metadata.streams[0].width * scale)
                height = int(metadata.streams[0].height * scale)
                if width % 2 == 1:
                    width = width + 1
                if height % 2 == 1:
                    height = height + 1
                options.width = width
                options.height = height
                slice_fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_slice_media",
                                                    environment_name=config.modal_environment)
                slice_result, sentry_trace = await slice_fn.remote.aio(media, [FFMpegSliceSegment(
                    start=TimeDelta.from_format_string(start_time), end=TimeDelta.from_format_string(end_time))], options,
                                                                       sentry_trace)
                video = MediaSource.from_str(slice_result[0].urn)
                logger.success("截取完成")
                video_path = os.path.join(config.S3_mount_dir, video.path)

                # 3、上传到gemini
                logger.info("3、视频文件开始上传到Gemini")
                video_gemini, code = upload(video_path)
                if code == 200:
                    video_gemini_uri = video_gemini
                else:
                    logger.error("视频文件上传Gemini失败")
                    raise Exception("视频文件上传Gemini失败")

                async def inference_api_first_stage():
                    try:
                        logger.info("🎈开始一阶段推理")
                        product_list_str = ""
                        for i, product in enumerate(product_title_list):
                            product_list_str += f"{i}. {product}\n"
                        image_parts = []
                        for i in product_grid_list:
                            image_parts.append(
                                {
                                    "file_data": {
                                        "mime_type": "image/jpeg",
                                        "file_uri": f"{i}"
                                    }
                                }
                            )
                        image_parts.append(
                            {
                                "text": IMAGE_PRODUCT_IDENTIFICATION_PROMPT.format(PRODUCT_LIST=product_list_str)
                            }
                        )
                        json_data = {
                            "safetySettings": [
                                {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
                                {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
                                {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
                                {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
                                {"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "BLOCK_NONE"}
                            ],
                            "contents": [
                                {
                                    "role": "user",
                                    "parts": image_parts
                                }
                            ],
                            "generationConfig": {
                                "temperature": 0.01,
                                "top_p": 0.7
                            }
                        }
                        resp = requests.post(
                            f"https://gateway.ai.cloudflare.com/v1/67720b647ff2b55cf37ba3ef9e677083/bowong-dev/google-vertex-ai/v1/projects/gen-lang-client-0413414134/locations/{random.choice(gemini_region)}/publishers/google/models/gemini-2.5-flash:generateContent",
                            headers={
                                "Content-Type": "application/json",
                                "Authorization": f"Bearer {google_api_key}"
                            },
                            json=json_data, timeout=900
                        )
                    except Exception as e:
                        logger.exception(f"😭Gemini一阶段推理请求失败: {e}")
                        raise e
                    if resp.status_code == 200:
                        target_json = resp.json()
                        reason = target_json["candidates"][0]["finishReason"]
                        if "usageMetadata" in target_json:
                            logger.info("😊一阶段用量信息"+json.dumps(target_json["usageMetadata"], indent=4))
                        if reason == "STOP":
                            result_text: str = target_json["candidates"][0]["content"]["parts"][0]["text"]
                            # 解析识别结果
                            identified_products = parse_stage1_result(result_text)
                            return identified_products
                        else:
                            logger.error(f"😭Gemini一阶段推理失败, Reason {reason}")
                            return None
                    else:
                        logger.error(f"😭Gemini一阶段推理失败, 状态码{resp.status_code}")
                        if resp.status_code == 429:
                            logger.warning("🥵请求负载过高, 随机更换地区重试")
                        return None

                async def inference_api_second_stage(identified_products):
                    try:
                        logger.info("🎈开始二阶段推理")
                        # 构建已识别商品清单字符串
                        identified_products_str = ""
                        for i, product in enumerate(identified_products):
                            identified_products_str += f"{i}. {json.dumps(product,ensure_ascii=False)}\n"
                        image_parts = []
                        image_parts.append(
                            {
                                "file_data": {
                                    "mime_type": "video/mp4",
                                    "file_uri": f"{video_gemini_uri}"
                                }
                            })
                        image_parts.append(
                            {
                                "text": VIDEO_TIMELINE_ANALYSIS_PROMPT.format(IDENTIFIED_PRODUCTS=identified_products_str)
                            }
                        )
                        json_data = {
                            "safetySettings": [
                                {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
                                {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
                                {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
                                {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
                                {"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "BLOCK_NONE"}
                            ],
                            "contents": [
                                {
                                    "role": "user",
                                    "parts": image_parts
                                }
                            ],
                            "generationConfig": {
                                "temperature": 0.01,
                                "top_p": 0.7
                            }
                        }
                        resp = requests.post(
                            f"https://gateway.ai.cloudflare.com/v1/67720b647ff2b55cf37ba3ef9e677083/bowong-dev/google-vertex-ai/v1/projects/gen-lang-client-0413414134/locations/{random.choice(gemini_region)}/publishers/google/models/gemini-2.5-flash:generateContent",
                            headers={
                                "Content-Type": "application/json",
                                "Authorization": f"Bearer {google_api_key}"
                            },
                            json=json_data, timeout=900
                        )
                    except Exception as e:
                        logger.exception(f"😭Gemini二阶段推理请求失败: {e}")
                        raise e
                    if resp.status_code == 200:
                        target_json = resp.json()
                        reason = target_json["candidates"][0]["finishReason"]
                        if "usageMetadata" in target_json:
                            logger.info("😊二阶段用量信息"+json.dumps(target_json["usageMetadata"], indent=4))
                        if reason == "STOP":
                            result_text: str = target_json["candidates"][0]["content"]["parts"][0]["text"]
                            if len(result_text)>10:
                                parts = result_text.split("```")[-2].replace("```", "").replace("json\n", "").replace("\n",
                                                                                                            "").replace("\\", "")
                                parts = json.loads(parts)
                                logger.info(f"👌合并前 {json.dumps(parts, indent=4, ensure_ascii=False)}")
                                # 合并产品和时间线
                                parts = merge_product_data(parts, start_time, end_time, merge_diff=5)
                                for part in parts:
                                    part["product"] = re.sub(r'^\x20*\d+\.\x20*', '', part["product"])
                                return parts
                            else:
                                return []
                        else:
                            logger.error(f"😭Gemini二阶段推理失败, Reason {reason}")
                            return None
                    else:
                        logger.error(f"😭Gemini二阶段推理失败, 状态码{resp.status_code}")
                        if resp.status_code == 429:
                            logger.warning("🥵请求负载过高, 随机更换地区重试")
                        return None

                logger.info("4、发起Gemini推理")

                # 一阶段推理
                product_info = None
                first_retry_time = retry_time
                while product_info is None and first_retry_time > 0:
                    product_info = await inference_api_first_stage()
                    first_retry_time -= 1
                if product_info is None:
                    raise Exception("一阶段推理失败")
                else:
                    logger.info(f"一阶段推理完成JSON \n{json.dumps(product_info, indent=4, ensure_ascii=False)}")

                # 二阶段推理
                product_timeline_info = None
                second_retry_time = retry_time
                while product_timeline_info is None and second_retry_time > 0:
                    product_timeline_info = await inference_api_second_stage(product_info)
                    second_retry_time -= 1
                if product_timeline_info is None:
                    raise Exception("二阶段推理失败")
                else:
                    logger.info(f"二阶段推理完成JSON \n{json.dumps(product_timeline_info, indent=4, ensure_ascii=False)}")

                return product_timeline_info, sentry_trace

            except Exception as e:
                logger.exception(f"推理失败, {e}")
                raise Exception(f"推理失败, {e}")

        return await _handler(media, google_api_key, product_grid_list, product_list, start_time, end_time, options,
                              sentry_trace, retry_time, scale)