import asyncio
import gc
import re
import time
from typing import Any

import torch # type: ignore
from diffusers import StableDiffusionPipeline, FluxPipeline # type: ignore

class ImageModelCache: # Stay loaded for 3 hours
    def __init__(self, timeout_seconds: float = 3 * 60 * 60):
        self._pipe = None
        self._model_name = None
        self._device = None
        self._last_access_time = 0
        self._timeout_seconds = timeout_seconds
        self._lock = asyncio.Lock()
        self._cleanup_task = None

    async def start(self):
        if self._cleanup_task is None:
            self._cleanup_task = asyncio.create_task(self._periodic_cleanup())

    def _get_model_type(self, model_name: str) -> str:
        if re.search(r"stable-diffusion", model_name, re.IGNORECASE):
            return "stable"
        return "flux"

    async def get_pipeline(self, model: str, device: str) -> Any:
        await self.start()  # Ensure cleanup task starts on first use

        async with self._lock:
            current_time = time.time()

            current_model_type = self._get_model_type(model)
            cached_model_type = self._get_model_type(self._model_name) if self._model_name else None

            if (
                self._pipe is not None and
                self._model_name == model and
                self._device == device and
                current_model_type == cached_model_type and
                current_time - self._last_access_time < self._timeout_seconds
            ):
                self._last_access_time = current_time
                return self._pipe

            await self._unload_model()

            if current_model_type == "stable":
                pipe = StableDiffusionPipeline.from_pretrained(
                    model,
                    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
                )
                def dummy_safety_checker(images, clip_input):
                    return images, [False] * len(images)
                pipe.safety_checker = dummy_safety_checker
            else:
                pipe = FluxPipeline.from_pretrained(
                    model,
                    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
                )
                try:
                    pipe.load_lora_weights('enhanceaiteam/Flux-uncensored', weight_name='lora.safetensors')
                except Exception as e:
                    raise Exception(f"Failed to load LoRA weights: {str(e)}")

            pipe = pipe.to(device)
            self._pipe = pipe
            self._model_name = model
            self._device = device
            self._last_access_time = current_time
            return pipe

    async def _unload_model(self):
        if self._pipe is not None:
            try:
                del self._pipe
                gc.collect()
                if self._device == "cuda":
                    torch.cuda.empty_cache()
                elif self._device == "xpu":
                    torch.xpu.empty_cache()
            except Exception:
                pass
            self._pipe = None
            self._model_name = None
            self._device = None

    async def cleanup_if_expired(self):
        async with self._lock:
            if (
                self._pipe is not None and
                time.time() - self._last_access_time >= self._timeout_seconds
            ):
                await self._unload_model()

    async def _periodic_cleanup(self):
        while True:
            await asyncio.sleep(self._timeout_seconds)
            await self.cleanup_if_expired()
