# Copyright contributors to the Terratorch project

import torch
import torch.nn as nn
from terratorch.registry import MODEL_FACTORY_REGISTRY
from typing import Optional

from pydantic import BaseModel
from typing import List, Dict
from enum import Enum

from packaging import version
from vllm import __version__ as vllm_version


def check_vllm_version(target_version: str, comparison: str):
    current_version = version.parse(vllm_version)
    target = version.parse(target_version)

    if comparison == "==":
        return current_version == target
    elif comparison == "!=":
        return current_version != target
    elif comparison == "<":
        return current_version < target
    elif comparison == "<=":
        return current_version <= target
    elif comparison == ">":
        return current_version > target
    elif comparison == ">=":
        return current_version >= target
    else:
        raise ValueError(f"Invalid comparison operator: {comparison}")


class InputTypeEnum(str, Enum):
    tensor = "torch.Tensor"


class data(BaseModel):
    type: InputTypeEnum
    shape: List[int]


class InputDefinition(BaseModel):
    data: Dict[str, data]
    target: Optional[str] = None


class MultiModalDataGenerator:
    def __init__(self, config: dict):
        # print("DummyDataGenerator called")
        # self.task_type= config["task_args"]["task"]
        self.config = config
        self.input_definition = InputDefinition(**config["input"])
        self.task_type = config["model"]["class_path"]

    def _get_mm_fields_config(self) -> dict[str:str]:
        fields = {}
        for input_name, input in self.input_definition.data.items():
            if input.type == InputTypeEnum.tensor:
                fields[input_name] = "image"
        return fields


class DummyDataGenerator:
    def __init__(self, config: dict):
        self.config = config
        self.input_definition = InputDefinition(**config["input"])
        self.task_type = config["model"]["class_path"]

    def get_dummy_mm_data(self) -> dict[str, torch.Tensor]:

        mm_data = {}
        for input_name, input in self.input_definition.data.items():
            if input.type == InputTypeEnum.tensor:
                mm_data[input_name] = torch.full(input.shape, 1.0, dtype=torch.float16)
        return mm_data


def lookup_task_name(class_path):
    if "SemanticSegmentationTask" in class_path:
        return "segmentation"
    elif "PixelwiseRegressionTask" in class_path:
        return "regression"
    else:
        return None


class InferenceRunner(nn.Module):
    def __init__(self, config: dict):

        super().__init__()
        model_conf = config["model"]
        self.task_type = config["model"]["class_path"]
        task_name = lookup_task_name(model_conf["class_path"])
        model_factory = MODEL_FACTORY_REGISTRY.build(model_conf["init_args"]["model_factory"])
        self.model = model_factory.build_model(task=task_name, **model_conf["init_args"]["model_args"])
        self.input_definition = InputDefinition(**config["input"])

    def _parse_and_validate_multimodal_data(self, **kwargs: object) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        mm_data = {}
        for input_name, input in self.input_definition.data.items():
            input_value = kwargs.pop(input_name, None)
            if input.type == InputTypeEnum.tensor:
                if not isinstance(input_value, torch.Tensor):
                    raise ValueError(
                        f"Incorrect type of {input_name}. Got type: {type(input_value)}, expected {input.type}"
                    )
                if self.input_definition.target and input_name != self.input_definition.target:
                    input_value = torch.unbind(input_value, dim=0)[0]

            mm_data[input_name] = input_value
        return mm_data

    def forward(self, **kwargs: object):

        input = self._parse_and_validate_multimodal_data(**kwargs)
        if self.input_definition.target:
            target_input = input.pop(self.input_definition.target)
            model_output = self.model(target_input, **input)
        else:
            model_output = self.model(input)
        return model_output
