from abc import ABC, abstractmethod
from typing import List, Dict, Any, Optional
import json
import time
from openharness.tools.tool import Tool
from openharness.trajectory.store import TrajectoryStore


class Observation:
    def __init__(self, data: Dict[str, Any]):
        self.data = data

    def __getitem__(self, key):
        return self.data.get(key)

    def __setitem__(self, key, value):
        self.data[key] = value

    def get(self, key, default=None):
        return self.data.get(key, default)


class Harness(ABC):
    def __init__(self, tools: List[Tool], tasks: List[Dict[str, Any]] = None):
        self.tools = {tool.name: tool for tool in tools}
        self.tasks = tasks or []
        self.current_task_index = 0
        self.trajectory_store = TrajectoryStore()
        self.episode_data = {
            "task": None,
            "actions": [],
            "reward": 0.0,
            "start_time": None,
            "end_time": None,
        }

    def reset(self) -> Observation:
        """Reset the environment to initial state"""
        # Clear episode data
        self.episode_data = {
            "task": None,
            "actions": [],
            "reward": 0.0,
            "start_time": None,
            "end_time": None,
        }

        # Reset tools if they have reset method
        for tool in self.tools.values():
            if hasattr(tool, "reset"):
                tool.reset()

        # Select next task if available
        if self.tasks and self.current_task_index < len(self.tasks):
            self.episode_data["task"] = self.tasks[self.current_task_index]
            self.current_task_index += 1
            self.episode_data["start_time"] = time.time()
        else:
            self.episode_data["task"] = None

        return self._get_observation()

    def step(
        self, action: Dict[str, Any]
    ) -> tuple[Observation, float, bool, Dict[str, Any]]:
        """Execute an action in the environment"""
        tool_name = action.get("tool")
        tool_input = action.get("input", {})

        if tool_name not in self.tools:
            observation_data = {
                "error": f"Tool {tool_name} not found",
                "available_tools": list(self.tools.keys()),
            }
            reward = -1.0  # Penalty for invalid action
            done = False
        else:
            tool = self.tools[tool_name]
            try:
                # Execute the tool
                tool_output = tool.run(**tool_input)

                # Record action in trajectory
                action_record = {
                    "tool": tool_name,
                    "input": tool_input,
                    "output": tool_output,
                    "timestamp": time.time(),
                }
                self.episode_data["actions"].append(action_record)

                observation_data = {"output": tool_output, "tool_used": tool_name}
                reward = 0.0  # Default reward, will be calculated later
                done = False

            except Exception as e:
                observation_data = {"error": str(e), "tool_used": tool_name}
                reward = -1.0  # Penalty for tool execution error
                done = False

        observation = Observation(observation_data)
        info = {}

        return observation, reward, done, info

    def calculate_reward(self) -> float:
        """Calculate reward based on task completion"""
        # This should be overridden by specific harness implementations
        return self.episode_data.get("reward", 0.0)

    def is_done(self) -> bool:
        """Check if the episode is complete"""
        # This should be overridden by specific harness implementations
        return False

    def _get_observation(self) -> Observation:
        """Get current observation"""
        if self.episode_data["task"]:
            observation_data = {
                "task": self.episode_data["task"],
                "message": "Task loaded. Ready for actions.",
            }
        else:
            observation_data = {"message": "No more tasks available", "task": None}
        return Observation(observation_data)

    def end_episode(self):
        """End the current episode and save trajectory"""
        self.episode_data["end_time"] = time.time()
        self.episode_data["reward"] = self.calculate_reward()

        # Save trajectory
        if self.episode_data["task"]:
            self.trajectory_store.save_trajectory(self.episode_data)

        return self.episode_data.copy()


class SimpleHarness(Harness):
    """A simple implementation of Harness for basic use cases"""

    def __init__(self, tools: List[Tool], tasks: List[Dict[str, Any]] = None):
        super().__init__(tools, tasks)
        self.max_steps = 10
        self.step_count = 0

    def reset(self) -> Observation:
        self.step_count = 0
        return super().reset()

    def step(
        self, action: Dict[str, Any]
    ) -> tuple[Observation, float, bool, Dict[str, Any]]:
        observation, reward, done, info = super().step(action)
        self.step_count += 1

        # Check if we've exceeded max steps
        if self.step_count >= self.max_steps:
            done = True

        # Check if task is done (to be overridden)
        if not done:
            done = self.is_done()

        return observation, reward, done, info

    def is_done(self) -> bool:
        """Simple completion check - in real harness this would be task-specific"""
        # For now, we'll consider it done if no error occurred in last action
        if not self.episode_data["actions"]:
            return False

        last_action = self.episode_data["actions"][-1]
        return "error" not in last_action.get("output", {})
