import random

from typing import Union, Dict, BinaryIO, List, Optional, Any
from urllib.parse import quote

import backoff
import httpx
import asyncio

from google.genai.types import GenerateContentResponse, ContentUnion
from loguru import logger
import aiofiles
from pathlib import Path
from pydantic import BaseModel, computed_field, Field, field_validator
from ..models.responses.models import GoogleUploadResponse


class HTTPDownloadUtils:
    # 创建一个类级别的 Semaphore 来控制并发
    _semaphore = asyncio.Semaphore(5)  # 限制最大并发数为5

    @staticmethod
    @backoff.on_exception(
        backoff.expo,
        (httpx.RequestError, httpx.HTTPStatusError, Exception),
        max_tries=10,
        max_time=300,  # 最大重试时间5分钟
        giveup=lambda e: isinstance(e, httpx.HTTPStatusError) and e.response.status_code == 404
    )
    async def async_download_file(
            url: str,
            output_path: Union[str, Path],
            timeout: float = 60.0,
            verify: bool = True
    ) -> Union[str, None]:
        """
        异步下载文件或获取内容

        :param url: 要下载的URL
        :param output_path: 输出文件路径，如果为None则返回内容
        :param timeout: 请求超时时间（秒）
        :param verify: 是否验证SSL证书

        :return: 如果output_path为None，返回下载的内容；否则返回保存的文件路径

        :exception httpx.RequestError: 请求错误
        :exception httpx.HTTPStatusError: HTTP状态错误
        :exception Exception: 其他错误
        """
        async with HTTPDownloadUtils._semaphore:  # 使用信号量控制并发
            try:
                logger.info(f"Starting download from {url}")

                async with httpx.AsyncClient(
                        timeout=timeout,
                        verify=verify,
                        follow_redirects=True
                ) as client:
                    response = await client.get(url)
                    response.raise_for_status()  # 检查HTTP状态码

                    output_path = Path(output_path)
                    output_path.parent.mkdir(parents=True, exist_ok=True)

                    async with aiofiles.open(output_path, 'wb') as f:
                        await f.write(response.content)
                    logger.success(f"Successfully downloaded to {output_path}")
                    return str(output_path)


            except httpx.HTTPStatusError as e:
                logger.error(f"HTTP error occurred: {e.response.status_code} - {url}")
                raise
            except httpx.RequestError as e:
                logger.error(f"Request error occurred: {str(e)} - {url}")
                raise
            except Exception as e:
                logger.error(f"Unexpected error occurred: {str(e)} - {url}")
                raise

    @staticmethod
    async def batch_download(
            urls: list[str],
            output_paths: list[Union[str, Path]] = None,
            max_concurrent: int = 5
    ) -> list[Union[str, None]]:
        """
        批量下载多个文件

        :param urls: URL列表
        :param output_paths: 输出路径列表
        :param max_concurrent: 最大并发数

        :return: 下载结果列表
        """
        # 更新信号量的值
        HTTPDownloadUtils._semaphore = asyncio.Semaphore(max_concurrent)

        tasks = []
        for i, url in enumerate(urls):
            output_path = output_paths[i]
            task = asyncio.create_task(
                HTTPDownloadUtils.async_download_file(url, output_path)
            )
            tasks.append(task)

        return await asyncio.gather(*tasks, return_exceptions=True)


from google.oauth2 import service_account
from google.genai import types


