# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2023-2024 @ CAMEL-AI.org. All Rights Reserved. =========

import random
from typing import Any, Dict, List, Optional, Tuple, Union

from camel.datasets import BaseGenerator, DataPoint, StaticDataset
from camel.logger import get_logger
from camel.verifiers.base import (
    BaseVerifier,
    VerificationResult,
)

from .models import Action, Observation, StepResult

logger = get_logger(__name__)


class SingleStepEnv:
    r"""A lightweight environment for single-step RL with LLMs as policy.

    This environment models a single interaction between an LLM-based agent
    and a problem drawn from a dataset—such as a question-answering or
    math problem—where the agent produces one response and receives feedback.

    Core Flow:
        - A question is sampled from a (possibly infinitely long) dataset.
        - The LLM generates a single-step response (the action).
        - The response is verified against the ground truth.
        - A reward is computed based on correctness and optional custom logic.

    Key Features:
        - Batched evaluation with per-sample state tracking.
        - Async setup and teardown for verifiers and related resources.
        - Supports deterministic sampling via local RNG (optional seed).
        - Extensible reward computation via subclassing.
    """

    PLACEHOLDER_OBS = Observation(
        question="Episode ended. This is just a placeholder."
    )

    ACCURACY_REWARD = 10

    def __init__(
        self,
        dataset: Union[StaticDataset, BaseGenerator],
        verifier: BaseVerifier,
        **kwargs,
    ) -> None:
        r"""Initialize the SingleStepEnv.

        Args:
            dataset (Union[StaticDataset, BaseGenerator]): Dataset to sample
                problems from.
            verifier (BaseVerifier): Verifier used to evaluate LLM responses
                against ground-truth answers.
            **kwargs: Optional metadata or configuration values.

        Notes:
            This class assumes all interactions are single-step: one question,
            one LLM response, one reward.
        """
        self.dataset = dataset
        self.verifier = verifier
        self._metadata = kwargs

        # State tracking
        self._is_setup: bool = False
        self._states: List[DataPoint] = []
        self._states_done: List[bool] = []
        self.current_batch_size: int = 0

    async def setup(self) -> None:
        r"""Set up the environment by initializing the verifier.

        This method ensures that the environment is ready for interaction.
        It sets up necessary components, including the verifier.

        Raises:
            Exception: If setup fails due to an internal error.
        """

        if self._is_setup:
            logger.warning("Environment has already been set up")
            return

        try:
            await self.verifier.setup()

            self._is_setup = True
            logger.info('Environment setup completed successfully')
        except Exception as e:
            logger.error(f'Failed to setup environment: {e}')
            raise

    async def close(self) -> None:
        r"""Clean up and close all resources used by the environment.

        This method shuts down the verifier, resets the internal
        state, and ensures that the environment is properly closed.

        Raises:
            Exception: If an error occurs while closing the environment.
        """

        if not self._is_setup:
            logger.warning(
                "Not closing environment - has not been set up yet."
            )
            return

        try:
            self._is_setup = False
            await self.verifier.cleanup()
            self._states = []
            self._states_done = []
            self.current_batch_size = 0
            logger.info('Environment closed successfully')
        except Exception as e:
            logger.error(f'Failed to close environment: {e}')
            raise

    async def reset(
        self, batch_size: int = 1, seed: Optional[int] = None
    ) -> Union[Observation, List[Observation]]:
        r"""Resets the environment and starts a new episode.

        This method samples a new batch of data points from the dataset and
        returns the corresponding initial observations.

        If a seed is provided, a local random number generator is initialized
        for deterministic sampling. The global random state is not affected.

        Args:
            batch_size (int): Number of data points to sample.
                (default: :obj:`1`)
            seed (Optional[int]): Seed for deterministic sampling. If None,
                sampling is non-deterministic. (default: :obj:`None`)

        Returns:
            Observation or List[Observation]: Initial observation(s) for the
                episode.

        Raises:
            RuntimeError: If called before all previous states are processed.
            ValueError: If batch size exceeds dataset size.
            TypeError: If the dataset is of an unsupported type.
        """
        if batch_size <= 0:
            raise ValueError("Batch size must be positive")

        if not self._is_setup:
            logger.warning(
                "reset() called on un-setup environment. Setting up..."
            )
            await self.setup()

        if self._batch_started() and not self._batch_done():
            logger.error(
                "Reset called before all states were processed. "
                "Call step on remaining states first."
            )
            raise RuntimeError(
                "reset() called before all states in batch were processed."
            )

        if seed is not None:
            rng = random.Random(seed)
        else:
            rng = random.Random()

        if isinstance(self.dataset, StaticDataset):
            dataset_len = len(self.dataset)

            if batch_size > dataset_len:
                raise ValueError(
                    f"Batch size {batch_size} is too large for dataset "
                    f"of size {dataset_len}"
                )

            start_idx = rng.randint(0, dataset_len - batch_size)
            idx_slice = slice(start_idx, start_idx + batch_size)
            val = self.dataset[idx_slice]
            self._states = [val] if isinstance(val, DataPoint) else val

            self.current_batch_size = len(self._states)
            self._states_done = [False] * self.current_batch_size

            observations = [
                Observation(question=sample.question, context={}, metadata={})
                for sample in self._states
            ]

            return observations[0] if batch_size == 1 else observations

        elif isinstance(self.dataset, BaseGenerator):
            raise NotImplementedError(
                "Reset not yet implemented for BaseGenerator datasets."
            )

        else:
            raise TypeError(f"Unsupported dataset type: {type(self.dataset)}")

    async def step(
        self, action: Union[Action, List[Action]]
    ) -> Union[
        Tuple[Observation, float, bool, Dict[str, Any]],
        List[Tuple[Observation, float, bool, Dict[str, Any]]],
    ]:
        r"""Process actions for a subset of states and update their finished
        status.

        Args:
            action: Single action (for batch_size=1 or micro-batch of size 1)
                or list of actions (for batch_size>=2 with multiple actions).
                Each action must have an index for batch_size>=2, indicating
                which state it corresponds to.

        Returns:
            Union[StepResult, List[StepResult]]: StepResult or list of
                StepResults for the processed states.

        Raises:
            RuntimeError: If environment isn't set up or episode has ended.
            ValueError: If indices are invalid, duplicate, or correspond to
                finished states.
        """
        if not self._is_setup:
            raise RuntimeError("Environment not set up. Call setup() first.")
        if self._batch_done():
            raise RuntimeError(
                "Episodes have ended for batch. Call reset() first."
            )
        if not self._states:
            raise RuntimeError("No current observation. Call reset() first.")

        # Normalize actions into a list for uniform processing
        if self.current_batch_size == 1:
            if isinstance(action, list):
                if len(action) != 1 or not isinstance(action[0], Action):
                    raise ValueError(
                        "For batch_size=1, expect a single Action or a "
                        "list containing exactly one Action"
                    )
            elif not isinstance(action, Action):
                raise ValueError(
                    "For batch_size=1, expect a single Action or a "
                    "list containing exactly one Action"
                )
            if isinstance(action, Action):
                actions = [action]
            else:
                actions = action
            if actions[0].index is None:
                actions[0].index = 0
            if actions[0].index != 0:
                raise ValueError("For batch_size=1, index must be None or 0")

        else:  # batch_size >= 2
            if isinstance(action, Action):
                if action.index is None:
                    raise ValueError(
                        "For batch_size>=2, each Action must have an index"
                    )
                if not isinstance(action.index, int):
                    raise ValueError("Index must be an integer")
                actions = [action]
            elif isinstance(action, list):
                if not action:  # Empty list
                    raise ValueError("Action list cannot be empty")
                actions = action
                for act in actions:
                    if not isinstance(act, Action):
                        raise ValueError(
                            "All elements in list must be Action objects"
                        )
                    if act.index is None:
                        raise ValueError(
                            "For batch_size>=2, each Action must have an index"
                        )
                    if not isinstance(act.index, int):
                        raise ValueError("Index must be an integer")
            else:
                raise ValueError(
                    "For batch_size>=2, expect an Action or list of Actions"
                )

        # Validate indices
        indices: List[int] = []
        for act in actions:
            assert act.index is not None
            indices.append(act.index)
        if len(set(indices)) != len(indices):
            raise ValueError("Duplicate state indices in actions.")
        for idx in indices:
            if idx < 0 or idx >= len(self._states):
                raise ValueError(f"Invalid state index {idx}.")
            if self._states_done[idx]:
                raise ValueError(f"State at index {idx} is already finished.")

        num_actions = len(actions)
        if self.current_batch_size % num_actions != 0:
            logger.warning(
                f"Number of actions ({num_actions}) is not a divisor of "
                f"total batch size ({self.current_batch_size})"
            )

        proposed_solutions = [act.llm_response for act in actions]
        ground_truths: List[str] = []
        for idx in indices:
            ground_truths.append(self._states[idx].final_answer)

        verification_results = await self.verifier.verify_batch(
            solutions=proposed_solutions,
            ground_truths=ground_truths,  # type: ignore [arg-type]
            raise_on_error=True,
        )

        total_rewards, rewards_dicts = await self._compute_reward_batch(
            proposed_solutions, verification_results
        )

        # TODO Batch this
        step_results = []
        for i, action in enumerate(actions):
            assert action.index is not None
            idx = action.index
            step_result = StepResult(
                observation=self.PLACEHOLDER_OBS,
                reward=total_rewards[i],
                rewards_dict=rewards_dicts[i],
                done=True,
                info={
                    "proposed_solution": proposed_solutions[i],
                    "verification_result": verification_results[i],
                    "state": self._states[idx],
                },
            )
            step_results.append(step_result.as_tuple())
            self._states_done[idx] = True

        return step_results[0] if len(step_results) == 1 else step_results

    async def _compute_reward_batch(
        self,
        proposed_solutions: List[str],
        verification_results: List[VerificationResult],
    ) -> Tuple[List[float], List[Dict[str, float]]]:
        r"""Compute rewards for a batch of proposed solutions based on
        verification results.

        Args:
            proposed_solutions (List[str]): List of LLM-generated responses to
                evaluate.
            verification_results (List[VerificationResult]): List of
                verification outcomes for each solution.

        Returns:
            Tuple containing:
                - List of total rewards for each solution.
                - List of reward component dictionaries for each solution.
        """
        if len(proposed_solutions) != len(verification_results):
            raise ValueError(
                f"Length mismatch: {len(proposed_solutions)} solutions vs "
                f"{len(verification_results)} verification results"
            )

        total_rewards = []
        rewards_dicts = []

        for solution, verification_result in zip(
            proposed_solutions, verification_results
        ):
            rewards: Dict[str, float] = {}

            rewards["correctness"] = (
                self.ACCURACY_REWARD if verification_result.status else 0.0
            )

            further_rewards = await self._compute_custom_reward(
                solution, verification_result
            )
            rewards = {**rewards, **further_rewards}

            total_reward = sum(rewards.values())
            total_rewards.append(total_reward)
            rewards_dicts.append(rewards)

        return total_rewards, rewards_dicts

    async def _compute_custom_reward(
        self, proposed_solution: str, verification_result: VerificationResult
    ) -> Dict[str, float]:
        r"""Compute additional custom reward components for a single solution.

        To be overridden by subclasses for domain-specific rewards.

        Args:
            proposed_solution (str): The LLM-generated response.
            verification_result (VerificationResult): The verification outcome.

        Returns:
            Dict[str, float]: Dictionary of custom reward components.
        """
        return {}

    def _batch_done(self) -> bool:
        r"""Check if all states in the current batch are done.

        Returns:
            bool: True if all states are marked as done, False otherwise.
        """
        return all(self._states_done)

    def _batch_started(self) -> bool:
        r"""Check if any state in the current batch is done.

        Returns:
            bool: True if at least one state is marked as done, False
                otherwise.
        """
        return any(self._states_done)

    @property
    def metadata(self) -> Dict[str, Any]:
        r"""Retrieve the metadata of the environment.

        This provides additional parameters and configuration details.

        Returns:
            Dict[str, Any]: A copy of the environment's metadata.
        """

        return self._metadata.copy()
