#!/usr/bin/env python3
"""
LangGraph Eval Generator.

Generate evaluation datasets and run evaluators for multi-agent systems.
Supports two modes:
  1. Dataset creation — build a LangSmith eval dataset from a JSON/YAML spec
  2. Evaluator generation — produce Python code for routing accuracy and
     resolution coverage evaluators

Usage:
    # Generate a dataset from a spec file
    python lg-eval-generator.py dataset --name my-evals --spec examples.json

    # Generate evaluator code
    python lg-eval-generator.py evaluator --name my-evals --output ./evals

    # Generate an example spec file to fill in
    python lg-eval-generator.py example-spec --output ./examples.json
"""

import argparse
import json
import os
from typing import Any, Dict, List


EXAMPLE_SPEC = [
    {
        "question": "I need to change my payment method to a credit card.",
        "expected_agents": ["billing"],
        "must_mention": ["payment", "credit card"],
    },
    {
        "question": "My SSO integration is returning error code SAML-401.",
        "expected_agents": ["tech_support"],
        "must_mention": ["SSO", "SAML"],
    },
    {
        "question": "I want to upgrade to Enterprise and also fix my broken SSO.",
        "expected_agents": ["tech_support", "account"],
        "must_mention": ["SSO", "upgrade"],
    },
    {
        "question": "Can you tell me who my account manager is?",
        "expected_agents": ["account"],
        "must_mention": ["account manager"],
    },
    {
        "question": "What's the status of my refund?",
        "expected_agents": ["billing"],
        "must_mention": ["refund"],
    },
    {
        "question": "Reset my password and downgrade my plan.",
        "expected_agents": ["tech_support", "account"],
        "must_mention": ["password", "downgrade"],
    },
]


