import asyncio
import datetime
import os
from typing import Annotated, Optional

import modal
from loguru import logger
from modal import current_function_call_id
import sentry_sdk
from fastapi import APIRouter, Depends, UploadFile, HTTPException, File, Form
from fastapi.responses import JSONResponse, RedirectResponse
from starlette import status
import boto3
from botocore.config import Config

from ..config import WorkerConfig
from ..middleware.authorization import verify_token
from ..models.media_model import (MediaSources,
                                  CacheResult,
                                  MediaSource,
                                  MediaCacheStatus,
                                  DownloadResult,
                                  UploadResultResponse,
                                  UploadBase64Request, UploadPresignRequest, UploadPresignResponse,
                                  UploadMultipartPresignRequest, UploadMultipartPresignResponse
                                  )
from ..models.web_model import SentryTransactionInfo, MonitorLiveRoomProductRequest, ModalTaskResponse, \
    LiveRoomProductCachesResponse
from ..utils.KVCache import MediaSourceKVCache, LiveProductKVCache
from ..utils.SentryUtils import SentryUtils

config = WorkerConfig()

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', )
                      )

router = APIRouter(prefix="/cache", tags=['缓存'], )

modal_kv_cache = MediaSourceKVCache(kv_name=config.modal_kv_name,
                                    environment=config.modal_environment)

modal_kv_product_cache = LiveProductKVCache(kv_name=config.modal_product_kv_name,
                                            environment=config.modal_environment)


@router.post("/",
             summary="缓存视频文件",
             description="异步缓存视频文件到S3存储桶和Modal Dict(KV)",
             dependencies=[Depends(verify_token)])
async def cache(medias: MediaSources) -> CacheResult:
    fn_id = current_function_call_id()
    caches: MediaSources
    sentry_trace = SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(),
                                         x_baggage=sentry_sdk.get_baggage())

    @SentryUtils.sentry_tracker(name="同步视频缓存", op="cache.get", fn_id=fn_id,
                                sentry_trace_id=None, sentry_baggage=None)
    async def cache_handler(media: MediaSource):
        cache_span = sentry_sdk.get_current_span()
        cache_span.set_data("runner_id", fn_id)
        cache_span.set_data("cache.key", [media.urn])
        cached_media = modal_kv_cache.get_cache(media.urn)
        cache_hit: bool = False
        if not cached_media:
            # start new download task
            with cache_span.start_child(name="视频缓存任务入队",
                                        op="queue.publish") as queue_publish_span:
                fn = modal.Function.from_name(config.modal_app_name, 'cache_submit')
                fn_task = fn.spawn(media, sentry_trace)
                queue_publish_span.set_data("cache.key", media.urn)
                queue_publish_span.set_data("messaging.message.id", fn_task.object_id)
                queue_publish_span.set_data("messaging.destination.name",
                                            "video-downloader.cache_submit")
                queue_publish_span.set_data("messaging.message.body.size", 0)
                media.status = MediaCacheStatus.downloading
                media.downloader_id = fn_task.object_id
                modal_kv_cache.set_cache(media)
        else:
            media = cached_media
        match media.status:
            case MediaCacheStatus.ready:
                cache_hit = True
            case MediaCacheStatus.downloading:  # 下载任务已经在进行
                cache_hit = True
            case _:
                # start new download task
                with cache_span.start_child(name="视频缓存任务入队",
                                            op="queue.publish") as queue_publish_span:
                    fn = modal.Function.from_name(config.modal_app_name, 'cache_submit')
                    fn_task = fn.spawn(media, sentry_trace)
                    queue_publish_span.set_data("cache.key", media.urn)
                    queue_publish_span.set_data("messaging.message.id", fn_task.object_id)
                    queue_publish_span.set_data("messaging.destination.name",
                                                "video-downloader.cache_submit")
                    queue_publish_span.set_data("messaging.message.body.size", 0)
                    media.status = MediaCacheStatus.downloading
                    media.downloader_id = fn_task.object_id
                    modal_kv_cache.set_cache(media)
                    cache_hit = False

        cache_span.set_data("cache.hit", cache_hit)
        return media

    async with asyncio.TaskGroup() as group:
        tasks = [group.create_task(cache_handler(media)) for media in medias.inputs]

    cache_task_result_dict = {}
    cache_task_result_list = []

    for task in tasks:
        result = task.result()
        cache_task_result_dict[result.urn] = result.model_dump_json()
        cache_task_result_list.append(result)

    modal_kv_cache.batch_update_cloudflare_kv(cache_task_result_dict)
    return CacheResult(caches={media.urn: media for media in cache_task_result_list})


