import os
import shutil
import modal


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

with ffmpeg_worker_image.imports():
    from BowongModalFunctions.models.web_model import SentryTransactionInfo, WebhookNotify, FFMPEGResult, \
        WebhookMethodEnum, BaseFFMPEGTaskStatusResponse, TaskStatus
    from BowongModalFunctions.utils.PathUtils import FileUtils
    from BowongModalFunctions.utils.VideoUtils import VideoUtils
    import sentry_sdk
    import backoff
    import httpx
    import m3u8
    from loguru import logger
    from typing import Optional, Tuple
    from modal import current_function_call_id
    from watchdog.events import FileMovedEvent, FileCreatedEvent, FileSystemEventHandler
    from watchdog.observers import Observer
    import boto3
    from botocore.client import BaseClient, Config

    hls_recording_volume = modal.Volume.from_name("stream_records", create_if_missing=True)
    hls_recording_mount_point = "/mnt/stream_records"


    class PlaylistEventHandler(FileSystemEventHandler):
        update_counter: int = 0
        fn_id: str
        webhook: Optional[WebhookNotify] = None
        s3_mount_output_dir: str
        s3_root_output_dir: str
        boto3_client: BaseClient

        def __init__(self, fn_id: str, s3_mount_output_dir: str, webhook: Optional[WebhookNotify] = None, *args,
                     **kwargs):
            super().__init__(*args, **kwargs)
            self.webhook = webhook
            self.fn_id = fn_id
            self.s3_mount_output_dir = s3_mount_output_dir
            self.boto3_client = boto3.client("s3",
                                             aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
                                             aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
                                             region_name=config.S3_region,
                                             endpoint_url="https://s3-accelerate.amazonaws.com",
                                             config=Config(
                                                 s3={'addressing_style': 'virtual'},
                                                 signature_version='s3v4', )
                                             )
            self.s3_root_output_dir = s3_mount_output_dir.replace(s3_mount, "")

        def on_created(self, event: FileCreatedEvent) -> None:
            logger.info(f"[created] {event.src_path}")
            if event.src_path.endswith(".ts"):
                filename = os.path.basename(event.src_path)
                mount_path = f"{self.s3_mount_output_dir}/{filename}"
                # 将Volume内的ts文件复制到S3挂载点
                shutil.copy(event.src_path, mount_path)
                logger.info(f"[copy] {event.src_path} -> {mount_path}")
            else:
                return

        def on_moved(self, event: FileMovedEvent) -> None:
            logger.info(f"[moved] {event.src_path} -> {event.dest_path}")
            if not event.dest_path.endswith(".m3u8"):
                return
            filename = os.path.basename(event.dest_path)
            filename.replace('.tmp', '')
            mount_path = f"{self.s3_mount_output_dir}/{filename}"
            # 将Volume内的playlist.m3u8.tmp复制到S3挂载点的playlist.m3u8文件
            try:
                shutil.copy(event.dest_path, mount_path)
            except Exception as e:
                logger.exception(e)
                filename = os.path.basename(event.dest_path)
                self.boto3_client.upload_file(event.dest_path, f"{self.s3_root_output_dir}/{filename}")
            logger.info(f"[copy] {event.dest_path} -> {mount_path}")
            self.update_counter += 1
            if self.webhook and self.update_counter == 1:
                logger.info("[Start] webhook trigger")
                try:
                    self.webhook_on_start()
                except Exception as e:
                    logger.exception(e)

        @backoff.on_exception(exception=Exception, wait_gen=backoff.constant,
                              max_time=5, max_tries=5, raise_on_giveup=True)
        def webhook_on_start(self):
            """
            开始录制第一次更新时回调
            """
            webhook = self.webhook
            if webhook.method is not WebhookMethodEnum.POST:
                logger.warning(f"webhook method {webhook.method.value} not supported")
            body = BaseFFMPEGTaskStatusResponse(taskId=self.fn_id,
                                                task_type="ffmpeg_stream_record_as_hls",
                                                status=TaskStatus.running).model_dump()
            response = httpx.post(url=webhook.endpoint.__str__(),
                                  json=body,
                                  headers=webhook.headers)
            response.raise_for_status()
            logger.info(f"[Start] webhook {response.status_code} {response.text}")


    @app.function(timeout=43200 + 300,  # 最长处理12h的录制任务 + 5分钟的清理缓存工作
                  cloud="aws",
                  # todo: 暂时不限制最大同时存在的录制数量
                  # 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),
                      ),
                      hls_recording_mount_point: hls_recording_volume
                  }, )
    @modal.concurrent(max_inputs=5)
    async def ffmpeg_stream_record_as_hls(stream_url: str, segment_duration: int, recording_timeout: int,
                                          first_segment_duration: int,
                                          webhook: Optional[WebhookNotify] = None, monitor_timeout: int = 36000,
                                          sentry_trace: Optional[SentryTransactionInfo] = None, ) -> Tuple[
        FFMPEGResult, Optional[SentryTransactionInfo]]:
        fn_id = current_function_call_id()
        output_dir = f"{config.modal_environment}/records/hls/{fn_id}"
        s3_mount_output_dir = f"{s3_mount}/{output_dir}"
        os.makedirs(s3_mount_output_dir, exist_ok=True)
        volume_output_dir = f"{hls_recording_mount_point}/{output_dir}"
        os.makedirs(volume_output_dir, exist_ok=True)
        logger.info(f"manifest = {volume_output_dir}/playlist.m3u8")

        @backoff.on_exception(exception=Exception, wait_gen=backoff.constant,
                              max_time=5, max_tries=5, raise_on_giveup=False)
        async def webhook_on_end(fn_id: str, webhook: WebhookNotify, result: FFMPEGResult):
            if webhook.method is not WebhookMethodEnum.POST:
                logger.warning(f"webhook method {webhook.method.value} not supported")
            body = BaseFFMPEGTaskStatusResponse(taskId=fn_id,
                                                task_type="ffmpeg_stream_record_as_hls",
                                                status=TaskStatus.success,
                                                results=[result])
            response = httpx.post(url=webhook.endpoint.__str__(),
                                  json=body.model_dump(),
                                  headers=webhook.headers)
            logger.info(f"[End] webhook {response.status_code} {response.text}")

        @backoff.on_exception(exception=Exception, wait_gen=backoff.constant,
                              max_time=5, max_tries=5, raise_on_giveup=False)
        async def webhook_on_error(fn_id: str, webhook: WebhookNotify, error: str):
            if webhook.method is not WebhookMethodEnum.POST:
                logger.warning(f"webhook method {webhook.method.value} not supported")
            body = BaseFFMPEGTaskStatusResponse(taskId=fn_id,
                                                task_type="ffmpeg_stream_record_as_hls",
                                                error=error,
                                                status=TaskStatus.failed).model_dump()
            response = httpx.post(url=webhook.endpoint.__str__(),
                                  json=body,
                                  headers=webhook.headers)
            logger.info(f"[End] webhook {response.status_code} {response.text}")

        playlist_handler = PlaylistEventHandler(webhook=webhook, fn_id=fn_id, s3_mount_output_dir=s3_mount_output_dir)
        playlist_observer = Observer()
        os.makedirs(volume_output_dir, exist_ok=True)
        # 监控本地Volume下录制缓存目录
        playlist_observer.schedule(playlist_handler, path=volume_output_dir, recursive=False)
        playlist_observer.start()
        try:
            playlist = await VideoUtils.ffmpeg_stream_record_as_hls(stream_url=stream_url,
                                                         first_segment_duration=first_segment_duration,
                                                         segment_duration=segment_duration,
                                                         stream_content_timeout=recording_timeout,
                                                         stream_monitor_timeout=monitor_timeout,
                                                         segments_output_dir=volume_output_dir,
                                                         playlist_output_dir=volume_output_dir)
            # 确保playlist有 #EXT-X-ENDLIST 结尾
            pl = m3u8.load(playlist)
            pl.is_endlist = True
            pl.dump(playlist)
            logger.info(f"[End] playlist = {playlist}")

        except Exception as e:
            logger.exception(e)
            playlist_observer.stop()
            if webhook:
                await webhook_on_error(webhook=webhook, fn_id=fn_id,
                                       error=e.message if hasattr(e, 'message') else str(e))
            raise e
        playlist_observer.stop()
        content_length = FileUtils.get_folder_size(volume_output_dir)
        metadata = VideoUtils.ffprobe_media_metadata(media_path=f"{volume_output_dir}/playlist.m3u8")
        result = FFMPEGResult(urn=f"s3://{config.S3_region}/{config.S3_bucket_name}/{output_dir}/playlist.m3u8",
                              metadata=metadata,
                              content_length=content_length, )
        if webhook:
            await webhook_on_end(webhook=webhook, fn_id=fn_id, result=result)
        if not sentry_trace:
            sentry_trace = SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(),
                                                 x_baggage=sentry_sdk.get_baggage())
        shutil.rmtree(volume_output_dir)
        return result, sentry_trace
