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')}")
        def calculate_target_dimensions(original_width, original_height, target_height=1280):
            """计算保持宽高比的目标尺寸，高度固定为1280"""
            aspect_ratio = original_width / original_height
            target_width = int(target_height * aspect_ratio)

            # 确保宽度是偶数（H.265编码要求）
            if target_width % 2 != 0:
                target_width -= 1

            return target_width, target_height

        def convert_video(input_file, output_file=None):
            """将视频转换为30fps、720p、H.264编码(CRF=16)"""
            if not output_file:
                # 生成默认输出文件名
                base_name, ext = os.path.splitext(input_file)
                base_name = base_name.split(os.sep)[-1]
                output_file = f"/root/{base_name}_720p_h265.mp4"

            try:
                # 第一步：获取原始视频宽高
                ffprobe = FFprobe(
                    inputs={input_file: ['-v', 'error', '-select_streams', 'v:0',
                                         '-show_entries', 'stream=width,height',
                                         '-of', 'csv=p=0']}
                )
                info = ffprobe.run(stdout=subprocess.PIPE, stderr=subprocess.PIPE)[0].decode('utf-8').strip()
                width, height = map(int, info.split(','))

                # 计算目标尺寸
                target_width, target_height = calculate_target_dimensions(width, height, target_height=int(
                    0.5 * height) if height > 1600 else 1138)

                # 第二步：执行视频转换
                ff = FFmpeg(
                    inputs={input_file: None},
                    outputs={output_file: [
                        '-vf', f'scale={target_width}:{target_height}',
                        '-r', '30',  # 帧率设置为30fps
                        '-c:v', 'libx264',  # 使用H.264编码
                        "-b:v", "2500K",
                        '-crf', '18',  # 质量设置（CRF 16）
                        '-c:a', 'copy',  # 音频编码复制
                        "-fs", "1536MB",
                        '-y'  # 覆盖已存在文件
                    ]}
                )

                logger.info(f"开始转换视频: {input_file}")
                logger.info(f"目标分辨率: {target_width}x{target_height}")
                logger.info(f"输出文件: {output_file}")

                ff.run()
                logger.success("视频转换完成!")
                return output_file

            except Exception as e:
                logger.exception(f"视频转换失败 {e}")
                return

        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 = 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:
                pre_upload_response = client.post(
                    url=f"https://generativelanguage.googleapis.com/upload/v1beta/files?key={google_api_key}",
                    headers={
                        "X-Goog-Upload-Protocol": "resumable",
                        "X-Goog-Upload-Command": "start",
                        "X-Goog-Upload-Header-Content-Length": str(content_length),
                        "X-Goog-Upload-Header-Content-Type": content_type
                    },
                    json={
                        "file": {
                            "display_name": filename.split(".")[0],
                        }
                    })
                pre_upload_response.raise_for_status()

                upload_url = pre_upload_response.headers.get("X-Goog-Upload-Url")

                upload_response = client.post(url=upload_url, content=video, headers={
                    "X-Goog-Upload-Offset": "0",
                    "X-Goog-Upload-Command": "upload, finalize",
                    "Content-Type": content_type
                })
                upload_response.raise_for_status()

            return upload_response.json(), upload_response.status_code

        @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
                product_info_txt = "\n".join(product_title_list)
                logger.info(f"\nmedia || {media.urn},"
                            f"\nstart_time || {start_time}, "
                            f"\nend_time || {end_time}, "
                            f"\nproduct_grid_list || {product_grid_list}, "
                            f"\nproduct_list || {product_info_txt}"
                            f"\noptions || {options.model_dump_json()}, "
                            f"\nscale || {scale}")
                image_parts = [{
                    "file_data": {
                        "mime_type": "image/jpeg",
                        "file_uri": f"{i}"
                    }
                } for i in product_grid_list]
                # 2、切20分钟的条
                logger.info("2、开始截取指定视频")
                # 计算缩放尺寸
                metadata = VideoUtils.ffprobe_media_metadata(media.path)
                width = metadata.streams[0].width * scale
                height = metadata.streams[0].height * scale
                if width % 2 == 1:
                    width = width + 1
                if height % 2 == 1:
                    height = height + 1
                options.width = int(width)
                options.height = int(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("截取完成")
                # 3、转换为720p h265
                # logger.info("3、转换视频为720p h264")
                # video_path = convert_video(os.path.join(config.S3_mount_dir, video.path))
                video_path = os.path.join(config.S3_mount_dir, video.path)
                # 4、上传到gemini
                logger.info("3、视频文件开始上传到Gemini")
                video_gemini, code = upload(video_path)
                if code == 200:
                    video_gemini_uri = video_gemini["file"]["uri"]
                else:
                    logger.error("视频文件上传Gemini失败")
                    raise Exception("视频文件上传Gemini失败")

                # 5、检查文件是否已处理完成
                def check():
                    with httpx.Client(timeout=Timeout(timeout=30)) as client:
                        file = video_gemini_uri.split("/")[-1]
                        response = client.get(
                            url=f"https://generativelanguage.googleapis.com/v1beta/files/{file}?key={google_api_key}")
                        response.raise_for_status()
                    if response.status_code == 200:
                        if response.json()["state"] == "ACTIVE":
                            return True
                    return False

                check_num = 60
                logger.info("4、开始检查Gemini文件是否处理完成")
                while check_num > 0:
                    logger.info(f"检查中, 剩余检查次数{check_num}")
                    ret = check()
                    if ret:
                        break
                    else:
                        check_num -= 1
                        time.sleep(5)
                if check_num <= 0:
                    raise Exception("Gemini文件上传处理状态检查超时")
                logger.success("Gemini文件处理完成")

                # 6、执行Gemini推理操作
                async def inference_api():
                    try:
                        logger.info("请求推理接口")
                        prompt = os.environ["prompt_text"]
                        logger.info(f"Prompt text: {prompt}")
                        image_parts.extend([
                            {
                                "file_data": {
                                    "mime_type": "video/mp4",
                                    "file_uri": f"{video_gemini_uri}"
                                }
                            },
                            {
                                "text": prompt.format(product_title_list)
                            }
                        ])
                        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": [
                                {
                                    "parts": image_parts
                                }
                            ],
                            "generationConfig": {
                                "temperature": 0.01
                            }
                        }
                        resp = requests.post(
                            "https://gateway.ai.cloudflare.com/v1/67720b647ff2b55cf37ba3ef9e677083/bowong-dev/google-ai-studio/v1beta/models/gemini-2.5-flash-preview-05-20:generateContent",
                            headers={
                                "Content-Type": "application/json",
                                "X-Goog-Api-Key": f"{google_api_key}"
                            },
                            json=json_data, timeout=900
                        )
                    except Exception as e:
                        logger.exception(f"Gemini推理请求失败: {e}")
                        raise e
                    if resp.status_code == 200:
                        return resp.json()
                    else:
                        logger.error(f"Gemini推理失败, 状态码{resp.status_code}")
                        if resp.status_code == 429:
                            logger.warning("请求频率过高, 等待50-70秒后重试")
                            await asyncio.sleep(random.randint(50, 70))
                        return None

                logger.info("5、发起Gemini推理")
                target_json = None
                while target_json is None and retry_time > 0:
                    target_json = await inference_api()
                    retry_time -= 1
                if target_json is None:
                    raise Exception("推理失败")
                logger.info(f"推理完成JSON \n{json.dumps(target_json, indent=4, ensure_ascii=False)}")

                reason = target_json["candidates"][0]["finishReason"]
                if reason == "STOP":
                    parts: str = target_json["candidates"][0]["content"]["parts"][0]["text"]
                    parts = parts.split("```")[-2].replace("```", "").replace("json\n", "").replace("\n", "").replace("\\", "")
                    parts = json.loads(parts)
                    # 合并产品和时间线
                    parts = merge_product_data(parts, start_time, end_time)
                    logger.info("\n最终结果: {}".format(json.dumps(parts, ensure_ascii=False, indent=2)))
                    return parts, sentry_trace
                else:
                    raise Exception(f"Gemini推理失败, 失败原因: {reason}")

            except Exception as e:
                logger.exception(f"推理失败, {e}")
                raise Exception(f"推理失败, {e}")
            finally:
                if video_gemini_uri:
                    logger.info("6、清除Gemini临时文件")
                    with httpx.Client(timeout=Timeout(timeout=120)) as client:
                        resp = client.delete(
                            f'https://bowongai-{config.modal_environment}--{config.modal_app_name}-fastapi-webapp.modal.run/google/delete',
                            params={"filename": video_gemini_uri.split("/")[-1]},
                            headers={"x-google-api-key": google_api_key})
                        resp.raise_for_status()
                    if resp.status_code == 200:
                        logger.success("Gemini临时文件清除成功")
                    else:
                        logger.warning("Gemini临时文件清除失败, 请自行清除")

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