#!/usr/bin/env python3
"""
Code Review Assistant -- Custom Scorer Evaluation

Demonstrates creating custom scorers and evaluating an agent with them:

  1. CascadeEval.create_scorer()  -- build a custom LLM-judge scorer
  2. Run the agent to produce a trace
  3. CascadeEval.evaluate()       -- evaluate the trace with the custom scorer
  4. evaluate_spans()             -- evaluate individual spans

Agent context:
  A code review assistant that analyzes code snippets, checks style,
  finds potential bugs, and suggests improvements through a multi-agent
  pipeline (AnalysisAgent → ReviewAgent).

Usage:
    export CASCADE_API_KEY="csk_live_..."
    export CASCADE_ENDPOINT="http://localhost:8000/v1/traces"
    export ANTHROPIC_API_KEY="sk-ant-..."
    python tests/test_custom_scorer.py
"""
import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

try:
    from dotenv import load_dotenv
    load_dotenv(os.path.join(os.path.dirname(__file__), "..", ".env"))
except ImportError:
    pass

from cascade import (
    init_tracing, trace_run, trace_agent, wrap_llm_client, tool, function,
    CascadeEval,
)
from anthropic import Anthropic


# ---------------------------------------------------------------------------
# Helper functions (@function)
# ---------------------------------------------------------------------------

@function
def count_lines(code: str) -> int:
    """Count non-empty lines in a code snippet."""
    return sum(1 for line in code.strip().splitlines() if line.strip())


@function
def extract_function_names(code: str) -> list:
    """Extract function/method names from Python code."""
    import re
    return re.findall(r'def\s+(\w+)\s*\(', code)


@function
def compute_complexity_score(branches: int, loops: int, nesting: int) -> float:
    """Compute a simple cyclomatic complexity estimate."""
    return round(1 + branches * 0.4 + loops * 0.3 + nesting * 0.3, 2)


# ---------------------------------------------------------------------------
# Tools (@tool)
# ---------------------------------------------------------------------------

@tool
def analyze_code(code: str, language: str) -> str:
    """Analyze a code snippet for structure and metrics."""
    lines = count_lines(code)
    funcs = extract_function_names(code)

    import re
    branches = len(re.findall(r'\b(if|elif|else|match|case)\b', code))
    loops = len(re.findall(r'\b(for|while)\b', code))
    max_indent = max((len(line) - len(line.lstrip()) for line in code.splitlines() if line.strip()), default=0) // 4
    complexity = compute_complexity_score(branches, loops, max_indent)

    return (
        f"Code Analysis ({language}):\n"
        f"  Lines: {lines}\n"
        f"  Functions: {', '.join(funcs) if funcs else 'none detected'}\n"
        f"  Branches: {branches}, Loops: {loops}, Max nesting: {max_indent}\n"
        f"  Complexity score: {complexity} ({'low' if complexity < 3 else 'medium' if complexity < 6 else 'high'})"
    )


@tool
def check_style(code: str, language: str) -> str:
    """Check code against style guidelines."""
    issues = []

    for i, line in enumerate(code.splitlines(), 1):
        if len(line) > 100:
            issues.append(f"  Line {i}: exceeds 100-char limit ({len(line)} chars)")
        if line.rstrip() != line:
            issues.append(f"  Line {i}: trailing whitespace")
        if "\t" in line and language == "python":
            issues.append(f"  Line {i}: tabs used (prefer spaces in Python)")

    import re
    if language == "python":
        if not re.search(r'""".*?"""|\'\'\'.*?\'\'\'', code, re.DOTALL):
            funcs = extract_function_names(code)
            if funcs:
                issues.append("  Missing docstrings in function definitions")
        bare_excepts = re.findall(r'except\s*:', code)
        if bare_excepts:
            issues.append(f"  {len(bare_excepts)} bare except clause(s) -- catch specific exceptions")

    if not issues:
        return "Style check: PASSED -- no issues found."
    return f"Style check: {len(issues)} issue(s) found:\n" + "\n".join(issues)


