import os
import json
import sys
import re
from openharness.core.harness import SimpleHarness, Observation
from openharness.tools.read_file import ReadFileTool
from openharness.tools.write_file import WriteFileTool
from openharness.tools.search import SearchTool
from openharness.tools.run_tests import RunTestsTool
from openharness.sandbox.docker_env import get_sandbox
from openharness.trajectory.store import TrajectoryStore


class BugFixHarness(SimpleHarness):
    """Harness for fixing failing unit tests"""
    
    def __init__(self, tasks=None, workspace_dir=None):
        # Initialize sandbox
        self.sandbox = get_sandbox("subprocess")
        
        # Store workspace directory (may be None)
        if workspace_dir is not None:
            self.workspace_dir = os.path.abspath(workspace_dir)
        else:
            self.workspace_dir = None
        self.tasks = tasks or []
        self.current_task_index = 0
        
        # Initialize tools (base_dir will be set in reset)
        tool_list = [
            ReadFileTool(),
            WriteFileTool(),
            SearchTool(),
            RunTestsTool(sandbox=self.sandbox)
        ]
        self.tools = {tool.name: tool for tool in tool_list}
        
        super().__init__(list(self.tools.values()), self.tasks)
        
        # Trajectory store
        self.trajectory_store = TrajectoryStore()
        self.episode_data = {
            "task": None,
            "actions": [],
            "reward": 0.0,
            "start_time": None,
            "end_time": None
        }
        
        # Current task's test file (if any)
        self.current_test_file = None
    
    def reset(self) -> Observation:
        """Reset the environment to initial state"""
        observation = super().reset()
        
        # If we have a task, set up the workspace for the tools
        if self.episode_data["task"]:
            # Use provided workspace_dir if given, otherwise use current directory
            if self.workspace_dir is not None:
                base_dir = self.workspace_dir
            else:
                base_dir = os.getcwd()
            # Set the base_dir for each tool that needs it
            for tool_name in ["read_file", "write_file", "search", "run_tests"]:
                self.tools[tool_name].base_dir = base_dir
            
            # Determine the test file for the current task
            self.current_test_file = None
            if "files" in self.episode_data["task"]:
                for file_name in self.episode_data["task"]["files"]:
                    if file_name.startswith("test_") and file_name.endswith(".py"):
                        self.current_test_file = file_name
                        break
        
        return observation
    
    def _parse_test_reward(self, test_output: str) -> float:
        """Parse test output to calculate a reward between 0 and 1"""
        try:
            # Pattern 1: "X passed, Y failed" (pytest style)
            passed_match = re.search(r'(\d+)\s+passed', test_output)
            failed_match = re.search(r'(\d+)\s+failed', test_output)
            
            if passed_match and failed_match:
                passed = int(passed_match.group(1))
                failed = int(failed_match.group(1))
                total = passed + failed
                if total > 0:
                    return passed / total
            
            # Pattern 2: Look for summary line like "4 passed in 0.00s"
            summary_match = re.search(r'(\d+)\s+passed', test_output)
            if summary_match and 'failed' not in test_output.lower():
                passed = int(summary_match.group(1))
                return 1.0 if passed > 0 else 0.0
            
            # If we see "all tests passed", return 1.0
            if 'all tests passed' in test_output.lower():
                return 1.0
            
            # If we see "failed", return low reward
            if 'failed' in test_output.lower():
                failed_numbers = re.findall(r'(\d+)\s+failed', test_output.lower())
                if failed_numbers:
                    failed = int(failed_numbers[0])
                    passed_numbers = re.findall(r'(\d+)\s+passed', test_output.lower())
                    if passed_numbers:
                        passed = int(passed_numbers[0])
                        total = passed + failed
                        if total > 0:
                            return passed / total
                    else:
                        return max(0.0, 1.0 - (failed * 0.2))
                return 0.1
            
            return 0.1
            
        except Exception:
            return 0.1
    
    def step(self, action):
        """Execute a step and return the step reward (not just episode reward)"""
        observation, reward, done, info = super().step(action)
        
        # If this was a run_tests action, calculate the step reward
        if action.get("tool") == "run_tests":
            # Get the output from the observation
            output = observation.data.get("output", "")
            step_reward = self._parse_test_reward(output)
        else:
            step_reward = 0.0
        
        return observation, step_reward, done, info
    
    def calculate_reward(self) -> float:
        """Calculate reward based on test results"""
        if not self.episode_data["actions"]:
            return 0.0
        
        # Find the last test action
        for action in reversed(self.episode_data["actions"]):
            if action["tool"] == "run_tests":
                return self._parse_test_reward(action["output"])
        
        return 0.0
    
    def is_done(self) -> bool:
        """Check if the task is complete (tests passing)"""
        if not self.episode_data["actions"]:
            return False
        
        # Find the last test action
        for action in reversed(self.episode_data["actions"]):
            if action["tool"] == "run_tests":
                return self._parse_test_reward(action["output"]) == 1.0
        
        return False