def generate_dataset(name: str, spec_path: str, output_path: str):
    """Generate LangSmith eval dataset creation script from a spec file."""
    with open(spec_path) as f:
        if spec_path.endswith(".json"):
            examples = json.load(f)
        else:
            # Assume JSON for simplicity — YAML support can be added
            examples = json.loads(f.read())

    # Validate
    for i, ex in enumerate(examples):
        if "question" not in ex:
            raise ValueError(f"Example {i} missing 'question'")
        if "expected_agents" not in ex:
            raise ValueError(f"Example {i} missing 'expected_agents'")

    dataset_code = f'''\
\"\"\"
Eval dataset: {name}

Generated by lg-eval-generator.py
Run this script to create the LangSmith dataset and run evaluators.
\"\"\"
from langsmith import Client, evaluate
from openevals.llm import create_llm_as_judge


# ── Dataset ────────────────────────────────────────────────────────────────

EXAMPLES = {json.dumps(examples, indent=2)}


def create_dataset(client: Client, dataset_name: str = "{name}"):
    \"\"\"Create or update the LangSmith eval dataset.\"\"\"
    try:
        dataset = client.create_dataset(
            dataset_name=dataset_name,
            description="Multi-agent routing and resolution evaluation dataset",
        )
    except Exception:
        # Dataset may already exist — use it
        dataset = [d for d in client.list_datasets() if d.name == dataset_name][0]

    inputs = [{{"question": ex["question"]}} for ex in EXAMPLES]
    outputs = [
        {{
            "expected_agents": ex.get("expected_agents", []),
            "must_mention": ex.get("must_mention", []),
        }}
        for ex in EXAMPLES
    ]

    client.create_examples(
        dataset_id=dataset.id,
        inputs=inputs,
        outputs=outputs,
    )
    print(f"Dataset '{{dataset_name}}' ready with {{len(EXAMPLES)}} examples")
    return dataset_name


# ── Evaluators ─────────────────────────────────────────────────────────────

ROUTING_QUALITY_PROMPT = """\\\\
Customer query: {inputs[question]}
Expected domains: {reference_outputs[expected_agents]}
Agent response: {outputs[final_response]}
Resolution notes: {outputs[resolution_notes]}

Rate 0.0-1.0 on whether the correct specialist agents handled the request
and the response fully addressed the customer's needs.
Return ONLY: {{"score": <float>, "reasoning": "<explanation>"}}"""

routing_judge = create_llm_as_judge(
    prompt=ROUTING_QUALITY_PROMPT,
    model="anthropic:claude-sonnet-4-5-20250929",
    feedback_key="routing_quality",
)


def resolution_coverage(inputs: dict, outputs: dict, reference_outputs: dict) -> dict:
    \"\"\"Did the agents address all parts of the customer's request?\"\"\"
    text = outputs.get("final_response", "").lower()
    notes = " ".join(outputs.get("resolution_notes", [])).lower()
    combined = text + " " + notes
    must_mention = reference_outputs.get("must_mention", [])
    hits = sum(1 for t in must_mention if t.lower() in combined)
    return {{
        "key": "resolution_coverage",
        "score": hits / len(must_mention) if must_mention else 1.0,
    }}


def agent_routing_accuracy(inputs: dict, outputs: dict, reference_outputs: dict) -> dict:
    \"\"\"Were the correct agents invoked?\"\"\"
    notes = " ".join(outputs.get("resolution_notes", [])).lower()
    expected = reference_outputs.get("expected_agents", [])
    hits = sum(1 for agent in expected if agent.lower() in notes)
    return {{
        "key": "routing_accuracy",
        "score": hits / len(expected) if expected else 1.0,
    }}


# ── Target Functions ───────────────────────────────────────────────────────

def supervisor_target(inputs: dict) -> dict:
    \"\"\"Invoke supervisor-pattern graph and return results for eval.

    Replace this with your actual graph invocation.
    \"\"\"
    from langchain_core.messages import HumanMessage

    # Import your graph — adjust the import path
    # from my_project.supervisor_graph import supervisor_graph

    # result = supervisor_graph.invoke({{
    #     "messages": [HumanMessage(content=inputs["question"])],
    #     "current_agent": "",
    #     "resolution_notes": [],
    # }})
    # return {{
    #     "final_response": result["messages"][-1].content,
    #     "resolution_notes": result.get("resolution_notes", []),
    # }}

    raise NotImplementedError("Replace with your supervisor graph invocation")


def swarm_target(inputs: dict) -> dict:
    \"\"\"Invoke swarm-pattern graph and return results for eval.

    Replace this with your actual graph invocation.
    \"\"\"
    from langchain_core.messages import HumanMessage

    # Import your graph — adjust the import path
    # from my_project.swarm_graph import swarm_graph

    # result = swarm_graph.invoke({{
    #     "messages": [HumanMessage(content=inputs["question"])],
    #     "current_agent": "",
    #     "resolution_notes": [],
    # }})
    # return {{
    #     "final_response": result["messages"][-1].content,
    #     "resolution_notes": result.get("resolution_notes", []),
    # }}

    raise NotImplementedError("Replace with your swarm graph invocation")


def run_evaluations():
    \"\"\"Run both supervisor and swarm evaluations for comparison.\"\"\"
    client = Client()

    dataset_name = create_dataset(client)

    print("\\\\nRunning supervisor evaluation...")
    supervisor_results = evaluate(
        supervisor_target,
        data=dataset_name,
        evaluators=[routing_judge, resolution_coverage, agent_routing_accuracy],
        experiment_prefix="supervisor-v1",
        max_concurrency=2,
        client=client,
    )

    print("\\\\nRunning swarm evaluation...")
    swarm_results = evaluate(
        swarm_target,
        data=dataset_name,
        evaluators=[routing_judge, resolution_coverage, agent_routing_accuracy],
        experiment_prefix="swarm-v1",
        max_concurrency=2,
        client=client,
    )

    print("\\\\n=== Results ===")
    print(f"Supervisor: {{supervisor_results}}")
    print(f"Swarm: {{swarm_results}}")


if __name__ == "__main__":
    import sys
    if "--compare" in sys.argv:
        run_evaluations()
    else:
        create_dataset(Client())
        print("Dataset created. Run with --compare to evaluate patterns.")
'''

    output_path = output_path or f"{name}_eval.py"
    with open(output_path, "w") as f:
        f.write(dataset_code)

    print(f"Eval script generated: {output_path}")
    print(f"Dataset name: {name}")
    print(f"Examples: {len(examples)}")