@tool
def find_bugs(code: str, language: str) -> str:
    """Scan code for common bug patterns."""
    import re
    warnings = []

    if language == "python":
        if re.search(r'def\s+\w+\([^)]*=\s*\[\]', code):
            warnings.append("  Mutable default argument (list). Use None and initialize inside.")
        if re.search(r'def\s+\w+\([^)]*=\s*\{\}', code):
            warnings.append("  Mutable default argument (dict). Use None and initialize inside.")
        if re.search(r'==\s*None|!=\s*None', code):
            warnings.append("  Use 'is None' / 'is not None' instead of '==' / '!='.")
        if re.search(r'except\s*:', code):
            warnings.append("  Bare except catches all errors including KeyboardInterrupt.")
        if "global " in code:
            warnings.append("  Global statement used -- consider passing values as parameters.")
        if re.search(r'open\(', code) and "with " not in code:
            warnings.append("  File opened without 'with' statement -- risk of resource leak.")

    if not warnings:
        return "Bug scan: CLEAN -- no common bug patterns detected."
    return f"Bug scan: {len(warnings)} potential issue(s):\n" + "\n".join(warnings)


@tool
def suggest_improvements(code: str, analysis: str) -> str:
    """Suggest concrete improvements based on the analysis."""
    suggestions = []

    if "complexity score" in analysis.lower():
        import re
        match = re.search(r'Complexity score:\s*([\d.]+)', analysis)
        if match and float(match.group(1)) > 4:
            suggestions.append("  - Extract helper functions to reduce complexity")
            suggestions.append("  - Consider early returns to flatten nested conditions")

    if "Mutable default" in analysis:
        suggestions.append("  - Replace mutable defaults: def f(items=None): items = items or []")

    if "bare except" in analysis.lower():
        suggestions.append("  - Replace bare except with specific exceptions (ValueError, TypeError, etc.)")

    lines = count_lines(code)
    if lines > 50:
        suggestions.append("  - Function is long -- consider splitting into smaller functions")

    if "docstring" in analysis.lower() or "Missing doc" in analysis:
        suggestions.append("  - Add docstrings to all public functions")

    if not suggestions:
        suggestions.append("  - Code looks solid. Consider adding type hints for better IDE support.")
        suggestions.append("  - Add unit tests if not already present.")

    return "Improvement suggestions:\n" + "\n".join(suggestions)


# ---------------------------------------------------------------------------
# Tool schema + dispatch
# ---------------------------------------------------------------------------

TOOLS_SCHEMA = [
    {"name": "analyze_code", "description": "Analyze a code snippet for structure and metrics.",
     "input_schema": {"type": "object", "properties": {"code": {"type": "string"}, "language": {"type": "string"}}, "required": ["code", "language"]}},
    {"name": "check_style", "description": "Check code against style guidelines.",
     "input_schema": {"type": "object", "properties": {"code": {"type": "string"}, "language": {"type": "string"}}, "required": ["code", "language"]}},
    {"name": "find_bugs", "description": "Scan code for common bug patterns.",
     "input_schema": {"type": "object", "properties": {"code": {"type": "string"}, "language": {"type": "string"}}, "required": ["code", "language"]}},
    {"name": "suggest_improvements", "description": "Suggest concrete improvements based on the analysis.",
     "input_schema": {"type": "object", "properties": {"code": {"type": "string"}, "analysis": {"type": "string"}}, "required": ["code", "analysis"]}},
]

TOOL_DISPATCH = {
    "analyze_code": analyze_code,
    "check_style": check_style,
    "find_bugs": find_bugs,
    "suggest_improvements": suggest_improvements,
}


# ---------------------------------------------------------------------------
# Agent runner
# ---------------------------------------------------------------------------