@router.delete("/",
               summary="清除指定的所有缓存",
               description="清除指定的所有缓存(包括KV记录和S3存储文件)",
               dependencies=[Depends(verify_token)])
async def purge_media_kv_file(medias: MediaSources):
    fn_id = current_function_call_id()
    fn = modal.Function.from_name(config.modal_app_name, "cache_delete", environment_name=config.modal_environment)

    @SentryUtils.sentry_tracker(name="清除媒体源缓存", op="cache.purge", fn_id=fn_id,
                                sentry_trace_id=None, sentry_baggage=None)
    async def purge_handle(media: MediaSource):
        cache_media = modal_kv_cache.pop(media.urn)
        if cache_media:
            deleted_cache: MediaSource = await fn.remote.aio(cache_media)
            return deleted_cache.urn
        return None

    async with asyncio.TaskGroup() as group:
        tasks = [group.create_task(purge_handle(media)) for media in medias.inputs]

    keys = [task.result() for task in tasks]
    modal_kv_cache.batch_remove_cloudflare_kv(keys)
    return JSONResponse(content={"success": True, "keys": keys})


@router.post("/download",
             summary="批量获取下载地址",
             description="获取已缓存的视频下载地址",
             dependencies=[Depends(verify_token)])
@sentry_sdk.trace
async def download_caches(medias: MediaSources) -> DownloadResult:
    cdn_endpoint = config.S3_cdn_endpoint
    urls = []
    for media in medias.inputs:
        urls.append(f"{cdn_endpoint}/{media.get_cdn_url()}")
    return DownloadResult(urls=urls)


@router.get("/download",
            summary="下载已缓存的视频",
            description="通过CDN下载已缓存的视频文件")
@sentry_sdk.trace
async def download_cache(media: str) -> RedirectResponse:
    cdn_endpoint = config.S3_cdn_endpoint
    media = MediaSource.from_str(media)
    return RedirectResponse(url=f"{cdn_endpoint}/{media.get_cdn_url()}", status_code=status.HTTP_302_FOUND)


@router.delete("/kv",
               summary="清除KV记录",
               description="清除当前环境下KV缓存过的所有数据(S3存储桶内的文件会保留)",
               dependencies=[Depends(verify_token)])
async def purge_kv_all():
    parent = sentry_sdk.get_current_span()
    span = parent.start_child(name="清除缓存KV", op="purge.flush")
    modal_kv_cache.clear()
    span.set_data("cache.success", True)
    span.finish()
    return JSONResponse(content={"success": True})


@router.post("/kv",
             summary="删除对应的KV记录",
             description="删除请求中对应的视频缓存记录",
             dependencies=[Depends(verify_token)])
async def purge_kv(medias: MediaSources):
    try:
        for media in medias.inputs:
            modal_kv_cache.pop(media.urn)
        keys = [media.urn for media in medias.inputs]
        modal_kv_cache.batch_remove_cloudflare_kv(keys)
        return JSONResponse(content={"success": True, "keys": keys})
    except Exception as e:
        return JSONResponse(content={"success": False, "error": str(e)})


@router.post("/media",
             summary="清除指定的所有缓存",
             description="清除指定的所有缓存(包括KV记录和S3存储文件), 将要被淘汰，使用DELETE /cache/替代",
             deprecated=True,
             dependencies=[Depends(verify_token)])
