from typing import Optional, Any, Union
from functools import wraps
import backoff
import httpx
import asyncio
from loguru import logger
import aiofiles
from pathlib import Path


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)