import sys
import os
import time

try:
    sys.stdout.reconfigure(encoding='utf-8')
except AttributeError:
    pass

GREEN = '\033[92m'
RED = '\033[91m'
RESET = '\033[0m'

def log_ok(msg):
    print(f"{GREEN}✅ {msg}{RESET}")

def log_err(msg):
    print(f"{RED}❌ {msg}{RESET}")

def verify():
    print("Starting End-to-End Verification...\n")
    failures = []
    
    # 1. Imports
    try:
        from data.bug_dataset import TRAINING_SCENARIOS
        from env.codedebugger_env import CodeDebuggerEnv
        from env.executor import CodeExecutor
        from agents.reward import calculate_reward
        from agents.critic import run_all_anti_hack_checks
        from agents.fixer import FixerAgent
        log_ok("Imports: All modules imported successfully")
    except Exception as e:
        log_err(f"Imports failed: {e}")
        failures.append("Imports")
        return # fatal
        
    # 2. Dataset length
    try:
        if len(TRAINING_SCENARIOS) == 30:
            log_ok("Dataset: 30 problems loaded")
        else:
            log_err(f"Dataset: Expected 30 problems, found {len(TRAINING_SCENARIOS)}")
            failures.append("Dataset length")
    except Exception as e:
        log_err(f"Dataset check failed: {e}")
        failures.append("Dataset")
        
    # 3. Create Env
    try:
        env = CodeDebuggerEnv(max_iterations=5)
        log_ok("Environment: CodeDebuggerEnv instantiated")
    except Exception as e:
        log_err(f"Environment creation failed: {e}")
        failures.append("Env instantiation")
        return # fatal
        
    # 4. reset()
    try:
        prob0 = TRAINING_SCENARIOS[0]
        obs = env.reset(prob0)
        if obs["problem_id"] == prob0["id"]:
            log_ok("Environment: reset() works properly")
        else:
            log_err("Environment: reset() observation mismatch")
            failures.append("Env reset mismatch")
    except Exception as e:
        log_err(f"Env.reset() failed: {e}")
        failures.append("Env reset")
        return # fatal
        
    # 5. step() with FIXED code
    try:
        fixed_code = prob0["fixed_code"]
        obs_fixed, reward_fixed, done_fixed, info_fixed = env.step(fixed_code)
        if reward_fixed["total"] > 50:
            log_ok(f"Environment: step() with FIXED code gives high reward ({reward_fixed['total']:.1f})")
        else:
            log_err(f"Environment: step() with FIXED code reward too low ({reward_fixed['total']:.1f})")
            failures.append("Fixed code reward")
    except Exception as e:
        log_err(f"Env.step() [FIXED] failed: {e}")
        failures.append("Fixed code step")
        reward_fixed = {"total": 0} # fallback
        
    # 6. step() with BUGGY code
    try:
        buggy_code = prob0["buggy_code"]
        obs_buggy, reward_buggy, done_buggy, info_buggy = env.step(buggy_code)
        if reward_buggy["total"] < reward_fixed["total"]:
            log_ok(f"Environment: step() with BUGGY code gives lower reward ({reward_buggy['total']:.1f})")
        else:
            log_err(f"Environment: BUGGY code reward not lower than fixed code")
            failures.append("Buggy code reward")
    except Exception as e:
        log_err(f"Env.step() [BUGGY] failed: {e}")
        failures.append("Buggy code step")
        
    # 7. step() with "pass"
    try:
        obs_pass, reward_pass, done_pass, info_pass = env.step("pass")
        if reward_pass.get("anti_hack_penalty", 0) < 0:
            log_ok(f"Environment: step() with 'pass' correctly applies anti-hack penalty ({reward_pass['anti_hack_penalty']})")
        else:
            log_err(f"Environment: 'pass' code did not trigger anti-hack penalty")
            failures.append("Anti-hack step penalty")
    except Exception as e:
        log_err(f"Env.step() ['pass'] failed: {e}")
        failures.append("Pass step")
        reward_pass = {}
        
    # 8. Check reward dict keys
    try:
        expected_keys = {"total", "test_score", "fix_quality", "format_score", "speed_score", "safety_score", "improvement_score", "anti_hack_penalty", "process_bonus"}
        actual_keys = set(reward_fixed.keys())
        # Subset check because our implementation might have exactly these or a few more internal ones
        if expected_keys.issubset(actual_keys):
            log_ok("Reward: Dict contains all expected scoring components")
        else:
            missing = expected_keys - actual_keys
            log_err(f"Reward: Missing keys {missing}")
            failures.append("Reward keys")
    except Exception as e:
        log_err(f"Reward Dict check failed: {e}")
        failures.append("Reward keys check")
        
    # 9. Run run_all_anti_hack_checks on empty code
    try:
        ah_res = run_all_anti_hack_checks(
            code="",
            original_code=prob0["buggy_code"],
            test_cases=prob0["test_cases"],
            test_results={"pass_rate": 0.0, "execution_time_total": 0.1},
            execution_time=0.1
        )
        if not ah_res["all_passed"] and ah_res["total_penalty"] < 0:
            log_ok(f"Critic: Standalone empty code caught with penalty {ah_res['total_penalty']}")
        else:
            log_err("Critic: Standalone empty code not caught")
            failures.append("Critic empty code")
    except Exception as e:
        log_err(f"Critic standalone check failed: {e}")
        failures.append("Critic check")
        
    # 10. Print final summary
    print("\n" + "="*40)
    if not failures:
        print(f"{GREEN}✅ All systems functional{RESET}")
    else:
        print(f"{RED}❌ Failures detected in: {', '.join(failures)}{RESET}")


if __name__ == "__main__":
    verify()