"""Evaluations Usage Examples

This example demonstrates evaluation management: creating evaluators, creating evaluations,
running them on traces or datasets, cancelling runs, and viewing results.

Evaluations apply an evaluator (LLM-as-judge, code-based, or agent-as-judge) to traces or
dataset records.
- Trace-based evaluations: Use filter_rules to select traces, can auto-run on new traces
- Dataset-based evaluations: Run on all records in a dataset

Requirements:
    pip install lumenova-beacon

    Set environment variables:
    export BEACON_ENDPOINT=https://your-endpoint.com
    export BEACON_API_KEY=your_api_key
"""

from lumenova_beacon import BeaconClient
from lumenova_beacon.evaluations import (
    Evaluation,
    EvaluationRun,
    EvaluationRunStatus,
    Evaluator,
    SourceType,
)

# Initialize the client (reads from BEACON_ENDPOINT and BEACON_API_KEY env vars)
client = BeaconClient()


# ============================================================================
# EVALUATOR MANAGEMENT
# ============================================================================

# === Create an LLM-as-Judge Evaluator ===
evaluator = Evaluator.create(
    name="Answer Quality",
    evaluator_type="llm-as-judge",
    prompt_template=(
        "Evaluate the quality of the following answer.\n\n"
        "Question: {{question}}\n"
        "Answer: {{answer}}\n\n"
        "Rate the answer quality."
    ),
    description="Evaluates answer quality on a percentage scale",
    category="quality",
    score_config={
        "type": "percent",
        "scoring_instructions": "Rate answer quality from 0 to 1",
    },
)
print(f"Created evaluator: {evaluator.name} ({evaluator.evaluator_type})")

# === Create a Code Evaluator ===
code_evaluator = Evaluator.create(
    name="Length Check",
    evaluator_type="code",
    code=(
        "def evaluate(answer: str) -> dict:\n"
        "    length = len(answer)\n"
        "    return {'score': min(length / 500, 1.0), 'reasoning': f'Length: {length} chars'}\n"
    ),
    description="Checks if the answer has sufficient length",
    category="quality",
    score_config={"type": "percent"},
)
print(f"Created code evaluator: {code_evaluator.name}")

# === List Available Evaluators ===
evaluators, pagination = Evaluator.list(page=1, page_size=20)
print(f"\nAvailable evaluators: {pagination.total}")

for ev in evaluators:
    print(f"  - {ev.name} ({ev.evaluator_type}, predefined={ev.is_predefined})")
    if ev.score_config:
        print(f"    Score type: {ev.score_config.get('type')}")
    if ev.default_config:
        print("    Has default config: Yes")

# Search evaluators by name
evaluators, _ = Evaluator.list(search="quality")

# === Get Evaluator Details ===
evaluator = Evaluator.get(evaluator.id)
print(f"\nEvaluator: {evaluator.name}")
print(f"  Type: {evaluator.evaluator_type}")
print(f"  Predefined: {evaluator.is_predefined}")
print(f"  Project ID: {evaluator.project_id}")

# LLM evaluators have prompt templates
if evaluator.prompt_template:
    print(f"  Prompt template: {evaluator.prompt_template[:100]}...")

# Code evaluators have Python code and parameters
if evaluator.code:
    print(f"  Code: {len(evaluator.code)} characters")
if evaluator.code_parameters:
    print(f"  Parameters: {[p['name'] for p in evaluator.code_parameters]}")

# === Update Evaluator ===
evaluator.update(
    description="Updated description for answer quality evaluation",
    score_config={
        "type": "percent",
        "scoring_instructions": "Rate from 0 (terrible) to 1 (perfect)",
    },
)
print(f"\nUpdated evaluator: {evaluator.description}")


# ============================================================================
# EVALUATION MANAGEMENT
# ============================================================================

# === Create Trace-Based Evaluation (Full Trace) ===
# Evaluates the entire trace
trace_evaluation = Evaluation.create(
    name="Response Quality Check",
    evaluator_id=evaluator.id,
    variable_mappings={
        "question": {
            "data_mode": "raw",
            "json_path": "$.messages[0].content",
        },
        "answer": {
            "data_mode": "raw",
            "json_path": "$.messages[-1].content",
        },
    },
    filter_rules={
        "logic": "AND",
        "rules": [
            {"field": "span.name", "operator": "contains", "value": "chat"}
        ],
    },
    source_type=SourceType.FULL_TRACE,
    active=True,  # Auto-run on new matching traces
    description="Evaluates chat response quality",
)
print(f"\nCreated trace evaluation: {trace_evaluation.name}")
print(f"  Source type: {trace_evaluation.source_type}")
print(f"  Active: {trace_evaluation.active}")

