"""
Lab 6.8: Build an Eval Harness -- SOLUTION

Golden Q&A pairs + automated pass/fail scoring with pass@k.
"""

from typing import Callable

GOLDEN_DATASET = [
    {"id": "t001", "input": "Find users created in last 24h",
     "expected_tools": ["query_database"],
     "expected_params": {"query": "SELECT * FROM users WHERE created_at > now() - interval '24 hours'"},
     "expected_output_contains": ["users"], "weight": 1.0},
    {"id": "t002", "input": "Search for AI agent frameworks",
     "expected_tools": ["search_web"],
     "expected_output_contains": ["framework"], "weight": 1.0},
    {"id": "t003", "input": "Deploy build to staging",
     "expected_tools": ["run_deploy"],
     "expected_params": {"environment": "staging"},
     "expected_output_contains": ["deploy"], "weight": 1.5},
    {"id": "t004", "input": "Explain recursion",
     "expected_tools": [],
     "expected_output_contains": ["base case"], "weight": 0.5},
    {"id": "t005", "input": "Fix vulnerability in auth.py",
     "expected_tools": ["read_file", "edit_file"],
     "expected_output_contains": ["fixed"], "weight": 2.0},
    {"id": "t006", "input": "What's the stock price of AAPL?",
     "expected_tools": ["fetch_stock_price", "search_web"],
     "expected_output_contains": ["AAPL", "$"], "weight": 1.0},
    {"id": "t007", "input": "Send a welcome email to new@user.com",
     "expected_tools": ["send_email"],
     "expected_params": {"to": "new@user.com"},
     "expected_output_contains": ["sent"], "weight": 1.5},
    {"id": "t008", "input": "Analyze this CSV: sales.csv",
     "expected_tools": ["read_file", "analyze_data"],
     "expected_output_contains": ["revenue", "average"], "weight": 1.5},
    {"id": "t009", "input": "What packages are in package.json?",
     "expected_tools": ["read_file"],
     "expected_output_contains": ["dependencies"], "weight": 0.5},
    {"id": "t010", "input": "Create a SQL migration for users table",
     "expected_tools": ["write_file"],
     "expected_params": {"path": "migrations/"},
     "expected_output_contains": ["CREATE TABLE"], "weight": 1.5},
]


class EvalHarness:
    def __init__(self, dataset: list[dict]):
        self.dataset = dataset
    
    def score_case(self, result: Any, test_case: dict) -> float:
        score = 0.0
        max_score = 0.0
        
        expected_tools = set(test_case.get("expected_tools", []))
        actual_tools = set(t.get("tool") for t in result.get("tool_calls", []))
        if expected_tools:
            overlap = expected_tools & actual_tools
            score += len(overlap) / len(expected_tools) * 0.4
            max_score += 0.4
        else:
            score += 0.4
            max_score += 0.4
        
        expected_params = test_case.get("expected_params", {})
        if expected_params:
            for at in result.get("tool_calls", []):
                for key, val in expected_params.items():
                    if at.get("params", {}).get(key) == val:
                        score += 0.3 / len(expected_params)
            max_score += 0.3
        else:
            score += 0.3
            max_score += 0.3
        
        output = result.get("output", "")
        expected_out = test_case.get("expected_output_contains", [])
        if expected_out:
            matches = sum(1 for s in expected_out if s.lower() in output.lower())
            score += (matches / len(expected_out)) * 0.3
            max_score += 0.3
        else:
            score += 0.3
            max_score += 0.3
        
        return score / max_score if max_score > 0 else 0.0
    
    def evaluate(self, agent_func: Callable, k_values: list[int] = None) -> dict:
        if k_values is None:
            k_values = [1, 3, 5]
        
        k_results = {k: [] for k in k_values}
        
        for case in self.dataset:
            all_attempts = []
            max_k = max(k_values)
            for _ in range(max_k):
                result = agent_func(case["input"])
                score = self.score_case(result, case)
                all_attempts.append(score)
            
            for k in k_values:
                attempts = all_attempts[:k]
                best = max(attempts)
                passed = best >= 0.7
                k_results[k].append({
                    "case_id": case["id"],
                    "best_score": round(best, 3),
                    "passed": passed,
                    "all_scores": [round(s, 3) for s in attempts],
                    "weight": case.get("weight", 1.0)
                })
        
        output = {}
        for k in k_values:
            results = k_results[k]
            weighted_pass = sum(r["weight"] for r in results if r["passed"])
            total_weight = sum(r["weight"] for r in results)
            pass_rate = weighted_pass / total_weight if total_weight > 0 else 0
            
            all_passed = sum(1 for r in results if r["passed"])
            consistency = sum(1 for r in results if all(s >= 0.7 for s in r["all_scores"])) / len(results) if results else 0
            
            output[f"pass@{k}"] = round(pass_rate, 3)
            output[f"pass^{k}"] = round(consistency, 3)
            output[f"results_{k}"] = results
        
        output["total_cases"] = len(self.dataset)
        output["weighted_scores"] = {f"pass@{k}": output[f"pass@{k}"] for k in k_values}
        
        return output


def example_agent(input_text: str) -> dict:
    """Example agent that simulates responses for testing."""
    from difflib import SequenceMatcher
    
    for case in GOLDEN_DATASET:
        ratio = SequenceMatcher(None, input_text.lower(), case["input"].lower()).ratio()
        if ratio > 0.3:
            return {
                "tool_calls": [{"tool": t, "params": case.get("expected_params", {}) or {}} for t in case["expected_tools"]],
                "output": f"Result for: {case['input']}. Found relevant data."
            }
    
    return {"tool_calls": [], "output": "I can help with that."}


if __name__ == "__main__":
    harness = EvalHarness(GOLDEN_DATASET)
    results = harness.evaluate(example_agent, k_values=[1, 3, 5])
    
    print("=" * 55)
    print(" EVAL HARNESS RESULTS")
    print("=" * 55)
    
    for k in [1, 3, 5]:
        print(f"\npass@{k}:  {results[f'pass@{k}']:.1%}")
        print(f"pass^{k}:  {results[f'pass^{k}']:.1%}")
    
    print(f"\nTotal test cases: {results['total_cases']}")
    print("\nDetailed results (pass@3):")
    for r in results["results_3"]:
        icon = "+" if r["passed"] else "-"
        print(f"  {icon} {r['case_id']}: best={r['best_score']:.1%} attempts={r['all_scores']}")