from env.codedebugger_env import CodeDebuggerEnv
from agents.fixer import FixerAgent

class Debugger:
    def __init__(self, max_iterations=5):
        self.env = CodeDebuggerEnv(max_iterations)
        self.fixer = FixerAgent()
    
    def run(self, problem: dict, verbose=True) -> dict:
        """
        Full debugging loop for one problem.
        """
        obs = self.env.reset(problem)
        
        best_reward = -float('inf')
        best_code = None
        best_iter = 0
        
        iter_data = []
        reward_gain = 0.0
        iter1_reward = 0.0
        final_reward = 0.0
        tests_passed_final = 0
        tests_total = len(problem.get("test_cases", []))
        solved = False
        
        if verbose:
            print(f"\n--- Debugging: {problem['id']} ({problem.get('title', '')}) ---")
            
        for i in range(1, self.env.max_iterations + 1):
            # Format history for fixer
            prev_explanation = None
            test_results = None
            if iter_data:
                # Use results from the immediate previous iteration
                last_env_hist = self.env.history[-1]
                test_results = last_env_hist["test_results"]
                prev_explanation = iter_data[-1]["fix_result"].get("explanation")
            
            # On first iteration, fix the original buggy code. 
            # On subsequent iterations, try to fix the previously generated code.
            code_to_fix = obs["buggy_code"] if i == 1 else self.env.history[-1]["code"]
            
            # Propose fix
            fix_result = self.fixer.fix_code(
                buggy_code=code_to_fix,
                error_type=problem.get("error_type", ""),
                description=problem.get("description", ""),
                test_cases=problem.get("test_cases", []),
                test_results=test_results,
                previous_explanation=prev_explanation,
                iteration=i
            )
            
            fixed_code = fix_result.get("fixed_code", code_to_fix)
            
            # Step environment
            obs, reward_dict, done, info = self.env.step(fixed_code)
            current_reward = reward_dict["total"]
            
            if i == 1:
                iter1_reward = current_reward
                
            final_reward = current_reward
            tests_passed_final = obs["test_results"].get("tests_passed", 0)
            
            # Track best separately
            if current_reward > best_reward:
                best_reward = current_reward
                best_code = fixed_code
                best_iter = i
                
            iter_info = {
                "iteration": i,
                "fix_result": fix_result,
                "reward": reward_dict,
                "done": done,
                "info": info
            }
            iter_data.append(iter_info)
            
            if verbose:
                print(f"Iter {i}: passed {tests_passed_final}/{tests_total} | "
                      f"reward: {current_reward:.1f} | method: {fix_result.get('method')}")
            
            if done:
                if info.get("solved"):
                    solved = True
                    if verbose:
                        print("-> Solved!")
                break
                
        # Calculate gain from first iteration
        if best_reward != -float('inf'):
            reward_gain = best_reward - iter1_reward
        else:
            best_reward = 0.0
            
        return {
            "problem_id": problem["id"],
            "difficulty": problem["difficulty"],
            "title": problem.get("title", ""),
            "n_iterations": len(iter_data),
            "final_reward": final_reward,
            "best_reward": best_reward,
            "iter1_reward": iter1_reward,
            "best_iter": best_iter,
            "best_code": best_code,
            "reward_gain": reward_gain,
            "tests_passed_final": tests_passed_final,
            "tests_total": tests_total,
            "solved": solved,
            "iterations": iter_data
        }
    
    def run_all(self, problems: list, verbose=True) -> list:
        """Run all problems, return list of results"""
        results = []
        for prob in problems:
            res = self.run(prob, verbose=verbose)
            results.append(res)
        return results
    
    def print_summary_table(self, results: list):
        """
        Print this exact format:
        
        Problem              Diff    iter1  final   best  gain  solved
        ----------------------------------------------------------------
        ...
        """
        print(f"\n{'Problem':<20} {'Diff':<7} {'iter1':>6} {'final':>6} {'best':>6} {'gain':>6} {'solved':>6}")
        print("-" * 64)
        
        solved_count = 0
        total_best_reward = 0.0
        
        for r in results:
            pid = str(r["problem_id"])
            if len(pid) > 20: 
                pid = pid[:17] + "..."
            
            diff = str(r["difficulty"])
            iter1 = r.get("iter1_reward", 0.0)
            final = r["final_reward"]
            best = r["best_reward"]
            gain = r["reward_gain"]
            solved_str = "YES" if r["solved"] else "NO"
            
            print(f"{pid:<20} {diff:<7} {iter1:>6.1f} {final:>6.1f} {best:>6.1f} {gain:>+6.1f} {solved_str:>6}")
            
            if r["solved"]: 
                solved_count += 1
            total_best_reward += best
            
        print("-" * 64)
        avg_best = total_best_reward / len(results) if results else 0.0
        print(f"TOTAL: {solved_count}/{len(results)} solved | Avg best reward: {avg_best:.1f}")