"""
Generate a template script for running agent evaluation.

This script creates a customized Python script that executes the agent
on an evaluation dataset and collects trace IDs for scoring.

Usage:
    python run_evaluation_template.py                                        # Auto-detect everything
    python run_evaluation_template.py --module my_agent.agent                # Specify module
    python run_evaluation_template.py --entry-point run_agent                # Specify entry point
    python run_evaluation_template.py --dataset-name my-dataset              # Specify dataset
    python run_evaluation_template.py --module my_agent --entry-point run_agent --dataset-name my-dataset
"""

import argparse
import os
import subprocess
import sys

from utils import validate_env_vars


def list_datasets() -> list[str]:
    """List available datasets in the experiment."""
    try:
        code = """
import os
from mlflow import MlflowClient

client = MlflowClient()
experiment_id = os.getenv("MLFLOW_EXPERIMENT_ID")

datasets = client.search_datasets(experiment_ids=[experiment_id])
for dataset in datasets:
    print(dataset.name)
"""
        result = subprocess.run(["python", "-c", code], capture_output=True, text=True, check=True)
        return [line.strip() for line in result.stdout.strip().split("\n") if line.strip()]
    except Exception:
        return []


def generate_evaluation_code(
    tracking_uri: str, experiment_id: str, dataset_name: str, agent_module: str, entry_point: str
) -> str:
    """Generate Python code for running evaluation."""

    return f'''#!/usr/bin/env python3
"""
Run agent on evaluation dataset and collect traces.

Generated by run_evaluation_template.py
"""

import os
import sys
import mlflow
from mlflow.genai.datasets import get_dataset

# Set environment variables
os.environ["MLFLOW_TRACKING_URI"] = "{tracking_uri}"
os.environ["MLFLOW_EXPERIMENT_ID"] = "{experiment_id}"

# Import agent
from {agent_module} import {entry_point}

# Configuration
DATASET_NAME = "{dataset_name}"

print("=" * 60)
print("Running Agent on Evaluation Dataset")
print("=" * 60)
print()

# Load dataset
# IMPORTANT: Do not modify this section. It uses the official MLflow API.
# Spark or databricks-sdk approaches are NOT recommended.
print("Loading evaluation dataset...")
try:
    dataset = get_dataset(DATASET_NAME)
    df = dataset.to_df()
    print(f"  Dataset: {{DATASET_NAME}}")
    print(f"  Total queries: {{len(df)}}")
    print()
except Exception as e:
    print(f"✗ Failed to load dataset: {{e}}")
    print()
    print("Common issues:")
    print("  1. Dataset name incorrect - check with: mlflow datasets list")
    print("  2. Not authenticated - run: databricks auth login")
    print("  3. Wrong experiment - verify MLFLOW_EXPERIMENT_ID")
    sys.exit(1)

# TODO: Configure your agent's LLM provider or other dependencies here
# Example:
# from your_agent.llm import LLMConfig, LLMProvider
# llm_config = LLMConfig(model="gpt-4", temperature=0.0)
# llm_provider = LLMProvider(config=llm_config)

print("⚠ IMPORTANT: Configure your agent's dependencies above before running!")
print("  Update the TODO section with your agent's setup code")
print()

# Run agent on each query
trace_ids = []
successful = 0
failed = 0

print("Running agent on dataset queries...")
print()

for index, row in df.iterrows():
    inputs = row['inputs']

    # Extract query from inputs
    query = inputs.get('query', inputs.get('question', str(inputs)))

    print(f"[{{index + 1}}/{{len(df)}}] Query: {{query[:80]}}{{'...' if len(query) > 80 else ''}}")

    try:
        # TODO: Adjust the function call to match your agent's signature
        # Examples:
        #   response = {entry_point}(query, llm_provider)
        #   response = {entry_point}(query)
        #   response = {entry_point}(**inputs)

        response = {entry_point}(query)  # <-- UPDATE THIS LINE

        # Capture trace ID
        trace_id = mlflow.get_last_active_trace_id()

        if trace_id:
            trace_ids.append(trace_id)
            successful += 1
            print(f"  ✓ Success (trace: {{trace_id}})")
        else:
            print(f"  ✗ No trace captured")
            failed += 1

    except Exception as e:
        print(f"  ✗ Error: {{str(e)[:100]}}")
        failed += 1

    print()

# Summary
print("=" * 60)
print("Execution Summary")
print("=" * 60)
print(f"  Total queries: {{len(df)}}")
print(f"  Successful: {{successful}}")
print(f"  Failed: {{failed}}")
print(f"  Traces collected: {{len(trace_ids)}}")
print()

# Save trace IDs
if trace_ids:
    traces_file = "evaluation_trace_ids.txt"
    with open(traces_file, 'w') as f:
        f.write(','.join(trace_ids))

    print(f"Trace IDs saved to: {{traces_file}}")
    print()

    # Print evaluation command
    print("=" * 60)
    print("Next Step: Evaluate Traces with Scorers")
    print("=" * 60)
    print()
    print("Run the following command to evaluate all traces:")
    print()
    print(f"  mlflow traces evaluate \\\\")
    print(f"    --trace-ids {{','.join(trace_ids[:3])}}{{',...' if len(trace_ids) > 3 else ''}} \\\\")
    print(f"    --scorers <scorer1>,<scorer2>,... \\\\")
    print(f"    --output json")
    print()
    print("Replace <scorer1>,<scorer2>,... with your registered scorers")
    print("  Example: RelevanceToQuery,Completeness,ToolUsageAppropriate")
    print()
else:
    print("✗ No traces were collected. Please check for errors above.")
    print()

print("=" * 60)
'''