async def purge_media(medias: MediaSources):
    fn_id = current_function_call_id()
    fn = modal.Function.from_name(config.modal_app_name, "cache_delete", environment_name=config.modal_environment)

    @SentryUtils.sentry_tracker(name="清除媒体源缓存", op="cache.purge", fn_id=fn_id,
                                sentry_trace_id=None, sentry_baggage=None)
    async def purge_handle(media: MediaSource):
        cache_media = modal_kv_cache.pop(media.urn)
        if cache_media:
            deleted_cache: MediaSource = await fn.remote.aio(cache_media)
            return deleted_cache.urn
        return None

    async with asyncio.TaskGroup() as group:
        tasks = [group.create_task(purge_handle(media)) for media in medias.inputs]

    keys = [task.result() for task in tasks]
    modal_kv_cache.batch_remove_cloudflare_kv(keys)
    return JSONResponse(content={"success": True, "keys": keys})


@router.post("/upload-s3",
             summary="上传文件到S3",
             description="上传文件到S3的文件必须小于200M",
             dependencies=[Depends(verify_token)])
async def s3_upload(file: Annotated[UploadFile, File(description="上传的文件")],
                    prefix: Annotated[Optional[str], Form()] = None) -> UploadResultResponse:
    fn_id = current_function_call_id()

    if file.size > 200 * 1024 * 1024:  # 上传文件不大于200M
        raise HTTPException(status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, detail="上传文件不可超过200MB")

    key = f"upload/{prefix}/{file.filename}" if prefix else f"upload/{file.filename}"
    local_path = f"{config.S3_mount_dir}/{key}"
    logger.info(f"s3上传到{key}, size={file.size}")
    os.makedirs(os.path.dirname(local_path), exist_ok=True)
    with open(local_path, 'wb') as f:
        f.write(file.file.read())
    logger.info(f"{local_path} 保存成功")
    media_source = MediaSource.from_str(f"s3://{config.S3_region}/{config.S3_bucket_name}/{key}")
    media_source.status = MediaCacheStatus.ready
    media_source.downloader_id = fn_id
    return UploadResultResponse(media=media_source)


@router.post('/upload-s3-b64',
             summary="基于Base64格式上传文件到S3",
             description="上传文件到S3当文件必须小于200M",
             dependencies=[Depends(verify_token)])
async def s3_upload_base64(body: UploadBase64Request) -> UploadResultResponse:
    fn_id = current_function_call_id()
    prefix = body.prefix
    file = body.file
    if file.size > 200 * 1024 * 1024:  # 上传文件不大于200M
        raise HTTPException(status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, detail="上传文件不可超过200MB")
    key = f"upload/{prefix}/{file.filename}" if prefix else f"upload/{file.filename}"
    local_path = f"{config.S3_mount_dir}/{key}"
    logger.info(f"s3上传到{key}, size={file.size}")
    os.makedirs(os.path.dirname(local_path), exist_ok=True)
    with open(local_path, 'wb') as f:
        f.write(file.raw_content)
    logger.info(f"{local_path} 保存成功")
    media_source = MediaSource.from_str(f"s3://{config.S3_region}/{config.S3_bucket_name}/{key}")
    media_source.status = MediaCacheStatus.ready
    media_source.downloader_id = fn_id
    return UploadResultResponse(media=media_source)


@router.post("/monitor_live_room_product_trigger",
             summary="触发监控直播间商品信息并缓存",
             description="触发监控直播间商品信息并缓存, 如果直播结束清除缓存, 触发间隔请控制在60s以上",
             dependencies=[Depends(verify_token)])
async def monitor_live_room_product(body: MonitorLiveRoomProductRequest) -> LiveRoomProductCachesResponse:
    fn = modal.Function.from_name(config.modal_app_name, "monitor_live_room_product_trigger",
                                  environment_name=config.modal_environment)
    status = await fn.remote.aio(body.cookie, body.room_id, body.author_id)
    if status == 0:
        product_list = modal_kv_product_cache.get_cache(body.room_id)
        return LiveRoomProductCachesResponse(status=status, cache_json=product_list.model_dump_json())
    elif status == 1:
        return LiveRoomProductCachesResponse(status=status, message="直播已结束")
    elif status == 2:
        return LiveRoomProductCachesResponse(status=status, message="分配到风控IP, 请稍后重试")
    elif status == 3:
        return LiveRoomProductCachesResponse(status=status, message="请求Tikhub API出现错误")
    else:
        return LiveRoomProductCachesResponse(status=4, message="内部错误")