class GoogleAuthUtils:
    class GoogleAuthResponse(BaseModel):
        access_token: str
        expires_in: int
        token_type: str

    class VertexAIRequestModel(types._common.BaseModel):
        contents: Union[types.ContentListUnion, types.ContentListUnionDict] = Field()
        generation_config: Optional[types.GenerateContentConfig] = Field(default=None, alias="generationConfig",
                                                                         serialization_alias="generationConfig",
                                                                         description="如果在JSON序列化时报MetadataClass")

        @field_validator("generation_config", mode="before")
        def validate_generation_config(cls, v):
            if not v:
                return None
            elif isinstance(v, types.GenerateContentConfig):
                if not v.response_schema:
                    return v
                if issubclass(v.response_schema, BaseModel):
                    flat_schema = v.response_schema.model_json_schema(schema_generator=FlatJsonSchemaGenerator)
                    v.response_schema = flat_schema
                    return v
                elif isinstance(v.response_schema, Dict):
                    return v
            raise TypeError("generation_config must be of type 'GenerateContentConfig'")

        @computed_field(alias="safetySettings")
        @property
        def safety_settings(self) -> Optional[List[types.SafetySetting]]:
            return self.generation_config.safety_settings

        @computed_field(alias="systemInstruction")
        @property
        def system_instruction(self) -> Optional[ContentUnion]:
            return self.generation_config.system_instruction

    class GoogleGenaiClient(BaseModel):
        cloudflare_project_id: str
        cloudflare_gateway_id: str
        google_project_id: str
        regions: List[str]
        access_token: str

        @computed_field(description="通过Cloudflare gateway调用的url")
        @property
        def gateway_url(self) -> str:
            return f"https://gateway.ai.cloudflare.com/v1/{self.cloudflare_project_id}/{self.cloudflare_gateway_id}/google-vertex-ai/v1/projects/{self.google_project_id}/locations/{random.choice(self.regions)}/publishers/google/models"

        def generate_content(self, model_id: str, contents: Union[types.ContentListUnion, types.ContentListUnionDict],
                             config: Optional[types.GenerateContentConfig] = None,
                             timeout: int = 30) -> tuple[dict[Any, Any], int] | tuple[GenerateContentResponse, int]:
            parameter_model = GoogleAuthUtils.VertexAIRequestModel(contents=contents, generationConfig=config)
            json_body = parameter_model.model_dump_json(indent=2, exclude_none=True, by_alias=True,
                                                        exclude={"generation_config": ["safety_settings", "system_instruction"]})
            logger.info(json_body)
            url = f"{self.gateway_url}/{model_id}:generateContent"
            logger.info(f"Authorization : Bearer {self.access_token}")
            logger.info(f"POST {url}")
            response = httpx.post(url=url, timeout=timeout,
                                  headers={"Content-Type": "application/json",
                                           "Authorization": f"Bearer {self.access_token}"},
                                  content=json_body)
            # response.raise_for_status()
            # logger.info(response.text)
            if response.status_code != 200:
                return {}, response.status_code
            return types.GenerateContentResponse.model_validate_json(response.text), response.status_code

    @staticmethod
    async def get_google_auth_jwt(service_account_info: Dict[str, str], scopes: list[str]) -> GoogleAuthResponse:
        credentials: service_account.Credentials = service_account.Credentials.from_service_account_info(
            service_account_info, scopes=scopes)
        assertion = credentials._make_authorization_grant_assertion().decode('utf-8')
        logger.info(f"jwt_assertion: {assertion}")
        params = {
            'grant_type': "urn:ietf:params:oauth:grant-type:jwt-bearer",
            'assertion': assertion
        }
        response = httpx.post(url="https://oauth2.googleapis.com/token",
                              headers={'Content-Type': 'application/x-www-form-urlencoded'},
                              data=params)
        response.raise_for_status()
        return GoogleAuthUtils.GoogleAuthResponse.model_validate_json(response.text)

    @staticmethod
    async def google_upload_file(file_stream: Union[BinaryIO, bytes], content_type: str, google_api_key: str,
                                 bucket_name: str,
                                 filename: str) -> GoogleUploadResponse:
        safe_filename = quote(filename)

        upload_response = httpx.post(
            url=f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={safe_filename}",
            content=file_stream,
            headers={
                "Authorization": f"Bearer {google_api_key}",
                "Content-Type": content_type
            }, timeout=900)
        upload_response.raise_for_status()

        # upload_url = f"gs://dy-media-storage/video/{filename}"
        return GoogleUploadResponse.model_validate_json(upload_response.text)


from pydantic.json_schema import GenerateJsonSchema, JsonSchemaMode
from pydantic import BaseModel, Field


class FlatJsonSchemaGenerator(GenerateJsonSchema):
    def generate(self, schema, mode: JsonSchemaMode = 'validation'):
        # 先用原始生成器生成 schema
        schema_dict = super().generate(schema, mode=mode)
        # 展开 $ref
        return self._flatten_refs(schema_dict)

    def _flatten_refs(self, schema, defs=None):
        if defs is None:
            defs = schema.get('$defs', {})
        if isinstance(schema, dict):
            if '$ref' in schema:
                ref_path = schema['$ref']
                # 只支持本地 $defs 引用
                if ref_path.startswith('#/$defs/'):
                    def_name = ref_path.split('/')[-1]
                    return self._flatten_refs(defs[def_name], defs)
                else:
                    raise ValueError(f'Unsupported $ref: {ref_path}')
            else:
                return {k: self._flatten_refs(v, defs) for k, v in schema.items() if k != '$defs'}
        elif isinstance(schema, list):
            return [self._flatten_refs(i, defs) for i in schema]
        else:
            return schema