# === Create Trace-Based Evaluation (Span-Level) ===
# Evaluates specific spans matching filter criteria
span_evaluation = Evaluation.create(
    name="LLM Output Quality",
    evaluator_id=evaluator.id,
    variable_mappings={
        "question": {
            "data_mode": "reduced",
            "reduced_fields": ["input"],
        },
        "answer": {
            "data_mode": "reduced",
            "reduced_fields": ["output"],
        },
    },
    filter_rules={"logic": "AND", "rules": []},
    source_type=SourceType.SPANS,
    source_config={
        "span_type": "llm",
        "match_mode": "first",  # 'first' or 'all'
    },
)
print(f"\nCreated span evaluation: {span_evaluation.name}")
print(f"  Source type: {span_evaluation.source_type}")
print(f"  Source config: {span_evaluation.source_config}")

# === Create Evaluation with LLM Config Override ===
# Override the project's default LLM config for this evaluation
evaluation_with_config = Evaluation.create(
    name="Custom LLM Evaluation",
    evaluator_id=evaluator.id,
    variable_mappings={"question": {"data_mode": "raw", "json_path": "$.input"}},
    filter_rules={"logic": "AND", "rules": []},
    source_type="full_trace",
    llm_config_id="your-llm-config-uuid",  # Override project default
)
print(f"\nCreated evaluation with LLM config: {evaluation_with_config.llm_config_name}")

# === Create Dataset-Based Evaluation ===
dataset_evaluation = Evaluation.create(
    name="Dataset QA Evaluation",
    evaluator_id=evaluator.id,
    variable_mappings={
        "question": {"column": "input"},
        "answer": {"column": "output"},
        "expected": {"column": "ground_truth"},
    },
    dataset_id="your-dataset-uuid",
    result_dataset_id="your-result-dataset-uuid",  # Optional: merge with experiment results
)
print(f"\nCreated dataset evaluation: {dataset_evaluation.name}")


# === Get Evaluation with Statistics ===
evaluation = Evaluation.get(
    trace_evaluation.id,
    include_evaluator=True,
    include_statistics=True,
)

if evaluation.statistics:
    stats = evaluation.statistics
    print("\nEvaluation statistics:")
    print(f"  Total runs: {stats['total_runs']}")
    print(f"  Completed: {stats['completed_runs']}")
    print(f"  Failed: {stats['failed_runs']}")
    print(f"  Pending: {stats['pending_runs']}")
    print(f"  Average score: {stats.get('avg_score')}")
    print(f"  Category counts: {stats.get('category_counts')}")


# === List Evaluations ===
evaluations, pagination = Evaluation.list(page=1, page_size=20)
print(f"\nTotal evaluations: {pagination.total}")

for ev in evaluations:
    eval_type = "trace" if ev.filter_rules else "dataset"
    print(f"  - {ev.name} ({eval_type}, active={ev.active}, source={ev.source_type})")

# Filter by evaluator
evaluations, _ = Evaluation.list(evaluator_id=evaluator.id)

# Filter by dataset
evaluations, _ = Evaluation.list(dataset_id="your-dataset-uuid")

# Search by name
evaluations, _ = Evaluation.list(search="quality")

# Include evaluator details and statistics
evaluations, _ = Evaluation.list(include_evaluator=True, include_statistics=True)


# === Update Evaluation ===
trace_evaluation.update(
    name="Updated Evaluation Name",
    description="Updated description",
    active=False,
    source_type="spans",
    source_config={"span_type": "llm", "match_mode": "all"},
)
print(f"\nUpdated evaluation: {trace_evaluation.name}")
print(f"  Source type: {trace_evaluation.source_type}")


# ============================================================================
# RUNNING EVALUATIONS
# ============================================================================

# === Run Evaluation on Single Trace ===
run = trace_evaluation.run(trace_id="your-trace-id")
print(f"\nCreated run: {run.id}")
print(f"  Status: {run.status.value}")

