"""
Lab 5.10: CI/CD Pipeline -- SOLUTION

Golden dataset evaluation + automated regression gate.
"""

import sys

GOLDEN_DATASET_EXAMPLE = [
    {"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": "What's the weather in Tokyo?",
     "expected_tools": ["search_web", "fetch_weather"],
     "expected_output_contains": ["Tokyo", "°"], "weight": 1.0},
    {"id": "t003", "input": "Deploy the latest build to staging",
     "expected_tools": ["run_deploy"],
     "expected_params": {"environment": "staging"},
     "expected_output_contains": ["deploy", "staging"], "weight": 1.5},
    {"id": "t004", "input": "Explain how recursion works",
     "expected_tools": [],
     "expected_output_contains": ["base case", "recursive"], "weight": 0.5},
    {"id": "t005", "input": "Find and fix the security vulnerability in auth.py",
     "expected_tools": ["read_file", "edit_file"],
     "expected_output_contains": ["fixed", "vulnerability"], "weight": 2.0},
]


def default_agent(input_text: str) -> dict:
    """Reference agent for testing. Replace with your actual agent."""
    return {
        "tool_calls": [{"tool": "query_database", "params": {"query": "SELECT * FROM users"}}],
        "output": "Found users in database"
    }


class AgentCIPipeline:
    def __init__(self, golden_dataset: list[dict], threshold: float = 0.8):
        self.dataset = golden_dataset
        self.threshold = threshold
    
    def score_case(self, agent_result: dict, test_case: dict) -> float:
        score = 0.0
        max_score = 0.0
        
        # Tool selection (40% of score)
        expected_tools = set(test_case.get("expected_tools", []))
        actual_tools = set(t.get("tool") for t in agent_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  # No tools expected, passed
            max_score += 0.4
        
        # Tool params (30% of score)
        expected_params = test_case.get("expected_params", {})
        if expected_params:
            actual_tools_list = agent_result.get("tool_calls", [])
            for at in actual_tools_list:
                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 content (30% of score)
        output = agent_result.get("output", "")
        expected_contains = test_case.get("expected_output_contains", [])
        if expected_contains:
            matches = sum(1 for s in expected_contains if s.lower() in output.lower())
            score += (matches / len(expected_contains)) * 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 run_eval(self, agent_func=default_agent, k: int = 3) -> dict:
        results = []
        for case in self.dataset:
            case_scores = []
            for attempt in range(k):
                result = agent_func(case["input"])
                case_scores.append(self.score_case(result, case))
            
            best_score = max(case_scores)
            passed = best_score >= 0.7
            
            results.append({
                "id": case["id"],
                "input": case["input"][:50],
                "best_score": round(best_score, 3),
                "passed": passed,
                "weight": case.get("weight", 1.0)
            })
        
        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
        
        return {
            "pass@1": pass_rate,
            "pass@k": pass_rate,
            "k": k,
            "results": results,
            "total_cases": len(results),
            "passed_cases": sum(1 for r in results if r["passed"])
        }
    
    def check_gate(self, eval_results: dict) -> bool:
        return eval_results["pass@k"] >= self.threshold
    
    def promote(self):
        print("  -> Copying config to production...")
        print("  -> Restarting agent service...")
        print("  -> Deploying new prompt version...")
        print("  + Production deployment complete")
    
    def rollback(self, reason: str):
        print(f"  -> Rolling back: {reason}")
        print("  -> Restoring previous agent config...")
        print("  -> Pinning model to previous version...")
        print("  + Rollback complete")
        print("  -> Alert sent to #agent-alerts")
    
    def run(self, agent_func=default_agent):
        print("=" * 55)
        print(" AGENT CI/CD PIPELINE")
        print("=" * 55)
        
        print("\n [CHECKOUT] Loading agent config from agent-config-v42.yaml")
        print("   Config hash: a1b2c3d4")
        print("   Model: claude-sonnet-4-20260501")
        
        print("\n [EVAL] Running golden dataset...")
        results = self.run_eval(agent_func, k=3)
        
        print(f"   pass@1: {results['pass@1']:.1%}")
        print(f"   pass@k: {results['pass@k']:.1%}")
        print(f"   Passed: {results['passed_cases']}/{results['total_cases']} cases")
        
        print("\n   Case breakdown:")
        for r in results["results"]:
            icon = "+" if r["passed"] else "-"
            print(f"   {icon} {r['id']}: score={r['best_score']:.1%} (weight={r['weight']})")
        
        print(f"\n [GATE] Threshold: {self.threshold:.0%}")
        if self.check_gate(results):
            print(f"   + GATE PASSED (pass@k={results['pass@k']:.1%} >= {self.threshold:.0%})")
            print("\n [PROMOTE] Promoting to production...")
            self.promote()
            print("\n [OK] Pipeline complete")
        else:
            print(f"   - GATE FAILED (pass@k={results['pass@k']:.1%} < {self.threshold:.0%})")
            print("\n [ROLLBACK] Reverting...")
            self.rollback(f"pass@k {results['pass@k']:.1%} below threshold")
            print("\n [FAIL] Pipeline failed")
            sys.exit(1)


if __name__ == "__main__":
    pipeline = AgentCIPipeline(GOLDEN_DATASET_EXAMPLE, threshold=0.8)
    pipeline.run()