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.cache_tasks.models import MediaSource
    from BowongModalFunctions.models.ffmpeg_tasks.models import SentryTransactionInfo, WebhookNotify, FFMPEGResult
    from BowongModalFunctions.utils.PathUtils import FileUtils
    from BowongModalFunctions.utils.SentryUtils import SentryUtils
    from BowongModalFunctions.utils.VideoUtils import VideoUtils
    import sentry_sdk
    import os
    import shutil
    from typing import Optional, Tuple, List
    from modal import current_function_call_id


    @app.function(timeout=600, cloud="aws",
                  max_containers=config.ffmpeg_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_subtitle_apply(media: MediaSource, subtitle: Optional[MediaSource],
                                    embed_subtitle: Optional[MediaSource],
                                    fonts: Optional[List[MediaSource]],
                                    sentry_trace: Optional[SentryTransactionInfo] = None,
                                    webhook: Optional[WebhookNotify] = None) -> Tuple[
        FFMPEGResult, Optional[SentryTransactionInfo]]:
        fn_id = current_function_call_id()

        @SentryUtils.sentry_tracker(name="视频叠加字幕", op="ffmpeg.subtitle.apply", 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_process(video: MediaSource, subtitle: Optional[MediaSource],
                                 embed_subtitle: Optional[MediaSource], func_id: str,
                                 fonts: Optional[List[MediaSource]], ) -> FFMPEGResult:
            media_path = f"{s3_mount}/{video.cache_filepath}"
            subtitle_path = f"{s3_mount}/{subtitle.cache_filepath}" if subtitle else None
            embed_subtitle_path = f"{s3_mount}/{embed_subtitle.cache_filepath}" if embed_subtitle else None
            output_path = f"{output_path_prefix}/{config.modal_environment}/subtitle_apply/{func_id}/output.mp4"
            if fonts:
                font_dir = f"{output_path_prefix}/{config.modal_environment}/subtitle_apply/{func_id}/fonts"
                os.makedirs(font_dir, exist_ok=True)
                local_fonts = [f"{s3_mount}/{font.cache_filepath}" for font in fonts]
                for font in local_fonts:
                    font_filename = os.path.basename(font)
                    shutil.copy(font, f"{font_dir}/{font_filename}")
            else:
                font_dir = None
            local_output, metadata = await VideoUtils.ffmpeg_subtitle_apply(media_path=media_path,
                                                                            embed_subtitle_path=embed_subtitle_path,
                                                                            subtitle_path=subtitle_path,
                                                                            font_dir=font_dir, output_path=output_path)
            s3_outputs = local_copy_to_s3([local_output])
            return FFMPEGResult(urn=s3_outputs[0], metadata=metadata,
                                content_length=FileUtils.get_file_size(local_output), )

        result = await ffmpeg_process(video=media, subtitle=subtitle, embed_subtitle=embed_subtitle, fonts=fonts,
                                      func_id=fn_id)
        if not sentry_trace:
            sentry_trace = SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(),
                                                 x_baggage=sentry_sdk.get_baggage())
        return result, sentry_trace