# Run on a specific span within a trace (for span-level evaluations)
run = span_evaluation.run(trace_id="your-trace-id", span_id="your-span-id")
print(f"  Span-level run: {run.id}")

# For dataset evaluations, use dataset_record_id:
# run = dataset_evaluation.run(dataset_record_id="your-record-id")


# === Execute Evaluation on All Targets ===
result = trace_evaluation.execute()
print("\nExecution result:")
print(f"  Total targets: {result['total_targets']}")
print(f"  Runs created: {result['runs_created']}")
print(f"  Message: {result['message']}")

# Execute with time filters (for trace-based evaluations)
result = trace_evaluation.execute(
    start_time_from="2024-01-01T00:00:00Z",
    start_time_to="2024-12-31T23:59:59Z",
)
print(f"\nFiltered execution: {result['runs_created']} runs created")


# ============================================================================
# VIEWING AND MANAGING RUNS
# ============================================================================

# === List Evaluation Runs ===
runs, pagination = trace_evaluation.list_runs(page=1, page_size=20)
print(f"\nEvaluation runs: {pagination.total}")

for run in runs:
    score_info = run.score_data if run.score_data else {"value": run.score}
    print(f"  - {run.id}: {run.status.value}, score={score_info.get('value')}")

# Sort runs by creation time
runs, _ = trace_evaluation.list_runs(sort_by="created_at", sort_order="desc")

# Filter by status
completed_runs, _ = trace_evaluation.list_runs(status=EvaluationRunStatus.COMPLETED)
print(f"\nCompleted runs: {len(completed_runs)}")

for run in completed_runs:
    print(f"  Score: {run.score_data}")
    if run.result:
        print(f"  Result: {run.result}")

# Get failed runs to investigate errors
failed_runs, _ = trace_evaluation.list_runs(status=EvaluationRunStatus.FAILED)
for run in failed_runs:
    print(f"  Error: {run.error_message}")


# === Get Evaluation Run Details ===
if runs:
    run = EvaluationRun.get(runs[0].id)
    print("\nRun details:")
    print(f"  ID: {run.id}")
    print(f"  Evaluation: {run.evaluation_name}")
    print(f"  Status: {run.status.value}")
    print(f"  Score: {run.score}")
    print(f"  Score data: {run.score_data}")
    print(f"  Trace ID: {run.trace_id}")
    print(f"  Span ID: {run.span_id}")
    print(f"  Created: {run.created_at}")
    print(f"  Completed: {run.completed_at}")


# === List All Runs (across evaluations) ===
all_runs, pagination = EvaluationRun.list(page=1, page_size=50)
print(f"\nAll evaluation runs: {pagination.total}")

# Filter and sort
runs, _ = EvaluationRun.list(evaluation_id=trace_evaluation.id)
runs, _ = EvaluationRun.list(trace_id="your-trace-id")
runs, _ = EvaluationRun.list(status=EvaluationRunStatus.COMPLETED)
runs, _ = EvaluationRun.list(sort_by="created_at", sort_order="desc")


# ============================================================================
# CANCELLING RUNS
# ============================================================================

# === Cancel a Single Run ===
if runs:
    pending_run = runs[0]
    pending_run.cancel()
    print(f"\nCancelled run: {pending_run.id}, status={pending_run.status.value}")

# === Bulk Cancel Runs ===
run_ids = [r.id for r in runs[:5]]
result = EvaluationRun.bulk_cancel(run_ids)
print(f"Bulk cancelled: {result['cancelled_count']} runs")

# === Cancel All Runs for an Evaluation ===
result = trace_evaluation.cancel_runs()
print(f"Cancelled all runs: {result['cancelled_count']} runs cancelled")


# ============================================================================
# CLEANUP
# ============================================================================

# === Bulk Delete Evaluation Runs ===
if runs:
    run_ids = [run.id for run in runs[:2]]
    result = EvaluationRun.bulk_delete(run_ids)
    print(f"\nBulk deleted runs: {result}")

# === Bulk Delete Evaluations ===
result = Evaluation.bulk_delete([trace_evaluation.id, dataset_evaluation.id])
print(f"Bulk deleted evaluations: {result}")

# === Delete Single Evaluation ===
# evaluation.delete()

# === Delete Evaluator ===
# Note: predefined evaluators cannot be deleted
code_evaluator.delete()
print(f"Deleted evaluator: {code_evaluator.name}")