def run_review_agent(client, code_snippet: str, language: str):
    """Run the code review agent. Returns (trace_id, root_span)."""

    with trace_run("CodeReviewSession", metadata={"language": language, "lines": len(code_snippet.splitlines())}) as root:
        trace_id = format(root.get_span_context().trace_id, "032x")
        messages = [{"role": "user", "content": f"Please review this {language} code:\n\n```{language}\n{code_snippet}\n```"}]

        # Phase 1 -- Analysis: run all checks
        with trace_agent("AnalysisAgent"):
            for _ in range(6):
                resp = client.messages.create(
                    model="claude-3-5-haiku-20241022",
                    max_tokens=700,
                    system=(
                        "You are a code analysis agent. Use the provided tools to "
                        "thoroughly analyze the code: run structural analysis, check "
                        "style, scan for bugs, and gather improvement suggestions. "
                        "Use ALL relevant tools before concluding."
                    ),
                    tools=TOOLS_SCHEMA,
                    messages=messages,
                )
                if resp.stop_reason == "end_turn":
                    break
                results = []
                for b in resp.content:
                    if b.type == "tool_use":
                        fn = TOOL_DISPATCH.get(b.name)
                        r = fn(**b.input) if fn else f"Unknown tool: {b.name}"
                        results.append({"type": "tool_result", "tool_use_id": b.id, "content": r})
                if not results:
                    break
                messages.append({"role": "assistant", "content": resp.content})
                messages.append({"role": "user", "content": results})

        # Phase 2 -- Review: compose a structured review
        analysis_text = "".join(b.text for b in resp.content if hasattr(b, "text"))
        with trace_agent("ReviewAgent"):
            client.messages.create(
                model="claude-3-5-haiku-20241022",
                max_tokens=900,
                system=(
                    "You are a senior code reviewer. Based on the analysis data, "
                    "write a structured review with: Summary, Issues (critical / "
                    "minor), Suggestions, and an overall quality rating (1-10)."
                ),
                messages=[{"role": "user", "content": f"Analysis results:\n{analysis_text}\n\nWrite a structured code review."}],
            )

    return trace_id, root


# ---------------------------------------------------------------------------
# Sample code snippets to review
# ---------------------------------------------------------------------------

SNIPPET_1 = '''\
def process_users(users=[]):
    result = {}
    for user in users:
        try:
            name = user["name"]
            email = user["email"]
            if email == None:
                continue
            age = user.get("age", 0)
            if age > 0:
                if age < 18:
                    result[name] = {"status": "minor", "email": email}
                else:
                    if age > 65:
                        result[name] = {"status": "senior", "email": email}
                    else:
                        result[name] = {"status": "adult", "email": email}
        except:
            print("Error processing user")
    return result
'''

