import asyncio
import hashlib
from pathlib import Path
from typing import ClassVar

import aiohttp

from yuxi.utils.logging_config import logger


class MediaPolicy:
    CACHE_DIR: ClassVar[Path] = Path("/tmp/yuxi_media_cache")
    SUPPORTED_IMAGE_FORMATS: ClassVar[frozenset[str]] = frozenset({"png", "jpg", "jpeg", "gif", "webp", "bmp"})

    def __init__(self):
        self._download_locks: dict[str, asyncio.Lock] = {}

    def _get_download_lock(self, url: str) -> asyncio.Lock:
        if url not in self._download_locks:
            self._download_locks[url] = asyncio.Lock()
        return self._download_locks[url]

    def validate_size(self, data: bytes, max_size_mb: int = 100) -> bool:
        limit = max_size_mb * 1024 * 1024
        return len(data) <= limit

    def detect_format(self, data: bytes) -> str | None:
        if data[:4] == b"\x89PNG":
            return "png"
        if data[:2] == b"\xff\xd8":
            return "jpg"
        if data[:6] in (b"GIF87a", b"GIF89a"):
            return "gif"
        if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
            return "webp"
        if data[:2] == b"BM":
            return "bmp"
        return None

    def is_supported_format(self, fmt: str) -> bool:
        return fmt in self.SUPPORTED_IMAGE_FORMATS

    async def cache_url_content(self, url: str) -> bytes | None:
        lock = self._get_download_lock(url)
        async with lock:
            cache_key = hashlib.md5(url.encode()).hexdigest()[:12]
            cache_file = self.CACHE_DIR / cache_key
            if cache_file.exists():
                return cache_file.read_bytes()

            self.CACHE_DIR.mkdir(parents=True, exist_ok=True)
            try:
                async with aiohttp.ClientSession() as session:
                    async with session.get(url, timeout=aiohttp.ClientTimeout(total=30)) as resp:
                        if resp.status != 200:
                            return None
                        data = await resp.read()
                cache_file.write_bytes(data)
                return data
            except Exception as e:
                logger.error(f"Failed to download media URL: {e}")
                return None