@router.post('/upload-s3/simple/presign',
             summary="S3简单上传预签名",
             description="利用S3就近接入点上传",
             dependencies=[Depends(verify_token)])
async def s3_presign_upload(body: UploadPresignRequest) -> UploadPresignResponse:
    expires_in = 3600
    expired_at = datetime.datetime.now() + datetime.timedelta(seconds=expires_in)
    signed_url = client.generate_presigned_url("put_object",
                                               Params={
                                                   'Bucket': config.S3_bucket_name,
                                                   'Key': f"upload/{body.key}",
                                                   "ContentType": body.content_type,
                                               }, ExpiresIn=expires_in, )
    return UploadPresignResponse(url=signed_url,
                                 urn=f"s3://{config.S3_region}/{config.S3_bucket_name}/upload/{body.key}",
                                 expired_at=expired_at)


@router.post("/upload-s3/multipart/presign",
             summary="S3分片上传预签名",
             description="""
            1. 本地按文件总大小分Chunk大小按urls链接内的顺序通过HTTP PUT请求上传文件分片; 并将上传完成后获得的返回头ETag值记录与PartNumber对应, PartNumber对应使用url在urls内的顺位, 以1开始
            \n2. 所有分片上传完成后使用XML格式拼装出用于确认上传的body; 并通过HTTP POST complete_url确认上传, ContentType 需要确保为application/xml\n\n
                <CompleteMultipartUpload>
	                <Part>
		                <PartNumber>1</PartNumber>
		                <ETag>"60575364b098a1a48765a28c3a48e0ef"</ETag>
	                </Part>
	                <Part>
		                <PartNumber>2</PartNumber>
		                <ETag>"a38691c31fd242faee5533c65b4501d7"</ETag>
	                </Part>
	                <Part>
		                <PartNumber>3</PartNumber>
		                <ETag>"7a7510cc83f98feea28f319567e4cf66"</ETag>
	                </Part>
                </CompleteMultipartUpload>
            \n3. 如无法确认上传，可使用HTTP GET list_url debug当前分片上传状态，确认成功后list_url无法返回有效数据
        """, dependencies=[Depends(verify_token)])
async def s3_presign_upload_multipart(body: UploadMultipartPresignRequest) -> UploadMultipartPresignResponse:
    chunk_count = body.parts_count
    multipart_upload_response = client.create_multipart_upload(Bucket=config.S3_bucket_name, Key=body.key,
                                                               ContentType=body.content_type, )
    upload_id = multipart_upload_response.get("UploadId")
    signed_urls = []
    expires_in = 3600
    expired_at = datetime.datetime.now() + datetime.timedelta(seconds=expires_in)
    for i in range(chunk_count):
        signed_url = client.generate_presigned_url("upload_part",
                                                   Params={
                                                       'Bucket': config.S3_bucket_name,
                                                       'Key': body.key,
                                                       'PartNumber': i + 1,
                                                       'UploadId': upload_id,
                                                   }, ExpiresIn=expires_in)
        signed_urls.append(signed_url)

    signed_completed_url = client.generate_presigned_url("complete_multipart_upload",
                                                         Params={
                                                             'Bucket': config.S3_bucket_name,
                                                             'Key': body.key,
                                                             'UploadId': upload_id,
                                                         }, ExpiresIn=expires_in)
    signed_list_url = client.generate_presigned_url("list_parts",
                                                    Params={
                                                        'Bucket': config.S3_bucket_name,
                                                        'Key': body.key,
                                                        'UploadId': upload_id,
                                                    }, ExpiresIn=expires_in)
    return UploadMultipartPresignResponse(urls=signed_urls,
                                          urn=f"s3://{config.S3_region}/{config.S3_bucket_name}/upload/{body.key}",
                                          expired_at=expired_at,
                                          complete_url=signed_completed_url,
                                          list_url=signed_list_url)
