"""Variable system for MMIRAGE pipeline."""

from __future__ import annotations

import abc
from dataclasses import dataclass
from typing import List, Dict, Any, Sequence, TypeVar
from types import MappingProxyType
from jmespath import search


@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.

    Input variables are extracted from source dataset samples using
    JMESPath queries to support flexible data access.

    Attributes:
        name: Name of the variable.
        key: JMESPath query to extract the variable from a sample.
    """

    key: str = ""


@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:
        """Check if this output variable can be computed given available variables.

        Args:
            vars: Sequence of currently available variables.

        Returns:
            True if the variable can be computed, False otherwise.
        """
        ...


C = TypeVar("C")


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

    Provides an immutable container for variable values, supporting
    creation of new environments with additional variables.

    Attributes:
        _vars_env: Internal dictionary storing variable values.
    """

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

        Args:
            var_env: Dictionary mapping variable names to their values.
        """
        self._vars_env = var_env

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

        Args:
            key: Name of the variable to add.
            value: Value of the variable.

        Returns:
            A new VariableEnvironment with the added variable.
        """
        return VariableEnvironment(self._vars_env | {key: value})

    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)

    @staticmethod
    def from_input_variables(
        sample: Dict[str, Any], input_vars: List[InputVar]
    ) -> VariableEnvironment:
        """Extract input variables from a dataset sample using JMESPath queries.

        Args:
            sample: Dictionary representing a dataset sample.
            input_vars: List of input variable definitions.

        Returns:
            A VariableEnvironment containing the extracted variables.

        Raises:
            ValueError: If an input variable cannot be found in the sample.
        """
        ret: Dict[str, Any] = {}
        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."
                )
            ret[input_var.name] = value

        return VariableEnvironment(ret)

    @staticmethod
    def from_batch_input_variables(
        batch: Dict[str, List[Any]], input_vars: List[InputVar]
    ) -> 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.

        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:
            current_vars = VariableEnvironment.from_input_variables(sample, input_vars)
            vars_samples.append(current_vars)

        return vars_samples
