"""Variable system for MMIRAGE pipeline with multimodal support."""

from __future__ import annotations

import abc
import os
from dataclasses import dataclass
from types import MappingProxyType
from typing import Any, Dict, List, Literal, Optional, Sequence

from jmespath import search
from PIL import Image


@dataclass
class BaseVar(abc.ABC):
    """Base class for variables in the MMIRAGE pipeline.

    Attributes:
        name: Name of the variable.
    """
    name: str = ""


@dataclass
class InputVar(BaseVar):
    """Input variable extracted from source datasets.

    Attributes:
        key: JMESPath query to extract the variable from a sample.
        type: Variable type - "text" or "image".
    """
    key: str = ""
    type: Literal["text", "image"] = "text"

    def is_image(self) -> bool:
        """Check if this input variable represents an image.

        Returns:
            bool: True if the variable type is "image", False otherwise.
        """
        return self.type == "image"


@dataclass
class OutputVar(BaseVar):
    """Output variable generated by processors.

    Output variables are created by processors (e.g., LLMs) and can
    depend on input variables and previously computed output variables.

    Attributes:
        name: Name of the variable.
        type: Type identifier for the processor that generates this variable.
    """
    type: str = ""

    @abc.abstractmethod
    def is_computable(self, vars: Sequence[BaseVar]) -> bool:
        raise NotImplementedError()


def _to_abs_path(path: str) -> str:
    """Normalize a local filesystem path to an absolute path."""
    return os.path.abspath(path)


def _resolve_image_input(value: Any, image_base_path: Optional[str] = None) -> Any:
    """Resolve image input to a format SGLang can use.

    In your SGLang setup, local images should be passed as plain filesystem paths
    (NOT file:// URIs). This function:
      - passes through PIL images and http(s) URLs
      - normalizes local paths to absolute paths
      - resolves relative paths against image_base_path (preferred) or CWD (fallback)
      - strips a leading file:// if it is present (for backward compatibility)

    Args:
        value: The image value to resolve (string path/URL, PIL image, etc.).
        image_base_path: Optional base directory for resolving relative image paths.

    Returns:
        A resolved image spec: PIL.Image.Image, http(s) URL string, or absolute path string.

    Raises:
        FileNotFoundError: If the resolved file does not exist.
        RuntimeError: If the resolved path exists but is not a file.
    """
    # Already a PIL Image - pass through
    if isinstance(value, Image.Image):
        return value

    # Not a string - pass through
    if not isinstance(value, str):
        return value

    # HTTP(S) URL - pass through
    if value.startswith(("http://", "https://")):
        return value

    # If value is a file:// URI, strip it to a local path
    if value.startswith("file://"):
        value = value[len("file://") :]

    # Absolute local path
    if os.path.isabs(value):
        path = os.path.realpath(value) if os.path.islink(value) else value
        path = _to_abs_path(path)
        if not os.path.exists(path):
            raise FileNotFoundError(f"Absolute image path does not exist: {path}")
        if not os.path.isfile(path):
            raise RuntimeError(f"The provided path exists but is not a file: {path}")
        return path

    # Relative path: resolve against base path if provided
    if image_base_path:
        resolved_path = _to_abs_path(os.path.join(image_base_path, value))
        if not os.path.exists(resolved_path):
            raise FileNotFoundError(
                f"Resolved image path '{resolved_path}' does not exist "
                f"(base='{image_base_path}', rel='{value}')."
            )
        if not os.path.isfile(resolved_path):
            raise RuntimeError(f"Resolved image path is not a file: {resolved_path}")
        return resolved_path

    # Fallback: resolve against current working directory
    cwd_path = _to_abs_path(value)
    if os.path.exists(cwd_path):
        if not os.path.isfile(cwd_path):
            raise RuntimeError(f"Relative image path exists but is not a file: {cwd_path}")
        return cwd_path

    raise FileNotFoundError(
        f"Relative image path '{value}' cannot be resolved. "
        "Set image_base_path in the dataset configuration or use an absolute path."
    )


