import gc
import torch
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Union, Optional
from PIL.Image import Image
from src.config import DeviceConfig
from src.logger import get_logger


logger = get_logger(__name__)


class ImageTextGenerator(ABC):
    def __init__(self, model) -> None:
        self.model = model

    @abstractmethod
    def generate(
        self, samples: list[str], images: list[Union[Image, list[Image], None]]
    ) -> list[str]:
        pass

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        del self.model

        torch.cuda.empty_cache()
        gc.collect()


@dataclass
class ImageTextModelConfig(DeviceConfig):
    model_name_or_path: str = field(
        default="", metadata={"help": "Path to the Llava-Next model."}
    )
    max_seq_len: int = field(default=None, metadata={"help": "Maximum sequence length."})
    max_new_tokens: int = field(
        default=512, metadata={"help": "Maximum number of new tokens to generate."}
    )
    top_p: float = field(default=0.95, metadata={"help": "Top-p sampling parameter."})
    temperature: float = field(default=0.0, metadata={"help": "Temperature for sampling."})
    prompt_template: str = field(
        default="{text}", metadata={"help": "Prompt template for generation."}
    )
    system_prompt: Optional[str] = field(
        default=None, metadata={"help": "System prompt to prepend to inputs."}
    )
    batch_size: Optional[int] = field(
        default=None, metadata={"help": "Batch size for generation if used."}
    )

    def __post_init__(self):
        if self.max_seq_len is None:
            if self.model_name_or_path.endswith("llava-v1.6-vicuna-13b-hf"):
                self.max_seq_len = 4096
            elif self.model_name_or_path.startswith("google/paligemma2"):
                self.max_seq_len = 8192
            else:
                self.max_seq_len = 16384

    @classmethod
    def from_yaml(cls, path: str) -> "ImageTextModelConfig":
        import yaml

        with open(path, "r") as file:
            config_data = yaml.safe_load(file)
        return cls(**config_data)


class AutoImageTextGenerator:
    @staticmethod
    def from_config(config: ImageTextModelConfig) -> ImageTextGenerator:
        if "pixtral" in config.model_name_or_path.lower():
            from src.modeling.mistral.predict import MistralGenerator

            logger.info("Using MistralGenerator for model: %s", config.model_name_or_path)
            return MistralGenerator(config)

        if "paligemma" in config.model_name_or_path.lower():
            logger.info("Using PaligemmaGenerator for model: %s", config.model_name_or_path)
            from src.modeling.gemma.predict import PaligemmaGenerator

            return PaligemmaGenerator(config)

        if "gemma-3" in config.model_name_or_path.lower():
            logger.info("Using GemmaGenerator for model: %s", config.model_name_or_path)
            from src.modeling.gemma.predict import GemmaGenerator

            return GemmaGenerator(config)

        logger.info("Using LlavaNextGenerator for model: %s", config.model_name_or_path)
        from src.modeling.llava.predict import LlavaNextGenerator

        return LlavaNextGenerator(config)
