import json
import os
from typing import Annotated, Optional, cast

import modal
import sentry_sdk
from fastapi.responses import Response
from loguru import logger
import httpx
from fastapi import APIRouter, UploadFile, Header, HTTPException
from pydantic import BaseModel, Field
from starlette import status
from starlette.responses import JSONResponse

from BowongModalFunctions.config import WorkerConfig
from BowongModalFunctions.models.web_model import SentryTransactionInfo, GeminiResultResponse, GeminiRequest, \
    ModalTaskResponse
from BowongModalFunctions.utils.ModalUtils import ModalUtils

config = WorkerConfig()

router = APIRouter(prefix="/google", tags=["Google"])


class GoogleAPIKeyHeaders(BaseModel):
    x_google_api_key: Optional[str] = Field(description="Google API Key", default=None)

class BundleHeaders(BaseModel):
    x_google_api_key: Optional[str] = Field(description="Google API Key")
    x_trace_id: str = Field(description="Sentry Transaction ID", default=None)
    x_baggage: str = Field(description="Sentry Transaction baggage", default=None)


@router.post("/upload",
             summary="上传文件到Google File",
             description="上传文件到Google File, 换取Google File URI, 不同Google API Key之间的URI不互通, 最多可为每个项目存储 20 GB 的文件，每个文件的大小上限为 2 GB。文件会存储 48 小时")
async def upload_file_multipart(file: UploadFile,
                                headers: Annotated[GoogleAPIKeyHeaders, Header()]):
    google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY")
    if not google_api_key:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google API Key")
    content_length = file.size
    content_type = file.content_type
    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 = {file.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": file.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=file.file.read(), headers={
            "X-Goog-Upload-Offset": "0",
            "X-Goog-Upload-Command": "upload, finalize",
            "Content-Type": content_type
        })
        upload_response.raise_for_status()

        return JSONResponse(content=upload_response.json(), status_code=upload_response.status_code)


@router.get("/status", summary="获取已上传文件的处理状态")
async def uploaded_file_status(filename: str,
                               headers: Annotated[GoogleAPIKeyHeaders, Header()]):
    google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY")
    if not google_api_key:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google API Key")
    with httpx.Client() as client:
        response = client.get(
            url=f"https://generativelanguage.googleapis.com/v1beta/files/{filename}?key={google_api_key}")
        response.raise_for_status()
    return JSONResponse(content=response.json(), status_code=response.status_code)


@router.delete('/delete', summary="删除已上传的文件")
async def delete_file(filename: str, headers: Annotated[GoogleAPIKeyHeaders, Header()]):
    google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY")
    if not google_api_key:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google API Key")
    with httpx.Client() as client:
        response = client.delete(
            url=f"https://generativelanguage.googleapis.com/v1beta/files/{filename}?key={google_api_key}")
        response.raise_for_status()
    return JSONResponse(content=response.json(), status_code=response.status_code)


@router.get('/list', summary="列出已上传的文件")
async def list_files(headers: Annotated[GoogleAPIKeyHeaders, Header()]):
    google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY")
    if not google_api_key:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google API Key")
    with httpx.Client() as client:
        response = client.get(
            url=f"https://generativelanguage.googleapis.com/v1beta/files?key={google_api_key}")
        response.raise_for_status()
    return JSONResponse(content=response.json(), status_code=response.status_code)

@router.post('/inference_gemini', summary="使用Gemini推理hls视频流指定时间段的打点情况")
async def inference_gemini(
                    data:GeminiRequest,
                    headers: Annotated[BundleHeaders, Header()],
                    ) -> ModalTaskResponse:
    google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY")
    if not google_api_key:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google API Key")
    fn = modal.Function.from_name(config.modal_app_name,"video_hls_slice_inference", environment_name=config.modal_environment)
    fn_call = fn.spawn(data.media_hls_url, google_api_key, data.product_list, data.start_time, data.end_time,
                     SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(), x_baggage=sentry_sdk.get_baggage())
                     if headers.x_trace_id is None else SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage))
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)

@router.get("/inference_gemini/{task_id}", summary="查询Gemini推理hls视频流指定时间段的打点任务")
async def gemini_status(task_id: str, response: Response) -> GeminiResultResponse:
    task_info = await ModalUtils.get_modal_task_status(task_id)

    if task_info.transaction:
        response.headers["x-trace-id"] = task_info.transaction.x_trace_id
        response.headers["x-baggage"] = task_info.transaction.x_baggage
    try:
        return GeminiResultResponse(taskId=task_id, status=task_info.status, code=cast(int, task_info.error_code.value),
                                    error=task_info.error_reason, result=json.dumps(task_info.results[0], ensure_ascii=False)
                                    if task_info.results and len(task_info.results) >= 1 else "")
    except Exception as e:
        logger.exception(f"获取Gemini状态发生错误 {e}")
        return GeminiResultResponse(taskId=task_id, status=task_info.status, code=cast(int, task_info.error_code.value),
                                    error=task_info.error_reason + f"获取Gemini状态发生错误 {e}"
                                    if task_info.error_reason else f"获取Gemini状态发生错误 {e}", result="")