import asyncio
from dataclasses import dataclass
from hashlib import md5
import time

from nonebot import logger
from nonebot_plugin_apscheduler import driver
from konabot.common.path import DATA_PATH
import os
from PIL import Image
from io import BytesIO

IMAGE_PATH = DATA_PATH / "temp" / "images"

@dataclass
class ImageResource:
    filename: str
    expire: int

@dataclass
class StorageImage:
    name: str
    resources: dict[str,
                    dict[str,ImageResource]] # {群号: {QQ号: ImageResource}}

class ImageStorager:
    images_pool: dict[str,StorageImage] = {}

    max_storage: int = 10 * 1024 * 1024  # 最大存储10MB
    max_image_count: int = 200  # 最大存储图片数量

    @staticmethod
    def init():
        if not IMAGE_PATH.exists():
            IMAGE_PATH.mkdir(parents=True, exist_ok=True)

    @staticmethod
    def delete_path_image(name: str):
        resource_path = IMAGE_PATH / name
        if resource_path.exists():
            os.remove(resource_path)

    @staticmethod
    async def clear_all_image():
        # 清理 temp 目录下的所有图片资源
        for file in os.listdir(IMAGE_PATH):
            file_path = IMAGE_PATH / file
            if file_path.is_file():
                os.remove(file_path)

    @classmethod
    async def clear_expire_image(cls):
        # 清理过期的图片资源，将未被删除的放入列表中，如果超过最大数量则删除最早过期的
        remaining_images = []
        current_time = time.time()
        for name, storage_image in list(ImageStorager.images_pool.items()):
            for group_id, resources in list(storage_image.resources.items()):
                for qq_id, resource in list(resources.items()):
                    if resource.expire < current_time:
                        del storage_image.resources[group_id][qq_id]
                        cls.delete_path_image(name)
                    else:
                        remaining_images.append((name, group_id, qq_id, resource.expire))
            if not storage_image.resources:
                del ImageStorager.images_pool[name]
        # 如果剩余图片超过最大数量，按过期时间排序并删除最早过期的
        if len(remaining_images) > ImageStorager.max_image_count:
            remaining_images.sort(key=lambda x: x[3])  # 按过期时间排序
            to_delete = len(remaining_images) - ImageStorager.max_image_count
            for i in range(to_delete):
                name, group_id, qq_id, _ = remaining_images[i]
                resource = ImageStorager.images_pool[name].resources[group_id][qq_id]
                del ImageStorager.images_pool[name].resources[group_id][qq_id]
                cls.delete_path_image(name)
        logger.info("过期图片清理完成")

    @classmethod
    def _add_to_pool(cls, filename: str, name: str, group_id: str, qq_id: str, expire: int = 36000):
        expire_time = time.time() + expire
        if name not in cls.images_pool:
            cls.images_pool[name] = StorageImage(name=name,resources={})
        if group_id not in cls.images_pool[name].resources:
            cls.images_pool[name].resources[group_id] = {}
        cls.images_pool[name].resources[group_id][qq_id] = ImageResource(filename=filename, expire=expire_time)
        logger.debug(f"{cls.images_pool}")

    @classmethod
    def save_image(cls, image: bytes, name: str, group_id: str, qq_id: str) -> None:
        """
        以哈希值命名保存图片，并返回图片资源信息
        """
        # 检测图像大小，不得超过 10 MB
        if len(image) > cls.max_storage:
            raise ValueError("图片大小超过 10 MB 限制")
        hash_name = md5(image).hexdigest()
        ext = os.path.splitext(name)[1]
        file_name = f"{hash_name}{ext}"
        full_path = IMAGE_PATH / file_name
        with open(full_path, "wb") as f:
            f.write(image)
        # 将文件写入 images_pool
        logger.debug(f"Image saved: {file_name} for group {group_id}, qq {qq_id}")
        cls._add_to_pool(file_name, name, group_id, qq_id)

    @classmethod
    def save_image_by_pil(cls, image: Image.Image, name: str, group_id: str, qq_id: str) -> None:
        """
        以哈希值命名保存图片，并返回图片资源信息
        """
        img_byte_arr = BytesIO()
        # 如果图片是动图，保存为 GIF 格式
        if getattr(image, "is_animated", False):
            image.save(img_byte_arr, format="GIF", save_all=True, loop=0)
        else:
            image.save(img_byte_arr, format=image.format or "PNG")
        img_bytes = img_byte_arr.getvalue()
        cls.save_image(img_bytes, name, group_id, qq_id)

    @classmethod
    def load_image(cls, name: str, group_id: str, qq_id: str) -> Image:
        logger.debug(f"Loading image: {name} for group {group_id}, qq {qq_id}")
        if name not in cls.images_pool:
            logger.debug(f"Image {name} not found in pool")
            return None
        if group_id not in cls.images_pool[name].resources:
            logger.debug(f"No resources for group {group_id} in image {name}")
            return None
        # 寻找对应 QQ 号 的资源，如果没有就返回相同群下的第一个资源
        if qq_id not in cls.images_pool[name].resources[group_id]:
            first_qq_id = next(iter(cls.images_pool[name].resources[group_id]))
            qq_id = first_qq_id
        resource = cls.images_pool[name].resources[group_id][qq_id]
        resource_path = IMAGE_PATH / resource.filename
        logger.debug(f"Image path: {resource_path}")
        return Image.open(resource_path)

class ImageStoragerManager:
    def __init__(self, interval: int = 300):  # 默认 5 分钟执行一次
        self.interval = interval
        self._clear_task = None
        self._running = False
        
    async def start_auto_clear(self):
        """启动自动任务"""
        # 先清理一次
        await ImageStorager.clear_all_image()
        self._running = True
        self._clear_task = asyncio.create_task(self._auto_clear_loop())
            
        logger.info(f"自动清理任务已启动，间隔: {self.interval}秒")
    
    async def stop_auto_clear(self):
        """停止自动清理任务"""
        if self._clear_task:
            self._running = False
            self._clear_task.cancel()
            try:
                await self._clear_task
            except asyncio.CancelledError:
                pass
            logger.info("自动清理任务已停止")
        else:
            logger.warning("没有正在运行的自动清理任务")
        
    async def _auto_clear_loop(self):
        """自动清理循环"""
        while self._running:
            try:
                await asyncio.sleep(self.interval)
                await ImageStorager.clear_expire_image()
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"定时清理失败: {e}")

image_manager = ImageStoragerManager(interval=300)  # 每5分钟清理一次

@driver.on_startup
async def init_image_storage():
    ImageStorager.init()
    # 启用定时任务清理过期图片
    await image_manager.start_auto_clear()