def main():
    """Main workflow."""
    # Parse command-line arguments
    parser = argparse.ArgumentParser(
        description="Generate evaluation execution template script",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--module", help="Agent module name (e.g., 'my_agent.agent')")
    parser.add_argument("--entry-point", help="Entry point function name (e.g., 'run_agent')")
    parser.add_argument("--dataset-name", help="Dataset name to use")
    parser.add_argument("--output", default="run_agent_evaluation.py", help="Output file name")
    args = parser.parse_args()

    print("=" * 60)
    print("MLflow Evaluation Execution Template Generator")
    print("=" * 60)
    print()

    # Check environment
    errors = validate_env_vars()
    if errors:
        print("✗ Environment validation failed:")
        for error in errors:
            print(f"  - {error}")
        print("\nRun scripts/setup_mlflow.py first")
        sys.exit(1)

    tracking_uri = os.getenv("MLFLOW_TRACKING_URI")
    experiment_id = os.getenv("MLFLOW_EXPERIMENT_ID")

    print(f"Tracking URI: {tracking_uri}")
    print(f"Experiment ID: {experiment_id}")
    print()

    # Get agent module (must be specified manually)
    print("Agent module configuration...")
    agent_module = args.module
    if not agent_module:
        print("  ✗ Agent module not specified")
        print("  Use --module to specify your agent module")
        print("  Example: --module my_agent.agent")
        print("\n  To find your agent module:")
        print("    grep -r 'def.*agent' . --include='*.py'")
        sys.exit(1)
    else:
        print(f"  ✓ Using specified: {agent_module}")

    # Get entry point (must be specified manually)
    print("\nEntry point configuration...")
    entry_point = args.entry_point
    if not entry_point:
        print("  ✗ Entry point not specified")
        print("  Use --entry-point to specify your agent's main function")
        print("  Example: --entry-point run_agent")
        print("\n  To find entry points with @mlflow.trace:")
        print("    grep -r '@mlflow.trace' . --include='*.py'")
        sys.exit(1)
    else:
        print(f"  ✓ Using specified: {entry_point}")

    # Get dataset name
    print("\nFetching available datasets...")
    dataset_name = args.dataset_name
    if not dataset_name:
        datasets = list_datasets()

        if datasets:
            print(f"\n✓ Found {len(datasets)} dataset(s):")
            for i, name in enumerate(datasets, 1):
                print(f"  {i}. {name}")

            # Auto-select first dataset
            dataset_name = datasets[0]
            print(f"\n✓ Auto-selected: {dataset_name}")
            print("  (Use --dataset-name to specify a different dataset)")
        else:
            print("  ✗ No datasets found")
            print("  Please create a dataset first or specify with --dataset-name")
            sys.exit(1)
    else:
        print(f"  ✓ Using specified: {dataset_name}")

    # Generate code
    print("\n" + "=" * 60)
    print("Generating Evaluation Execution Script")
    print("=" * 60)

    code = generate_evaluation_code(
        tracking_uri, experiment_id, dataset_name, agent_module, entry_point
    )

    # Write to file
    output_file = args.output
    with open(output_file, "w") as f:
        f.write(code)

    print(f"\n✓ Script generated: {output_file}")
    print()

    # Make executable
    try:
        os.chmod(output_file, 0o755)
        print(f"✓ Made executable: chmod +x {output_file}")
    except Exception:
        pass

    print()
    print("=" * 60)
    print("Next Steps")
    print("=" * 60)
    print()
    print(f"1. Review the generated script: {output_file}")
    print("2. Update the TODO sections with your agent's setup code")
    print("3. Update the agent call to match your signature")
    print(f"4. Execute it: python {output_file}")
    print("5. Use the trace IDs to run evaluation with scorers")
    print()
    print("=" * 60)


if __name__ == "__main__":
    main()
