import modal

from ..ffmpeg_app import ffmpeg_worker_image, app, config, s3_mount, local_copy_to_s3, output_path_prefix

with ffmpeg_worker_image.imports():
    from BowongModalFunctions.models.ffmpeg_worker_model import FFMpegSliceSegment
    from BowongModalFunctions.models.media_model import MediaSources, MediaProtocol, MediaSource
    from BowongModalFunctions.models.web_model import SentryTransactionInfo, WebhookNotify, FFMPEGResult
    from BowongModalFunctions.utils.PathUtils import FileUtils
    from BowongModalFunctions.utils.SentryUtils import SentryUtils
    from BowongModalFunctions.utils.VideoUtils import VideoUtils, FFMPEGSliceOptions
    import sentry_sdk
    from loguru import logger
    from typing import Optional, Tuple, List
    from modal import current_function_call_id


    @app.function(
        timeout=900,
        cloud="aws",
        cpu=(0.5, 64),
        max_containers=config.ffmpeg_slice_worker_concurrency,
        volumes={
            s3_mount: modal.CloudBucketMount(
                bucket_name=config.S3_bucket_name,
                secret=modal.Secret.from_name("aws-s3-secret",
                                              environment_name=config.modal_environment),
            ),
        },
    )
    @modal.concurrent(max_inputs=1)
    async def ffmpeg_slice_media(media: MediaSource,
                                 markers: List[FFMpegSliceSegment],
                                 options: FFMPEGSliceOptions,
                                 sentry_trace: Optional[SentryTransactionInfo] = None,
                                 webhook: Optional[WebhookNotify] = None) -> Tuple[
        List[FFMPEGResult], Optional[SentryTransactionInfo]]:
        fn_id = current_function_call_id()

        @SentryUtils.sentry_tracker(name="视频文件/直播切割任务", op="ffmpeg.slice.media", fn_id=fn_id,
                                    sentry_trace_id=sentry_trace.x_trace_id if sentry_trace else None,
                                    sentry_baggage=sentry_trace.x_baggage if sentry_trace else None)
        @SentryUtils.webhook_handler(webhook=webhook, func_id=fn_id)
        async def ffmpeg_slice_process(media_source: MediaSource, media_markers: List[FFMpegSliceSegment],
                                       fn_id: str, options: FFMPEGSliceOptions = None, is_streams: bool = False,
                                       ) -> List[FFMPEGResult]:
            if not is_streams:
                cache_filepath = f"{s3_mount}/{media_source.cache_filepath}"
                logger.info(f"从{media_source.urn}切割")
                for i, marker in enumerate(media_markers):
                    logger.info(f"[{i}] {marker.start.toFormatStr()} --> {marker.end.toFormatStr()}")
            else:
                cache_filepath = media_source.path

            segments = await VideoUtils.ffmpeg_slice_media(media_path=cache_filepath,
                                                           media_markers=media_markers,
                                                           options=options,
                                                           is_streams=is_streams,
                                                           output_path=f"{output_path_prefix}/{config.modal_environment}/slice/outputs/{fn_id}/output.mp4")
            return [FFMPEGResult(urn=local_copy_to_s3([segment[0]])[0], metadata=segment[1],
                                 content_length=FileUtils.get_file_size(segment[0])) for segment in segments]

        match media.protocol:
            case MediaProtocol.hls:
                outputs = await ffmpeg_slice_process(media_source=media,
                                                     media_markers=markers,
                                                     options=options,
                                                     is_streams=True,
                                                     fn_id=fn_id)
            case MediaProtocol.vod:
                outputs = await ffmpeg_slice_process(media_source=media,
                                                     media_markers=markers,
                                                     options=options,
                                                     is_streams=False,
                                                     fn_id=fn_id)
            case MediaProtocol.s3:
                outputs = await ffmpeg_slice_process(media_source=media,
                                                     media_markers=markers,
                                                     options=options,
                                                     is_streams=False,
                                                     fn_id=fn_id)
            case _:  # webhook不会报错，需要确认
                raise NotImplementedError("暂不支持的协议")

        return [result for result in outputs], sentry_trace