SNIPPET_2 = '''\
import re
from typing import Optional

def validate_email(email: str) -> bool:
    """Validate an email address format."""
    pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$'
    return bool(re.match(pattern, email))

def sanitize_input(text: str, max_length: int = 500) -> str:
    """Sanitize user input by removing dangerous characters."""
    cleaned = re.sub(r'[<>&\\'"]', '', text)
    return cleaned[:max_length].strip()

def parse_config(filepath: str) -> Optional[dict]:
    """Parse a JSON config file safely."""
    import json
    try:
        with open(filepath, 'r') as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError) as e:
        print(f"Config error: {e}")
        return None
'''


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    api_key = os.getenv("ANTHROPIC_API_KEY")
    if not api_key:
        print("ERROR: Set ANTHROPIC_API_KEY"); return
    if not os.getenv("CASCADE_API_KEY"):
        print("ERROR: Set CASCADE_API_KEY"); return

    print("=" * 60)
    print("  Code Review Assistant -- Custom Scorer Evaluation")
    print("=" * 60)

    # ------------------------------------------------------------------
    # 1. Init tracing (no prebuilt evals -- we'll use custom ones)
    # ------------------------------------------------------------------
    print("\n[1] Initializing tracing...")
    init_tracing(project="code_review_assistant_trial")
    client = wrap_llm_client(Anthropic(api_key=api_key))
    print("    project = code_review_assistant\n")

    # ------------------------------------------------------------------
    # 2. Create custom scorers BEFORE running the agent
    # ------------------------------------------------------------------
    print("[2] Creating custom scorers...")
    evals = CascadeEval()
    created_ids = []

    # Scorer A: Trace-level -- did the agent use all review tools?
    tool_coverage = evals.create_scorer(
        name="Review Tool Coverage",
        scorer_type="llm_judge",
        scope="trace",
        threshold=0.5,
        description="Checks whether the code review agent used all relevant analysis tools",
        llm_provider="anthropic",
        llm_model="claude-3-5-haiku-20241022",
        output_type="categorical",
        choices=[
            {"label": "Y", "score": 1.0, "description": "All relevant tools were used"},
            {"label": "N", "score": 0.0, "description": "Some tools were skipped"},
        ],
        llm_template=(
            "Evaluate whether the code review agent was thorough in its analysis.\n\n"
            "Full execution trajectory:\n{{trajectory}}\n\n"
            "Available tools: analyze_code, check_style, find_bugs, suggest_improvements.\n"
            "A thorough review should use at least analyze_code, check_style, and find_bugs.\n\n"
            "Did the agent use all relevant tools? Answer Y if thorough, N if it skipped important checks.\n\n"
            "Answer (Y/N):"
        ),
        variable_mappings={"trajectory": "trajectory"},
        tags=["code-review", "coverage", "trace-level"],
    )
    created_ids.append(tool_coverage["scorer_id"])
    print(f"    Created: {tool_coverage['name']} (trace-level)")

    # Scorer B: Trace-level -- how actionable is the final review? (numeric 1-10)
    actionability = evals.create_scorer(
        name="Review Actionability",
        scorer_type="llm_judge",
        scope="trace",
        threshold=0.6,
        description="Rates how specific and actionable the review feedback is on a 0-1 scale",
        llm_provider="anthropic",
        llm_model="claude-3-5-haiku-20241022",
        output_type="numeric",
        min_score=0,
        max_score=1,
        llm_template=(
            "Rate the actionability of this code review on a scale of 0.0 to 1.0.\n\n"
            "Full execution trajectory:\n{{trajectory}}\n\n"
            "Scoring criteria:\n"
            "  0.0-0.3: Vague, generic advice with no specific references\n"
            "  0.3-0.6: Some specific feedback but missing concrete fix suggestions\n"
            "  0.6-0.8: Points to specific patterns, provides concrete fixes, "
            "distinguishes severity\n"
            "  0.8-1.0: Exceptional -- line-level references, ready-to-apply fixes, "
            "clear priority ranking\n\n"
            "Score (0.0-1.0):"
        ),
        variable_mappings={"trajectory": "trajectory"},
        tags=["code-review", "actionability", "trace-level"],
    )
    created_ids.append(actionability["scorer_id"])
    print(f"    Created: {actionability['name']} (trace-level, numeric 0-1)")

    # Scorer C: Span-level -- are individual LLM responses well-reasoned?
    llm_reasoning = evals.create_scorer(
        name="LLM Reasoning Quality",
        scorer_type="llm_judge",
        scope="span",
        threshold=0.5,
        description="Evaluates if individual LLM responses show clear reasoning",
        llm_provider="anthropic",
        llm_model="claude-3-5-haiku-20241022",
        output_type="categorical",
        choices=[
            {"label": "Y", "score": 1.0, "description": "Good reasoning"},
            {"label": "N", "score": 0.0, "description": "Poor reasoning"},
        ],
        llm_template=(
            "Evaluate the reasoning quality of this LLM response.\n\n"
            "Prompt:\n{{prompt}}\n\n"
            "Response:\n{{completion}}\n\n"
            "Does the response show clear, logical reasoning? "
            "Does it make appropriate tool calls or provide well-structured analysis?\n"
            "Answer Y if reasoning is sound, N if confused or shallow.\n\n"
            "Answer (Y/N):"
        ),
        variable_mappings={"prompt": "llm.prompt", "completion": "llm.completion"},
        tags=["code-review", "reasoning", "span-level"],
    )
    created_ids.append(llm_reasoning["scorer_id"])
    print(f"    Created: {llm_reasoning['name']} (span-level)")
    print(f"    Total: {len(created_ids)} custom scorers\n")

    # ------------------------------------------------------------------
    # 3. Run agent on snippet 1 (has several issues)
    # ------------------------------------------------------------------
    print("[3] Running code review agent on snippet 1 (buggy code)...")
    trace_id_1, span_1 = run_review_agent(client, SNIPPET_1, "python")
    print(f"    Trace: {trace_id_1}")

    # ------------------------------------------------------------------
    # 4. Run agent on snippet 2 (cleaner code)
    # ------------------------------------------------------------------
    print("\n[4] Running code review agent on snippet 2 (clean code)...")
    trace_id_2, span_2 = run_review_agent(client, SNIPPET_2, "python")
    print(f"    Trace: {trace_id_2}")

    # ------------------------------------------------------------------
    # 5. Wait, then run custom evaluations
    # ------------------------------------------------------------------
    print("\n[5] Waiting 5s for trace ingestion...")
    time.sleep(5)

    # Trace-level eval on snippet 1
    print("\n[6] Evaluating trace 1 with custom scorers (trace-level)...")
    trace_scorers = [tool_coverage["scorer_id"], actionability["scorer_id"]]
    eval_resp = evals.evaluate(trace_id=trace_id_1, scorer_ids=trace_scorers)
    print(f"    Success: {eval_resp.get('success')}")
    for r in eval_resp.get("results", []):
        tag = "PASS" if r.get("passed") else "FAIL"
        print(f"    [{tag}] {r['scorer_name']}: score={r.get('score')}")
        reasoning = r.get("reasoning", "")
        if reasoning:
            print(f"           {reasoning[:120]}...")

    # Span-level eval on snippet 1
    print("\n[7] Evaluating LLM spans with custom scorer (span-level)...")
    span_resp = evals.evaluate_spans(
        trace_id=trace_id_1,
        scorer_ids=[llm_reasoning["scorer_id"]],
        span_type="llm",
    )
    print(f"    Success: {span_resp.get('success')}")
    for r in span_resp.get("results", []):
        tag = "PASS" if r.get("passed") else "FAIL"
        span_label = r.get("span_name", "?")[:30]
        print(f"    [{tag}] {span_label}: {r['scorer_name']} score={r.get('score')}")

    # Trace-level eval on snippet 2
    print("\n[8] Evaluating trace 2 with custom scorers...")
    eval_resp_2 = evals.evaluate(trace_id=trace_id_2, scorer_ids=trace_scorers)
    print(f"    Success: {eval_resp_2.get('success')}")
    for r in eval_resp_2.get("results", []):
        tag = "PASS" if r.get("passed") else "FAIL"
        print(f"    [{tag}] {r['scorer_name']}: score={r.get('score')}")

    # ------------------------------------------------------------------
    # 9. List all results for trace 1
    # ------------------------------------------------------------------
    print("\n[9] Listing all evaluation results for trace 1...")
    all_results = evals.list_results(trace_id=trace_id_1)
    print(f"    Found {len(all_results)} result(s):")
    for r in all_results:
        tag = "PASS" if r.get("passed") else "FAIL"
        print(f"    [{tag}] {r.get('scorer_name', '?')[:35]:35s} score={r.get('score')}")

    # ------------------------------------------------------------------
    # Done
    # ------------------------------------------------------------------
    print("\n" + "=" * 60)
    print("  Done!")
    print("=" * 60)
    print("  - Created 3 custom scorers (2 trace-level [1 categorical, 1 numeric 0-1], 1 span-level)")
    print("  - Ran code review agent on 2 different snippets")
    print("  - Evaluated both traces with custom scorers")
    print("  - Evaluated LLM spans for reasoning quality")
    print("  - View results at: http://localhost:3000  (project: code_review_assistant)")
    print()


if __name__ == "__main__":
    main()
