from __future__ import annotations
from typing import TYPE_CHECKING

from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, TypeAlias, Any
from mindor.dsl.schema.component import ModelComponentConfig
from mindor.dsl.schema.action import ModelActionConfig, SdxlImageGenerationModelActionConfig
from mindor.core.logger import logging
from ....base import ComponentActionContext
from ..common import ImageGenerationTaskService, ImageGenerationTaskAction
from PIL import Image as PILImage
import asyncio

if TYPE_CHECKING:
    import torch

class SdxlImageGenerationTaskAction(ImageGenerationTaskAction):
    def __init__(
        self,
        config: SdxlImageGenerationModelActionConfig,
        pipeline: Any,
        device: Optional[torch.device]
    ):
        super().__init__(config, device)

        self.config: SdxlImageGenerationModelActionConfig = config
        self.pipeline = pipeline

    async def _generate(self, params: Dict[str, Any]) -> List[PILImage.Image]:
        pass

class SdxlImageGenerationTaskService(ImageGenerationTaskService):
    def __init__(self, id: str, config: ModelComponentConfig, daemon: bool):
        super().__init__(id, config, daemon)

        self.pipeline: Optional[Any] = None
        self.device: Optional[torch.device] = None

    def get_setup_requirements(self) -> Optional[List[str]]:
        return None

    def _load_model(self) -> None:
        self.pipeline, self.device = self._load_pretrained_pipeline()

    def _unload_model(self) -> None:
        self.pipeline = None
        self.device = None

    def _load_pretrained_pipeline(self) -> Tuple[Any, torch.device]:
        pass

    async def _run(
        self,
        action: ModelActionConfig,
        context: ComponentActionContext,
        loop: asyncio.AbstractEventLoop
    ) -> Any:
        return SdxlImageGenerationModelActionConfig(action, self.model).run(context)
