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

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

from ..models.settings.cluster import WorkerConfig
from ..middleware.authorization import verify_token
from ..models.responses.models import ModalTaskResponse, GeminiResultResponse, GoogleUploadResponse
from ..models.requests.models import GeminiRequest, MakeGridGeminiRequest, GeminiRequestFirstOnly, \
    GeminiRequestSecondOnly
from ..models.ffmpeg_tasks.models import SentryTransactionInfo
from ..utils.ModalUtils import ModalUtils
from ..utils.HTTPUtils import GoogleAuthUtils

config = WorkerConfig()

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


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


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)


class GoogleAuthHeaders(BaseModel):
    Authorization: str = Field(description="Google Auth Bearer Token")

    @computed_field(description="Google Auth Token")
    @property
    def auth_token(self) -> str:
        return self.Authorization[len("Bearer "):]


@router.post("/upload",
             summary="上传文件到Google AI Studio提供的File API, 将被VertexAI新规范替代，统一使用Google CloudStorage存储",
             deprecated=True,
             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
    # 允许上传的视频
    google_content_type_allow_list = ['video/mp4', 'video/mpeg', 'video/mov', 'video/avi', 'video/x-flv', 'video/mpg',
                                      'video/webm', 'video/wmv', 'video/3gpp', ]
    # 允许上传的图片
    google_content_type_allow_list.extend(["image/png", "image/jpeg", "image/webp", "image/heic", "image/heif"])
    if content_type not in google_content_type_allow_list:
        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, 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.post("/vertex-ai/upload", summary="上传文件到Vertex AI可访问的Google cloud storage", )
async def upload_file(bucket: str, file: UploadFile,
                      headers: Annotated[GoogleAuthHeaders, Header()],
                      prefix: Optional[str] = None) -> GoogleUploadResponse:
    key = f"{prefix}/{file.filename}" if prefix else file.filename
    return await GoogleAuthUtils.google_upload_file(file_stream=file.file,
                                                    content_type=file.content_type,
                                                    google_api_key=headers.auth_token,
                                                    bucket_name=bucket, filename=key)


@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.delete('/delete_all', summary="删除所有已上传的文件/第一页")
async def delete_all(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(timeout=60) as client:
        resp = client.get(f"https://generativelanguage.googleapis.com/v1beta/files?key={google_api_key}")
        if resp.status_code != 200:
            return JSONResponse(content={"msg": "列出文件失败"}, status_code=resp.status_code)
        filename_list = [i["name"].split("/")[1] for i in resp.json()["files"]]
        for filename in filename_list:
            logger.info(f"删除 {filename}...")
            response = client.delete(
                url=f"https://generativelanguage.googleapis.com/v1beta/files/{filename}?key={google_api_key}")
            response.raise_for_status()
            if resp.status_code != 200:
                return JSONResponse(content={"msg": f"删除文件{filename}失败"}, status_code=resp.status_code)
    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(timeout=60) as client:
        response = client.get(
            url=f"https://generativelanguage.googleapis.com/v1beta/files?key={google_api_key}")
    if response.status_code != 200:
        return JSONResponse(content={}, status_code=response.status_code)
    return JSONResponse(content=response.json() if len(response.text) > 0 else "", status_code=response.status_code)


@router.get('/access-token', summary="获取一个新的Google Access Token",
            description="获取的Access Token最多同时存在35000个",
            dependencies=[Depends(verify_token)])
async def get_access_token() -> GoogleAuthUtils.GoogleAuthResponse:
    # GOOGLE_AUTH_JSON 通过modal secrets 配置
    service_account_json = os.environ.get("GOOGLE_AUTH_JSON", None)
    if not service_account_json:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google IAM Service Account JSON")
    service_account_info = json.loads(service_account_json)
    access_token = await GoogleAuthUtils.get_google_auth_jwt(service_account_info=service_account_info,
                                                             scopes=['https://www.googleapis.com/auth/cloud-platform'])
    return access_token


@router.post("/make_grid_gemini", summary="将输入图拼为网格上传到Gemini网盘")
async def make_grid_gemini_upload(data: MakeGridGeminiRequest, headers: Annotated[BundleHeaders, 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")
    fn = modal.Function.from_name(config.modal_app_name, "make_image_grid_upload",
                                  environment_name=config.modal_environment)
    image_grid_uri = await fn.remote.aio(data.pic_info_list, data.image_size, data.text_height, data.font_size,
                                         data.padding, data.separator, google_api_key,
                                         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 JSONResponse(content={"uri": image_grid_uri}, status_code=200)


@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_cover_grid_uri_list, data.product_list,
                       data.start_time, data.end_time, data.options,
                       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),
                       data.webhook, 3, data.scale, data.last_product_text)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post('/inference_gemini_first_only', summary="使用Gemini推理hls视频流指定时间段的打点情况-仅第一阶段")
async def inference_gemini_first_only(data: GeminiRequestFirstOnly,
                                      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_inference_first_only",
                                  environment_name=config.modal_environment)
    fn_call = fn.spawn(google_api_key, data.product_cover_grid_uri_list, data.product_list)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post('/inference_gemini_second_only', summary="使用Gemini推理hls视频流指定时间段的打点情况-仅第二阶段")
async def inference_gemini_second_only(data: GeminiRequestSecondOnly,
                                       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_inference_second_only",
                                  environment_name=config.modal_environment)
    fn_call = fn.spawn(data.media_hls_url, google_api_key, data.identified_product_list, data.start_time,
                       data.end_time, data.scale, data.options, data.prompt_label, data.last_product_text)
    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, ensure_ascii=False)
                                    if task_info.results is not None 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="")