class VariableEnvironment:
    """Environment for storing and accessing variables during processing."""

    def __init__(self, var_env: Dict[str, Any], image_vars: Optional[set] = None) -> None:
        """Initialize a variable environment.

        Args:
            var_env: Dictionary mapping variable names to their values.
            image_vars: Set of variable names that represent images. Defaults to empty set.
        """
        self._vars_env = var_env
        self._image_vars = image_vars or set()

    def with_variable(self, key: str, value: Any, is_image: bool = False) -> "VariableEnvironment":
        """Create a new environment with an additional variable.

        Args:
            key: Name of the variable to add.
            value: Value of the variable.
            is_image: Whether the variable represents an image.

        Returns:
            VariableEnvironment: New environment with the added variable.
        """
        new_image_vars = self._image_vars.copy()
        if is_image:
            new_image_vars.add(key)
        return VariableEnvironment(self._vars_env | {key: value}, new_image_vars)

    def to_dict(self) -> MappingProxyType:
        """Get an immutable view of the variable dictionary.

        Returns:
            MappingProxyType providing read-only access to variables.
        """
        return MappingProxyType(self._vars_env)

    def get(self, key: str, default: Any = None) -> Any:
        """Get a variable value by name.

        Args:
            key: Name of the variable to retrieve.
            default: Default value to return if variable is not found.

        Returns:
            Variable value, or default if not found.
        """
        return self._vars_env.get(key, default)

    def is_image_var(self, key: str) -> bool:
        """Check if a variable represents an image.

        Args:
            key: Name of the variable to check.

        Returns:
            bool: True if the variable is an image variable, False otherwise.
        """
        return key in self._image_vars

    def get_image_vars(self) -> set:
        """Get all image variable names.

        Returns:
            set: Copy of the set containing names of all image variables.
        """
        return self._image_vars.copy()

    def get_images(self) -> List[Any]:
        """Get image values in a deterministic order."""
        return [self._vars_env[k] for k in sorted(self._image_vars) if k in self._vars_env]

    def has_images(self) -> bool:
        """Check if the environment contains any image variables.

        Returns:
            bool: True if at least one image variable is present, False otherwise.
        """
        return any(k in self._vars_env for k in self._image_vars)

    @staticmethod
    def from_input_variables(sample: Dict[str, Any], input_vars: List[InputVar], image_base_path: Optional[str] = None) -> "VariableEnvironment":
        """Create a variable environment from a single sample.

        Args:
            sample: Dictionary containing the data for one sample.
            input_vars: List of input variable definitions to extract.
            image_base_path: Optional base directory for resolving relative image paths.

        Returns:
            VariableEnvironment: Environment populated with extracted variables.

        Raises:
            ValueError: If a required input variable is not found in the sample.
        """
        ret: Dict[str, Any] = {}
        image_vars: set = set()

        for input_var in input_vars:
            value = search(input_var.key, sample)
            if value is None:
                raise ValueError(
                    f"Input variable '{input_var.name}' with key '{input_var.key}' not found in the sample."
                )

            if input_var.is_image():
                value = _resolve_image_input(value, image_base_path)
                image_vars.add(input_var.name)

            ret[input_var.name] = value

        return VariableEnvironment(ret, image_vars)

    @staticmethod
    def from_batch_input_variables(
        batch: Dict[str, List[Any]], input_vars: List[InputVar], image_base_path: Optional[str] = None
    ) -> List["VariableEnvironment"]:
        """Extract input variables from a batch of samples.

        Args:
            batch: Dictionary mapping column names to lists of values.
            input_vars: List of input variable definitions.
            image_base_path: Optional base directory for resolving relative image paths.

        Returns:
            List of VariableEnvironments, one for each sample in the batch.
        """
        vars_samples: List[VariableEnvironment] = []

        batch_size = len(next(iter(batch.values())))
        batch_list: List[Dict[str, Any]] = [
            {k: batch[k][i] for k in batch.keys()} for i in range(batch_size)
        ]

        for sample in batch_list:
            vars_samples.append(VariableEnvironment.from_input_variables(sample, input_vars, image_base_path))

        return vars_samples