def generate_evaluator(name: str, output_dir: str):
    """Generate standalone evaluator module."""
    os.makedirs(output_dir, exist_ok=True)
    path = os.path.join(output_dir, f"{name}_evaluators.py")

    code = '''\
"""Multi-agent evaluators for routing accuracy and resolution coverage."""

from typing import Any, Dict


def resolution_coverage(
    inputs: Dict[str, Any],
    outputs: Dict[str, Any],
    reference_outputs: Dict[str, Any],
) -> Dict[str, Any]:
    """Evaluate whether all required topics were addressed in the response.

    Args:
        inputs: The original input (e.g., {"question": "..."})
        outputs: The graph output (e.g., {"final_response": "...", "resolution_notes": [...]})
        reference_outputs: Expected results (e.g., {"must_mention": ["...", "..."]})

    Returns:
        {"key": "resolution_coverage", "score": 0.0-1.0}
    """
    text = outputs.get("final_response", "").lower()
    notes = " ".join(outputs.get("resolution_notes", [])).lower()
    combined = text + " " + notes

    must_mention = reference_outputs.get("must_mention", [])
    if not must_mention:
        return {"key": "resolution_coverage", "score": 1.0}

    hits = sum(1 for term in must_mention if term.lower() in combined)
    return {"key": "resolution_coverage", "score": hits / len(must_mention)}


def agent_routing_accuracy(
    inputs: Dict[str, Any],
    outputs: Dict[str, Any],
    reference_outputs: Dict[str, Any],
) -> Dict[str, Any]:
    """Evaluate whether the correct specialist agents were invoked.

    Args:
        inputs: The original input
        outputs: Graph output with resolution_notes containing "AgentName: ..." entries
        reference_outputs: Expected agents (e.g., {"expected_agents": ["billing", "tech"]})

    Returns:
        {"key": "routing_accuracy", "score": 0.0-1.0}
    """
    notes = " ".join(outputs.get("resolution_notes", [])).lower()
    expected = reference_outputs.get("expected_agents", [])
    if not expected:
        return {"key": "routing_accuracy", "score": 1.0}

    hits = sum(1 for agent in expected if agent.lower() in notes)
    return {"key": "routing_accuracy", "score": hits / len(expected)}


def handoff_chain_length(
    inputs: Dict[str, Any],
    outputs: Dict[str, Any],
    reference_outputs: Dict[str, Any] = None,
) -> Dict[str, Any]:
    """Measure the number of handoffs in a swarm pattern execution.

    Long chains (>3) indicate routing problems or ping-pong behavior.

    Returns:
        {"key": "handoff_chain_length", "score": N} where N is the handoff count
    """
    count = outputs.get("handoff_count", 0)
    # Score = 1.0 if <= 3 handoffs, decreasing linearly after that
    score = min(1.0, 3.0 / max(count, 1)) if count > 0 else 1.0
    return {"key": "handoff_chain_length", "score": score, "info": {"count": count}}


def routing_efficiency(
    inputs: Dict[str, Any],
    outputs: Dict[str, Any],
    reference_outputs: Dict[str, Any] = None,
) -> Dict[str, Any]:
    """Evaluate routing efficiency — token cost vs accuracy tradeoff.

    Returns:
        {"key": "routing_efficiency", "score": 0.0-1.0}
    """
    # This is a placeholder — real implementation would compare token counts
    # from LangSmith traces against routing accuracy
    messages = outputs.get("messages", [])
    total_tokens = sum(
        getattr(m, "usage_metadata", {}).get("total_tokens", 0)
        for m in messages
        if hasattr(m, "usage_metadata")
    )
    return {"key": "routing_efficiency", "score": 1.0, "info": {"estimated_tokens": total_tokens}}
'''
    with open(path, "w") as f:
        f.write(code)
    print(f"Evaluator module generated: {path}")


def generate_example_spec(output_path: str):
    """Generate an example JSON spec file to fill in."""
    with open(output_path, "w") as f:
        json.dump(EXAMPLE_SPEC, f, indent=2)
    print(f"Example spec generated: {output_path}")
    print("Fill in the question/expected_agents/must_mention fields for your domain.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="LangGraph Eval Generator")
    subparsers = parser.add_subparsers(dest="mode", help="Mode: dataset, evaluator, or example-spec")

    # Dataset mode
    ds_parser = subparsers.add_parser("dataset", help="Generate eval dataset creation script")
    ds_parser.add_argument("--name", required=True, help="Dataset name")
    ds_parser.add_argument("--spec", required=True, help="Path to JSON spec file")
    ds_parser.add_argument("--output", default=None, help="Output Python file path")

    # Evaluator mode
    ev_parser = subparsers.add_parser("evaluator", help="Generate evaluator module")
    ev_parser.add_argument("--name", required=True, help="Evaluator name prefix")
    ev_parser.add_argument("--output", default=".", help="Output directory")

    # Example spec mode
    ex_parser = subparsers.add_parser("example-spec", help="Generate example JSON spec")
    ex_parser.add_argument("--output", default="eval_examples.json", help="Output path")

    args = parser.parse_args()

    if args.mode == "dataset":
        generate_dataset(args.name, args.spec, args.output)
    elif args.mode == "evaluator":
        generate_evaluator(args.name, args.output)
    elif args.mode == "example-spec":
        generate_example_spec(args.output)
    else:
        parser.print_help()