import json, os
from datetime import datetime
from dotenv import load_dotenv
load_dotenv()

import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from data.bug_dataset import TRAINING_SCENARIOS
from orchestrator import Debugger

def run_baseline():
    debugger = Debugger(max_iterations=3)
    all_results = []
    
    print(f"Running baseline on {len(TRAINING_SCENARIOS)} problems...")
    
    for i, problem in enumerate(TRAINING_SCENARIOS):
        print(f"\n[{i+1}/{len(TRAINING_SCENARIOS)}] {problem['title']} [{problem['difficulty']}]")
        result = debugger.run(problem, verbose=True)
        all_results.append(result)
    
    # Build summary — OUTSIDE the loop
    by_diff = {"easy": [], "medium": [], "hard": []}
    for r in all_results:
        by_diff[r["difficulty"]].append(r)
    
    summary = {
        "total_problems": len(all_results),
        "solved": sum(1 for r in all_results if r["solved"]),
        "avg_best_reward": sum(r["best_reward"] for r in all_results) / len(all_results),
        "by_difficulty": {
            diff: {
                "solved": sum(1 for r in problems if r["solved"]),
                "total": len(problems),
                "avg_reward": sum(r["best_reward"] for r in problems) / max(len(problems), 1)
            }
            for diff, problems in by_diff.items()
        }
    }
    
    output = {
        "format_version": 1,
        "model": "llama-3.1-8b-instant",
        "timestamp": datetime.now().isoformat(),
        "results": all_results,
        "summary": summary
    }
    
    os.makedirs("outputs", exist_ok=True)
    with open("outputs/baseline_scores.json", "w") as f:
        json.dump(output, f, indent=2)
    
    # Print summary ONCE here
    print("\n" + "="*60)
    print("BASELINE SUMMARY")
    print("="*60)
    debugger.print_summary_table(all_results)
    print(f"\nSolved: {summary['solved']}/{summary['total_problems']}")
    print(f"Avg best reward: {summary['avg_best_reward']:.1f}")
    print("Saved to outputs/baseline_scores.json")

if __name__ == "__main__":
    run_baseline()