"""Experiments Usage Examples

This example demonstrates experiment management: creating pipeline experiments
with steps (prompt, LLM, evaluation, external agent) against datasets,
monitoring progress, viewing results, and managing experiment lifecycle.

Requirements:
    pip install lumenova-beacon

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

import time

from lumenova_beacon import BeaconClient
from lumenova_beacon.datasets import Dataset
from lumenova_beacon.experiments import (
    Experiment,
    ExperimentRunStatus,
    ExperimentStatus,
    ExperimentStep,
    ExperimentStepType,
)
from lumenova_beacon.llm_configs import LLMConfig

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


# === List All Experiments ===
experiments, pagination = Experiment.list(page=1, page_size=10)
print(f"Total experiments: {pagination.total}")

for exp in experiments:
    print(f"  - {exp.name}")
    print(f"    ID: {exp.id}")
    print(f"    Status: {exp.status.name}")
    print(f"    Records: {exp.record_count}, Steps: {exp.step_count}")


# === Filter by Status ===
running, _ = Experiment.list(status=ExperimentStatus.RUNNING)
print(f"\nRunning experiments: {len(running)}")

completed, _ = Experiment.list(status=ExperimentStatus.COMPLETED)
print(f"Completed experiments: {len(completed)}")

failed, _ = Experiment.list(status=ExperimentStatus.FAILED)
print(f"Failed experiments: {len(failed)}")


# === Search Experiments ===
search_results, _ = Experiment.list(search="comparison")
print(f"\nExperiments matching 'comparison': {len(search_results)}")


# === Get Experiment by ID ===
if experiments:
    experiment = Experiment.get(experiments[0].id)
    print("\nExperiment details:")
    print(f"  Name: {experiment.name}")
    print(f"  ID: {experiment.id}")
    print(f"  Status: {experiment.status.name}")
    print(f"  Dataset: {experiment.dataset_name} ({experiment.dataset_id})")
    print(f"  Records: {experiment.record_count}")
    print(f"  Steps: {experiment.step_count}")
    if experiment.description:
        print(f"  Description: {experiment.description}")
    if experiment.steps:
        print("  Pipeline steps:")
        for step in experiment.steps:
            print(f"    {step.position}. [{step.step_type}] -> {step.output_column}")
    print(f"  Created: {experiment.created_at}")
    if experiment.started_at:
        print(f"  Started: {experiment.started_at}")
    if experiment.completed_at:
        print(f"  Completed: {experiment.completed_at}")


# === List Available LLM Configs ===
llm_configs = LLMConfig.list()
print(f"\nAvailable LLM configs: {len(llm_configs)}")
for cfg in llm_configs:
    print(f"  - {cfg.name}: {cfg.provider}/{cfg.litellm_model}")


# === Create New Experiment (Pipeline) ===
# Experiments use pipeline steps that execute sequentially per record.
datasets, _ = Dataset.list(page=1, page_size=1)

if datasets and datasets[0].id and llm_configs:
    dataset = datasets[0]
    llm = llm_configs[0]
    print(f"\nUsing dataset: {dataset.name} (ID: {dataset.id})")
    print(f"Using LLM config: {llm.name}")

    # Create a pipeline with prompt rendering + LLM generation
    experiment = Experiment.create(
        name="Pipeline Experiment",
        dataset_id=str(dataset.id),
        description="Render prompts then generate with LLM",
        steps=[
            ExperimentStep(
                step_type=ExperimentStepType.PROMPT,
                output_column="rendered_prompt",
                config={
                    "prompt_content": "Summarize the following text: {{input}}",
                    "variable_mappings": {"input": "text_column"},
                },
            ),
            ExperimentStep(
                step_type=ExperimentStepType.LLM,
                output_column="llm_response",
                config={
                    "model_config_id": str(llm.id),
                    "model_parameters": {"temperature": 0.7, "max_tokens": 200},
                    "variable_mappings": {"input": "rendered_prompt"},
                },
            ),
        ],
    )

    print(f"\nCreated experiment: {experiment.name}")
    print(f"  ID: {experiment.id}")
    print(f"  Status: {experiment.status.name}")
    print(f"  Steps: {experiment.step_count}")


# === Update Draft Experiment ===
# Modify steps, dataset, or metadata before starting
if experiment.status == ExperimentStatus.DRAFT:
    experiment.update_draft(
        name="Updated Pipeline Experiment",
        description="Updated description",
    )
    print(f"\nUpdated draft: {experiment.name}")


# === Preview Output ===
# Test configurations against sample records without creating an experiment
if datasets and llm_configs:
    dataset = datasets[0]
    records, _ = dataset.list_records(page=1, page_size=2)
    if records:
        preview = Experiment.preview_output(
            dataset_id=str(dataset.id),
            configurations=[
                {
                    "label": "A",
                    "prompt_content": {"text": "Summarize: {{input}}"},
                    "model_config_id": str(llm_configs[0].id),
                    "model_parameters": {"temperature": 0.7},
                },
            ],
            record_ids=[str(r.id) for r in records[:3]],
            variable_mappings={"input": "text_column"},
        )
        print(f"\nPreview results: {len(preview['results'])} runs")
        print(f"Total time: {preview['total_execution_time_ms']}ms")
        for result in preview["results"]:
            print(f"  [{result['configuration_label']}] {result['status']}: "
                  f"{result.get('output', '')[:50]}...")


# === Start Experiment ===
# Experiments are created in DRAFT status - start to begin execution
experiment.start()
print(f"\nExperiment started! Status: {experiment.status.name}")


# === Monitor Progress ===
while experiment.status in [ExperimentStatus.QUEUED, ExperimentStatus.RUNNING]:
    progress = experiment.get_progress()
    print(f"Progress: {progress['progress_percent']:.1f}% "
          f"({progress['completed_runs']}/{progress['total_runs']})")
    experiment = Experiment.get(experiment.id)  # Refresh status
    time.sleep(5)
print(f"Finished! Status: {experiment.status.name}")


# === View Run Statistics ===
stats = experiment.get_run_stats()
print(f"\nRun statistics:")
print(f"  Total: {stats.total_runs}")
print(f"  Completed: {stats.completed_runs}")
print(f"  Failed: {stats.failed_runs}")
if stats.avg_execution_time_ms:
    print(f"  Avg time: {stats.avg_execution_time_ms:.0f}ms")


# === View Enriched Records ===
# Records with per-step execution details
records, pagination = experiment.list_records(page=1, page_size=5)
print(f"\nRecords ({pagination.total} total):")
for record in records:
    print(f"  Record {record.record_id}: {record.overall_status}")
    for step_run in record.step_runs:
        status_icon = "OK" if step_run.status == ExperimentRunStatus.COMPLETED else "FAIL"
        print(f"    Step {step_run.step_position} [{step_run.step_type}] "
              f"-> {step_run.output_column}: {status_icon}")
        if step_run.output:
            print(f"      Output: {str(step_run.output)[:80]}...")


# === View Individual Runs ===
runs, _ = experiment.list_runs(page=1, page_size=5, status=ExperimentRunStatus.FAILED)
print(f"\nFailed runs: {len(runs)}")
for run in runs:
    print(f"  Run {run.id}: {run.error}")


# === Get Source & Result Datasets ===
source_dataset = experiment.get_dataset()
print(f"\nSource dataset: {source_dataset.name}")

if experiment.status == ExperimentStatus.COMPLETED and experiment.result_dataset_id:
    result_dataset = experiment.get_result_dataset()
    if result_dataset:
        print(f"Result dataset: {result_dataset.name} ({result_dataset.record_count} records)")


# === Retry Failed Runs ===
if experiment.status == ExperimentStatus.COMPLETED_WITH_ERRORS:
    result = experiment.retry_failed()
    print(f"\nRetried {result['retried_count']} failed runs")


# === Reset Experiment ===
# Reset a completed/failed experiment back to DRAFT to re-run with changes
if experiment.status in {
    ExperimentStatus.COMPLETED,
    ExperimentStatus.FAILED,
    ExperimentStatus.COMPLETED_WITH_ERRORS,
    ExperimentStatus.STOPPED,
}:
    experiment.reset()
    print(f"\nReset experiment to: {experiment.status.name}")
    # Can now update draft and start again
    experiment.start()


# === Continue Stopped Experiment ===
# Resume a stopped experiment from where it left off (only incomplete records re-run)
if experiment.status == ExperimentStatus.STOPPED:
    experiment.continue_experiment()
    print(f"\nContinuing experiment! Status: {experiment.status.name}")

# For experiments with external agent steps, pass external_agents to continue:
# experiment.continue_experiment(external_agents={"My Agent": my_agent})

# continue_experiment() also works on WAITING_FOR_EXTERNAL experiments,
# equivalent to calling start() — useful to resume the agent polling loop:
# if experiment.status == ExperimentStatus.WAITING_FOR_EXTERNAL:
#     experiment.continue_experiment(
#         external_agents={"My Agent": my_agent_fn},
#         trace_calls=True,
#     )


# === Clone Experiment ===
# Create a copy in DRAFT status
cloned = experiment.clone()
print(f"\nCloned experiment: {cloned.name} (ID: {cloned.id})")
print(f"  Status: {cloned.status.name}")


# === Stop Experiment ===
# Stop a running or queued experiment (can be continued later)
# running, _ = Experiment.list(status=ExperimentStatus.RUNNING)
# if running:
#     exp = running[0]
#     exp.cancel()
#     print(f"Stopped experiment: {exp.name}")
#     # Resume later with exp.continue_experiment()


# === Update Experiment Metadata ===
if experiments:
    exp = experiments[0]
    exp.update(
        name="Updated Experiment Name",
        description="Updated description",
    )
    print(f"\nUpdated experiment: {exp.name}")


# === Delete Experiment ===
# Delete an experiment (commented out to prevent accidental deletion)
# if experiments:
#     exp = experiments[0]
#     exp.delete()
#     print(f"Deleted experiment: {exp.id}")


# === External Agent Experiment ===
# Create an experiment with an external agent step that the SDK executes locally
def my_agent(input_text: str) -> str:
    """Example agent function that processes input and returns output."""
    return f"Processed: {input_text}"


if datasets and datasets[0].id:
    agent_experiment = Experiment.create(
        name="Agent Evaluation",
        dataset_id=str(datasets[0].id),
        steps=[
            ExperimentStep(
                step_type=ExperimentStepType.EXTERNAL_AGENT,
                output_column="agent_output",
                config={
                    "agent_name": "My Agent",
                    "variable_mappings": {"input": "text_column"},
                },
            ),
        ],
    )

    # start() with external_agents blocks until all matching agent steps are done
    agent_experiment.start(
        external_agents={"My Agent": my_agent},
        trace_calls=True,  # Creates tracing spans for each invocation
    )
    print(f"\nAgent experiment finished: {agent_experiment.status.name}")
