use clap::{Args, ValueEnum};
use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};

use homeboy::component::{Component, ScopedExtensionConfig};
use homeboy::engine::baseline::BaselineFlags;
use homeboy::engine::execution_context::{self, ResolveOptions};
use homeboy::engine::run_dir::RunDir;
use homeboy::extension::trace as extension_trace;
use homeboy::extension::trace::{
    TraceCommandOutput, TraceListWorkflowArgs, TraceRunWorkflowArgs, TraceSpanDefinition,
};
use homeboy::extension::ExtensionCapability;
use homeboy::observation::{
    NewRunRecord, NewTraceRunRecord, NewTraceSpanRecord, ObservationStore, RunStatus,
};
use homeboy::rig::{self, RigSpec};

use super::utils::args::{BaselineArgs, HiddenJsonArgs, PositionalComponentArgs, SettingArgs};
use super::{CmdResult, GlobalArgs};

mod output;

use output::{
    aggregate_span, render_aggregate_markdown, render_compare_markdown, run_compare,
    TraceAggregateSpanSample,
};

#[cfg(test)]
use output::{
    compare_trace_aggregates, compare_trace_aggregates_with_focus, parse_trace_aggregate_input,
    TraceAggregateInput, TraceAggregateSpanInput,
};

#[derive(Args, Clone)]
pub struct TraceArgs {
    #[command(flatten)]
    comp: PositionalComponentArgs,

    /// Scenario ID to run, or `list` to discover available scenarios.
    pub scenario: String,

    /// After aggregate JSON when running `homeboy trace compare before.json after.json`.
    #[arg(value_name = "AFTER_JSON")]
    pub compare_after: Option<PathBuf>,

    /// Run trace against a rig-pinned component path after `rig check` passes.
    #[arg(long, value_name = "RIG_ID")]
    pub rig: Option<String>,

    #[command(flatten)]
    pub setting_args: SettingArgs,

    #[command(flatten)]
    pub _json: HiddenJsonArgs,

    /// Print compact machine-readable summary.
    #[arg(long)]
    pub json_summary: bool,

    /// Render a Markdown trace report instead of the JSON envelope.
    #[arg(long, value_parser = ["markdown"])]
    pub report: Option<String>,

    /// Run the same trace scenario multiple times.
    #[arg(long, value_name = "N", default_value_t = 1)]
    pub repeat: usize,

    /// Aggregate repeated trace output.
    #[arg(long, value_parser = ["spans"])]
    pub aggregate: Option<String>,

    /// Run order for repeated trace executions.
    #[arg(long, value_enum, default_value_t = TraceSchedule::Grouped)]
    pub schedule: TraceSchedule,

    /// Highlight a span in aggregate and compare reports. Repeatable.
    #[arg(long = "focus-span", value_name = "SPAN_ID")]
    pub focus_spans: Vec<String>,

    /// Add a span definition as `id:source.event:source.event`.
    #[arg(long = "span", value_name = "ID:FROM:TO", value_parser = extension_trace::spans::parse_span_definition)]
    pub spans: Vec<TraceSpanDefinition>,

    /// Add an ordered phase milestone as `[label:]source.event`.
    #[arg(long = "phase", value_name = "[LABEL:]SOURCE.EVENT", value_parser = extension_trace::spans::parse_phase_milestone)]
    pub phases: Vec<extension_trace::spans::TracePhaseMilestone>,

    /// Use a named phase preset declared by the selected rig/workload.
    #[arg(long = "phase-preset", value_name = "NAME")]
    pub phase_preset: Option<String>,

    #[command(flatten)]
    pub baseline_args: BaselineArgs,

    /// Span regression tolerance as a percentage.
    #[arg(long, value_name = "PERCENT", default_value_t = extension_trace::baseline::DEFAULT_REGRESSION_THRESHOLD_PERCENT)]
    pub regression_threshold: f64,

    /// Minimum span slowdown in milliseconds before a regression can fail.
    #[arg(long, value_name = "MS", default_value_t = extension_trace::baseline::DEFAULT_REGRESSION_MIN_DELTA_MS)]
    pub regression_min_delta_ms: u64,

    /// Apply a patch file for this trace run, then reverse it afterward.
    #[arg(long = "overlay", value_name = "PATCH_FILE")]
    pub overlays: Vec<String>,

    /// Leave overlay changes in place after the trace run.
    #[arg(long)]
    pub keep_overlay: bool,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)]
pub enum TraceSchedule {
    Grouped,
    Interleaved,
}

impl TraceSchedule {
    fn as_str(self) -> &'static str {
        match self {
            Self::Grouped => "grouped",
            Self::Interleaved => "interleaved",
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
struct TraceRunPlanEntry {
    index: usize,
    group: String,
    iteration: usize,
}

fn plan_trace_run_order(
    repeat: usize,
    schedule: TraceSchedule,
    groups: &[&str],
) -> Vec<TraceRunPlanEntry> {
    let mut entries = Vec::new();
    let mut push_entry = |group: &str, iteration: usize| {
        entries.push(TraceRunPlanEntry {
            index: entries.len() + 1,
            group: group.to_string(),
            iteration,
        });
    };
    match schedule {
        TraceSchedule::Grouped => {
            for group in groups {
                for iteration in 1..=repeat {
                    push_entry(group, iteration);
                }
            }
        }
        TraceSchedule::Interleaved => {
            for iteration in 1..=repeat {
                for group in groups {
                    push_entry(group, iteration);
                }
            }
        }
    }
    entries
}

pub fn is_markdown_mode(args: &TraceArgs) -> bool {
    args.report.as_deref() == Some("markdown")
}

pub fn run_markdown(args: TraceArgs, global: &GlobalArgs) -> CmdResult<String> {
    let (output, exit_code) = run(args, global)?;
    match output {
        TraceCommandOutput::Run(run_output) => {
            let Some(results) = run_output.results else {
                return Ok(("# Trace\n\nNo trace results were produced.\n".to_string(), exit_code));
            };
            Ok((
                extension_trace::render_markdown(&results, &run_output.overlays),
                exit_code,
            ))
        }
        TraceCommandOutput::Summary(summary) => Ok((
            format!(
                "# Trace Summary\n\n- **Component:** `{}`\n- **Status:** `{}`\n- **Exit code:** `{}`\n",
                summary.component, summary.status, summary.exit_code
            ),
            exit_code,
        )),
        TraceCommandOutput::Aggregate(aggregate) => {
            Ok((render_aggregate_markdown(&aggregate), exit_code))
        }
        TraceCommandOutput::Compare(compare) => Ok((render_compare_markdown(&compare), exit_code)),
        TraceCommandOutput::List(list) => {
            let mut markdown = format!("# Trace Scenarios: `{}`\n\n", list.component_id);
            for scenario in list.scenarios {
                markdown.push_str(&format!("- `{}`", scenario.id));
                if let Some(summary) = scenario.summary {
                    markdown.push_str(&format!(": {}", summary));
                }
                markdown.push('\n');
            }
            Ok((markdown, exit_code))
        }
    }
}

pub fn run(args: TraceArgs, _global: &GlobalArgs) -> CmdResult<TraceCommandOutput> {
    let ((stdout_output, _artifact_output), exit_code) = run_outputs(args)?;
    Ok((stdout_output, exit_code))
}

pub fn run_json_with_output_artifact(
    args: TraceArgs,
    _global: &GlobalArgs,
) -> (
    homeboy::Result<serde_json::Value>,
    i32,
    Option<homeboy::Result<serde_json::Value>>,
) {
    crate::commands::utils::tty::status("homeboy is working...");
    let output_to_json = |output: TraceCommandOutput| {
        serde_json::to_value(output).map_err(|err| {
            homeboy::Error::internal_json(err.to_string(), Some("serialize response".to_string()))
        })
    };
    match run_outputs(args) {
        Ok(((stdout_output, artifact_output), exit_code)) => (
            output_to_json(stdout_output),
            exit_code,
            artifact_output.map(output_to_json),
        ),
        Err(err) => {
            let (json_result, exit_code) = crate::commands::utils::response::map_cmd_result_to_json::<
                TraceCommandOutput,
            >(Err(err));
            (json_result, exit_code, None)
        }
    }
}

fn run_outputs(args: TraceArgs) -> CmdResult<(TraceCommandOutput, Option<TraceCommandOutput>)> {
    if args.comp.component.as_deref() == Some("compare") {
        let (output, exit_code) = run_compare(args)?;
        return Ok(((output, None), exit_code));
    }

    if args.compare_after.is_some() {
        return Err(homeboy::Error::validation_invalid_argument(
            "AFTER_JSON",
            "extra positional argument is only supported by `homeboy trace compare before.json after.json`",
            None,
            None,
        ));
    }

    if args.repeat == 0 {
        return Err(homeboy::Error::validation_invalid_argument(
            "--repeat",
            "repeat must be at least 1",
            None,
            None,
        ));
    }

    if args.scenario == "list" {
        let (output, exit_code) = run_list(args)?;
        return Ok(((output, None), exit_code));
    }

    if args.repeat > 1 || args.aggregate.as_deref() == Some("spans") {
        let (output, exit_code) = run_repeat(args)?;
        return Ok(((output, None), exit_code));
    }

    let summary_only = args.json_summary;
    let execution = execute_trace_run(args)?;

    let (stdout_output, artifact_output, exit_code) = extension_trace::from_main_workflow_outputs(
        execution.workflow,
        execution.rig_state,
        summary_only,
    );
    Ok(((stdout_output, artifact_output), exit_code))
}

struct TraceRunExecution {
    workflow: extension_trace::TraceRunWorkflowResult,
    run_dir: RunDir,
    rig_state: Option<rig::RigStateSnapshot>,
}

fn execute_trace_run(args: TraceArgs) -> homeboy::Result<TraceRunExecution> {
    let rig_context = load_rig_context(args.rig.as_deref())?;
    let effective_id = resolve_component_id(&args.comp, rig_context.as_ref().map(|c| &c.rig_spec))?;
    let path_override = args.comp.path.clone().or_else(|| {
        rig_context
            .as_ref()
            .and_then(|context| rig_component_path(&context.rig_spec, &effective_id))
    });
    let component_override = rig_context
        .as_ref()
        .and_then(|context| rig_component_for_trace(&context.rig_spec, &effective_id));

    let ctx = execution_context::resolve_with_component(
        &ResolveOptions::with_capability_and_json(
            &effective_id,
            path_override.clone(),
            ExtensionCapability::Trace,
            args.setting_args.setting.clone(),
            args.setting_args.setting_json.clone(),
        ),
        component_override,
    )?;
    if let Some(context) = rig_context.as_ref() {
        run_rig_workload_preflight(
            &context.rig_spec,
            ctx.extension_id.as_deref(),
            rig::RigWorkloadKind::Trace,
        )?;
    }
    let span_definitions = span_definitions_for_args(
        &args,
        rig_context.as_ref(),
        ctx.extension_id.as_deref(),
        args.aggregate.as_deref() == Some("spans")
            && args.phase_preset.is_none()
            && args.phases.is_empty()
            && args.spans.is_empty(),
    )?;

    let rig_state = rig_context
        .as_ref()
        .map(|context| rig::snapshot_state(&context.rig_spec));
    let run_dir = RunDir::create()?;
    let scenario_id = args.scenario.clone();
    let rig_id = args.rig.clone();
    let requested_overlays = args.overlays.clone();
    let component_path_for_observation = path_override
        .clone()
        .unwrap_or_else(|| ctx.component.local_path.clone());
    let observation = ObservationStore::open_initialized().ok().and_then(|store| {
        store
            .start_run(NewRunRecord {
                kind: "trace".to_string(),
                component_id: Some(ctx.component_id.clone()),
                command: Some(std::env::args().collect::<Vec<_>>().join(" ")),
                cwd: std::env::current_dir()
                    .ok()
                    .map(|path| path.to_string_lossy().to_string()),
                homeboy_version: Some(env!("CARGO_PKG_VERSION").to_string()),
                git_sha: homeboy::git::short_head_revision_at(Path::new(
                    &component_path_for_observation,
                )),
                rig_id: rig_id.clone(),
                metadata_json: serde_json::json!({
                    "scenario_id": scenario_id,
                    "component_path": component_path_for_observation,
                    "requested_overlays": requested_overlays,
                    "span_definitions": span_definitions.clone(),
                    "phase_preset": args.phase_preset.clone(),
                    "phase_milestones": args.phases.clone().into_iter().map(|phase| {
                        serde_json::json!({ "label": phase.label, "key": phase.key })
                    }).collect::<Vec<_>>(),
                    "baseline": {
                        "baseline": args.baseline_args.baseline,
                        "ignore_baseline": args.baseline_args.ignore_baseline,
                        "ratchet": args.baseline_args.ratchet,
                        "regression_threshold_percent": args.regression_threshold,
                        "regression_min_delta_ms": args.regression_min_delta_ms
                    }
                }),
            })
            .ok()
            .map(|run| ActiveTraceObservation {
                store,
                run_id: run.id,
                component_id: ctx.component_id.clone(),
                rig_id: rig_id.clone(),
                scenario_id: scenario_id.clone(),
            })
    });
    let extra_workloads = rig_context
        .as_ref()
        .and_then(|context| {
            ctx.extension_id.as_deref().map(|id| {
                rig::workloads_for_extension(
                    &context.rig_spec,
                    rig::RigWorkloadKind::Trace,
                    context.rig_package_root.as_deref(),
                    id,
                )
            })
        })
        .unwrap_or_default();
    let workflow = match extension_trace::run_trace_workflow(
        &ctx.component,
        TraceRunWorkflowArgs {
            component_label: effective_id.clone(),
            component_id: ctx.component_id.clone(),
            path_override,
            settings: settings_as_strings(&ctx.settings),
            settings_json: settings_as_json(&ctx.settings),
            scenario_id: args.scenario,
            json_summary: args.json_summary,
            rig_id: args.rig,
            overlays: args.overlays,
            keep_overlay: args.keep_overlay,
            extra_workloads,
            span_definitions,
            baseline_flags: BaselineFlags {
                baseline: args.baseline_args.baseline,
                ignore_baseline: args.baseline_args.ignore_baseline,
                ratchet: args.baseline_args.ratchet,
            },
            regression_threshold_percent: args.regression_threshold,
            regression_min_delta_ms: args.regression_min_delta_ms,
        },
        &run_dir,
        rig_state.clone(),
    ) {
        Ok(workflow) => workflow,
        Err(error) => {
            if let Some(observation) = observation.as_ref() {
                persist_trace_workflow_error(observation, &run_dir, &error);
            }
            return Err(error);
        }
    };
    if let Some(observation) = observation.as_ref() {
        persist_trace_workflow_result(observation, &run_dir, &workflow, rig_state.as_ref());
    }

    Ok(TraceRunExecution {
        workflow,
        run_dir,
        rig_state,
    })
}

fn run_repeat(args: TraceArgs) -> CmdResult<TraceCommandOutput> {
    let repeat = args.repeat;
    let scenario_id = args.scenario.clone();
    let mut runs = Vec::new();
    let mut overlays = Vec::new();
    let mut span_samples: BTreeMap<String, Vec<TraceAggregateSpanSample>> = BTreeMap::new();
    let mut span_failures: BTreeMap<String, usize> = BTreeMap::new();
    let mut all_span_ids: BTreeSet<String> = cli_span_definitions_for_args(&args)?
        .into_iter()
        .map(|span| span.id)
        .collect();
    let mut rig_state = None;
    let mut component = None;
    let mut failure_count = 0;

    let run_order = plan_trace_run_order(repeat, args.schedule, &["run"]);

    for plan_entry in &run_order {
        let index = plan_entry.index;
        let mut run_args = args.clone();
        run_args.repeat = 1;
        match execute_trace_run(run_args) {
            Ok(execution) => {
                if rig_state.is_none() {
                    rig_state = execution.rig_state.clone();
                }
                if component.is_none() {
                    component = Some(execution.workflow.component.clone());
                }
                if overlays.is_empty() && !execution.workflow.overlays.is_empty() {
                    overlays = execution.workflow.overlays.clone();
                }
                let passed =
                    execution.workflow.exit_code == 0 && execution.workflow.status == "pass";
                if !passed {
                    failure_count += 1;
                }
                let artifact_path = execution
                    .run_dir
                    .step_file(homeboy::engine::run_dir::files::TRACE_RESULTS)
                    .to_string_lossy()
                    .to_string();
                let mut seen_span_ids = BTreeSet::new();
                if let Some(results) = execution.workflow.results.as_ref() {
                    for span in &results.span_results {
                        all_span_ids.insert(span.id.clone());
                        seen_span_ids.insert(span.id.clone());
                        if span.status == extension_trace::parsing::TraceSpanStatus::Ok {
                            if let Some(duration) = span.duration_ms {
                                span_samples.entry(span.id.clone()).or_default().push(
                                    TraceAggregateSpanSample {
                                        duration_ms: duration,
                                        run_index: index,
                                        artifact_path: artifact_path.clone(),
                                    },
                                );
                                continue;
                            }
                        }
                        *span_failures.entry(span.id.clone()).or_default() += 1;
                    }
                    for span_id in all_span_ids.difference(&seen_span_ids) {
                        *span_failures.entry(span_id.clone()).or_default() += 1;
                    }
                } else {
                    for span_id in &all_span_ids {
                        *span_failures.entry(span_id.clone()).or_default() += 1;
                    }
                }
                runs.push(extension_trace::TraceAggregateRunOutput {
                    index,
                    passed,
                    status: execution.workflow.status,
                    exit_code: execution.workflow.exit_code,
                    artifact_path,
                    scenario_id: execution
                        .workflow
                        .results
                        .as_ref()
                        .map(|r| r.scenario_id.clone()),
                    summary: execution
                        .workflow
                        .results
                        .as_ref()
                        .and_then(|r| r.summary.clone()),
                    failure: execution
                        .workflow
                        .failure
                        .as_ref()
                        .map(|failure| failure.stderr_excerpt.clone())
                        .or_else(|| {
                            execution
                                .workflow
                                .results
                                .as_ref()
                                .and_then(|r| r.failure.clone())
                        }),
                });
            }
            Err(error) => {
                failure_count += 1;
                for span_id in &all_span_ids {
                    *span_failures.entry(span_id.clone()).or_default() += 1;
                }
                runs.push(extension_trace::TraceAggregateRunOutput {
                    index,
                    passed: false,
                    status: "error".to_string(),
                    exit_code: 1,
                    artifact_path: String::new(),
                    scenario_id: Some(scenario_id.clone()),
                    summary: None,
                    failure: Some(error.message),
                });
            }
        }
    }

    let spans = all_span_ids
        .into_iter()
        .map(|id| {
            let samples = span_samples.remove(&id).unwrap_or_default();
            let failures = span_failures.remove(&id).unwrap_or(0);
            aggregate_span(id, samples, failures)
        })
        .collect::<Vec<_>>();
    let focus_spans = focus_aggregate_spans(&spans, &args.focus_spans);
    let exit_code = if failure_count == 0 { 0 } else { 1 };
    let output = extension_trace::TraceAggregateOutput {
        command: "trace.aggregate.spans",
        passed: failure_count == 0,
        status: if failure_count == 0 { "pass" } else { "fail" }.to_string(),
        component: component.unwrap_or_else(|| args.comp.component.clone().unwrap_or_default()),
        scenario_id,
        repeat,
        run_count: runs.len(),
        failure_count,
        exit_code,
        schedule: Some(args.schedule.as_str().to_string()),
        run_order: run_order
            .into_iter()
            .map(|entry| extension_trace::TraceRunOrderEntryOutput {
                index: entry.index,
                group: entry.group,
                iteration: entry.iteration,
            })
            .collect(),
        rig_state,
        overlays,
        runs,
        spans,
        focus_span_ids: args.focus_spans.clone(),
        focus_spans,
    };

    Ok((TraceCommandOutput::Aggregate(output), exit_code))
}

fn focus_aggregate_spans(
    spans: &[extension_trace::TraceAggregateSpanOutput],
    focus_span_ids: &[String],
) -> Vec<extension_trace::TraceAggregateSpanOutput> {
    if focus_span_ids.is_empty() {
        return Vec::new();
    }
    let focus = focus_span_ids.iter().collect::<BTreeSet<_>>();
    spans
        .iter()
        .filter(|span| focus.contains(&span.id))
        .cloned()
        .collect()
}

const DEFAULT_TRACE_PHASE_PRESET: &str = "default";

fn cli_span_definitions_for_args(args: &TraceArgs) -> homeboy::Result<Vec<TraceSpanDefinition>> {
    let mut definitions = args.spans.clone();
    let phase_definitions =
        extension_trace::spans::phase_span_definitions(&args.phases).map_err(|message| {
            homeboy::Error::validation_invalid_argument("--phase", message, None, None)
        })?;
    definitions.extend(phase_definitions);
    Ok(definitions)
}

fn span_definitions_for_args(
    args: &TraceArgs,
    rig_context: Option<&TraceRigContext>,
    extension_id: Option<&str>,
    use_default_preset: bool,
) -> homeboy::Result<Vec<TraceSpanDefinition>> {
    let mut definitions = cli_span_definitions_for_args(args)?;
    let Some(preset_name) = args.phase_preset.as_deref().or_else(|| {
        if use_default_preset {
            default_trace_phase_preset_for_args(args, rig_context, extension_id)
        } else {
            None
        }
    }) else {
        return Ok(definitions);
    };

    let preset_phases = trace_phase_preset_for_args(args, rig_context, extension_id, preset_name)?;
    let phase_definitions = extension_trace::spans::phase_span_definitions(&preset_phases)
        .map_err(|message| {
            homeboy::Error::validation_invalid_argument("--phase-preset", message, None, None)
        })?;
    definitions.extend(phase_definitions);
    Ok(definitions)
}

fn default_trace_phase_preset_for_args<'a>(
    args: &TraceArgs,
    rig_context: Option<&'a TraceRigContext>,
    extension_id: Option<&str>,
) -> Option<&'a str> {
    let context = rig_context?;
    let extension_id = extension_id?;
    let workload = context
        .rig_spec
        .trace_workloads
        .get(extension_id)
        .and_then(|workloads| {
            workloads
                .iter()
                .find(|workload| trace_workload_scenario_id(workload.path()) == args.scenario)
        })?;
    workload.trace_default_phase_preset().or_else(|| {
        workload
            .trace_phase_preset(DEFAULT_TRACE_PHASE_PRESET)
            .map(|_| DEFAULT_TRACE_PHASE_PRESET)
    })
}

fn trace_phase_preset_for_args(
    args: &TraceArgs,
    rig_context: Option<&TraceRigContext>,
    extension_id: Option<&str>,
    preset_name: &str,
) -> homeboy::Result<Vec<extension_trace::spans::TracePhaseMilestone>> {
    let context = rig_context.ok_or_else(|| {
        homeboy::Error::validation_invalid_argument(
            "--phase-preset",
            "phase presets require --rig so Homeboy can read rig/workload metadata",
            None,
            None,
        )
    })?;
    let extension_id = extension_id.ok_or_else(|| {
        homeboy::Error::validation_invalid_argument(
            "--phase-preset",
            "phase presets require a resolved trace extension",
            None,
            None,
        )
    })?;

    let workloads = context
        .rig_spec
        .trace_workloads
        .get(extension_id)
        .map(|workloads| workloads.as_slice())
        .unwrap_or(&[]);
    let workload = workloads
        .iter()
        .find(|workload| trace_workload_scenario_id(workload.path()) == args.scenario);
    let phases = workload
        .and_then(|workload| workload.trace_phase_preset(preset_name))
        .ok_or_else(|| {
            homeboy::Error::validation_invalid_argument(
                "--phase-preset",
                format!(
                    "trace phase preset '{}' is not declared for scenario '{}'",
                    preset_name, args.scenario
                ),
                None,
                None,
            )
        })?;

    phases
        .iter()
        .map(|phase| {
            extension_trace::spans::parse_phase_milestone(phase).map_err(|message| {
                homeboy::Error::validation_invalid_argument("--phase-preset", message, None, None)
            })
        })
        .collect()
}

fn trace_workload_scenario_id(path: &str) -> String {
    let file_name = Path::new(path)
        .file_name()
        .and_then(|name| name.to_str())
        .unwrap_or(path);
    if let Some((stem, _)) = file_name.split_once(".trace.") {
        return stem.to_string();
    }
    Path::new(file_name)
        .file_stem()
        .and_then(|stem| stem.to_str())
        .unwrap_or(file_name)
        .to_string()
}

fn run_list(args: TraceArgs) -> CmdResult<TraceCommandOutput> {
    let rig_context = load_rig_context(args.rig.as_deref())?;
    let effective_id = resolve_component_id(&args.comp, rig_context.as_ref().map(|c| &c.rig_spec))?;
    let path_override = args.comp.path.clone().or_else(|| {
        rig_context
            .as_ref()
            .and_then(|context| rig_component_path(&context.rig_spec, &effective_id))
    });
    let component_override = rig_context
        .as_ref()
        .and_then(|context| rig_component_for_trace(&context.rig_spec, &effective_id));

    let ctx = execution_context::resolve_with_component(
        &ResolveOptions::with_capability_and_json(
            &effective_id,
            path_override.clone(),
            ExtensionCapability::Trace,
            args.setting_args.setting.clone(),
            args.setting_args.setting_json.clone(),
        ),
        component_override,
    )?;
    if let Some(context) = rig_context.as_ref() {
        run_rig_workload_preflight(
            &context.rig_spec,
            ctx.extension_id.as_deref(),
            rig::RigWorkloadKind::Trace,
        )?;
    }

    let run_dir = RunDir::create()?;
    let extra_workloads = rig_context
        .as_ref()
        .and_then(|context| {
            ctx.extension_id.as_deref().map(|id| {
                rig::workloads_for_extension(
                    &context.rig_spec,
                    rig::RigWorkloadKind::Trace,
                    context.rig_package_root.as_deref(),
                    id,
                )
            })
        })
        .unwrap_or_default();
    let list = extension_trace::run_trace_list_workflow(
        &ctx.component,
        TraceListWorkflowArgs {
            component_label: effective_id.clone(),
            component_id: ctx.component_id.clone(),
            path_override,
            settings: settings_as_strings(&ctx.settings),
            settings_json: settings_as_json(&ctx.settings),
            rig_id: args.rig,
            extra_workloads,
        },
        &run_dir,
    )?;

    Ok(extension_trace::from_list_workflow(effective_id, list))
}

struct TraceRigContext {
    rig_spec: RigSpec,
    rig_package_root: Option<PathBuf>,
}

fn load_rig_context(rig_id: Option<&str>) -> homeboy::Result<Option<TraceRigContext>> {
    let Some(rig_id) = rig_id else {
        return Ok(None);
    };
    let spec = rig::load(rig_id)?;
    let package_root =
        rig::read_source_metadata(&spec.id).map(|metadata| PathBuf::from(metadata.package_path));
    Ok(Some(TraceRigContext {
        rig_spec: spec,
        rig_package_root: package_root,
    }))
}

fn run_rig_workload_preflight(
    spec: &RigSpec,
    extension_id: Option<&str>,
    kind: rig::RigWorkloadKind,
) -> homeboy::Result<()> {
    let groups =
        extension_id.and_then(|id| rig::check_groups_for_extension_workloads(spec, kind, id));
    let check = match groups {
        Some(groups) => rig::run_check_groups(spec, &groups)?,
        None => rig::run_check(spec)?,
    };
    if !check.success {
        return Err(homeboy::Error::validation_invalid_argument(
            "--rig",
            format!(
                "rig '{}' check failed; fix the rig before running trace",
                spec.id
            ),
            None,
            None,
        ));
    }
    Ok(())
}

fn resolve_component_id(
    comp: &PositionalComponentArgs,
    rig_spec: Option<&RigSpec>,
) -> homeboy::Result<String> {
    if let Some(id) = comp.id() {
        return Ok(id.to_string());
    }
    if let Some(spec) = rig_spec {
        if spec.components.len() == 1 {
            return Ok(spec.components.keys().next().unwrap().clone());
        }
        return Err(homeboy::Error::validation_invalid_argument(
            "component",
            format!(
                "rig '{}' has multiple components; pass the component id to trace",
                spec.id
            ),
            None,
            None,
        ));
    }
    comp.resolve_id()
}

fn rig_component_path(spec: &RigSpec, component_id: &str) -> Option<String> {
    let component = spec.components.get(component_id)?;
    Some(homeboy::rig::expand::expand_vars(spec, &component.path))
}

fn rig_component_for_trace(spec: &RigSpec, component_id: &str) -> Option<Component> {
    let component = spec.components.get(component_id)?;
    let mut extensions = component.extensions.clone().unwrap_or_default();
    for extension_id in rig::extension_ids_for_workloads(spec, rig::RigWorkloadKind::Trace) {
        extensions
            .entry(extension_id)
            .or_insert_with(ScopedExtensionConfig::default);
    }
    Some(Component {
        id: component_id.to_string(),
        local_path: rig_component_path(spec, component_id)
            .unwrap_or_else(|| component.path.clone()),
        remote_url: component.remote_url.clone(),
        triage_remote_url: component.triage_remote_url.clone(),
        extensions: if extensions.is_empty() {
            None
        } else {
            Some(extensions)
        },
        ..Default::default()
    })
}

fn settings_as_strings(settings: &[(String, serde_json::Value)]) -> Vec<(String, String)> {
    settings
        .iter()
        .filter_map(|(key, value)| match value {
            serde_json::Value::String(s) => Some((key.clone(), s.clone())),
            _ => None,
        })
        .collect()
}

fn settings_as_json(settings: &[(String, serde_json::Value)]) -> Vec<(String, serde_json::Value)> {
    settings
        .iter()
        .filter_map(|(key, value)| match value {
            serde_json::Value::String(_) => None,
            other => Some((key.clone(), other.clone())),
        })
        .collect()
}

struct ActiveTraceObservation {
    store: ObservationStore,
    run_id: String,
    component_id: String,
    rig_id: Option<String>,
    scenario_id: String,
}

fn persist_trace_workflow_result(
    observation: &ActiveTraceObservation,
    run_dir: &RunDir,
    workflow: &extension_trace::TraceRunWorkflowResult,
    rig_state: Option<&rig::RigStateSnapshot>,
) {
    let run_status = trace_run_status(workflow);
    let baseline_status = baseline_status(workflow);
    let results = workflow.results.as_ref();
    let _ = observation.store.record_trace_run(NewTraceRunRecord {
        run_id: observation.run_id.clone(),
        component_id: observation.component_id.clone(),
        rig_id: observation.rig_id.clone(),
        scenario_id: results
            .map(|results| results.scenario_id.clone())
            .unwrap_or_else(|| observation.scenario_id.clone()),
        status: run_status.as_str().to_string(),
        baseline_status: baseline_status.clone(),
        metadata_json: serde_json::json!({
            "status": &workflow.status,
            "exit_code": workflow.exit_code,
            "summary": results.and_then(|results| results.summary.clone()),
            "failure": &workflow.failure,
            "overlays": &workflow.overlays,
            "baseline_comparison": &workflow.baseline_comparison,
            "baseline_status": baseline_status,
            "hints": &workflow.hints,
            "rig_state": rig_state,
            "assertion_count": results.map(|results| results.assertions.len()).unwrap_or(0),
            "artifact_count": results.map(|results| results.artifacts.len()).unwrap_or(0),
            "span_count": results.map(|results| results.span_results.len()).unwrap_or(0),
        }),
    });

    if let Some(results) = results {
        for span in &results.span_results {
            let _ = observation.store.record_trace_span(NewTraceSpanRecord {
                run_id: observation.run_id.clone(),
                span_id: span.id.clone(),
                status: format!("{:?}", span.status).to_ascii_lowercase(),
                duration_ms: span.duration_ms.map(|value| value as f64),
                from_event: Some(span.from.clone()),
                to_event: Some(span.to.clone()),
                metadata_json: serde_json::json!({
                    "from_t_ms": span.from_t_ms,
                    "to_t_ms": span.to_t_ms,
                    "missing": span.missing,
                    "message": &span.message,
                }),
            });
        }
    }

    record_trace_artifacts(&observation.store, &observation.run_id, run_dir, results);
    let _ = observation.store.finish_run(
        &observation.run_id,
        run_status,
        Some(trace_run_finish_metadata(workflow)),
    );
}

fn persist_trace_workflow_error(
    observation: &ActiveTraceObservation,
    run_dir: &RunDir,
    error: &homeboy::Error,
) {
    let error_metadata = serde_json::json!({
        "error": {
            "code": error.code.as_str(),
            "message": &error.message,
            "details": &error.details,
        }
    });
    let _ = observation.store.record_trace_run(NewTraceRunRecord {
        run_id: observation.run_id.clone(),
        component_id: observation.component_id.clone(),
        rig_id: observation.rig_id.clone(),
        scenario_id: observation.scenario_id.clone(),
        status: RunStatus::Error.as_str().to_string(),
        baseline_status: None,
        metadata_json: error_metadata.clone(),
    });
    record_trace_artifacts(&observation.store, &observation.run_id, run_dir, None);
    let _ =
        observation
            .store
            .finish_run(&observation.run_id, RunStatus::Error, Some(error_metadata));
}

fn trace_run_status(workflow: &extension_trace::TraceRunWorkflowResult) -> RunStatus {
    if workflow.failure.is_some() || workflow.status == "error" {
        RunStatus::Error
    } else if workflow.exit_code == 0 && workflow.status == "pass" {
        RunStatus::Pass
    } else {
        RunStatus::Fail
    }
}

fn baseline_status(workflow: &extension_trace::TraceRunWorkflowResult) -> Option<String> {
    workflow.baseline_comparison.as_ref().map(|comparison| {
        if comparison.regression {
            "regression"
        } else if comparison.has_improvements {
            "improvement"
        } else {
            "pass"
        }
        .to_string()
    })
}

fn trace_run_finish_metadata(
    workflow: &extension_trace::TraceRunWorkflowResult,
) -> serde_json::Value {
    serde_json::json!({
        "status": &workflow.status,
        "exit_code": workflow.exit_code,
        "failure": &workflow.failure,
        "overlays": &workflow.overlays,
        "baseline_comparison": &workflow.baseline_comparison,
        "hints": &workflow.hints,
        "results": &workflow.results,
    })
}

fn record_trace_artifacts(
    store: &ObservationStore,
    run_id: &str,
    run_dir: &RunDir,
    results: Option<&extension_trace::TraceResults>,
) {
    let trace_results_path = run_dir.step_file(homeboy::engine::run_dir::files::TRACE_RESULTS);
    record_artifact_if_file(store, run_id, "trace-results", &trace_results_path);
    if let Some(results) = results {
        for artifact in &results.artifacts {
            let path = PathBuf::from(&artifact.path);
            let resolved = if path.is_absolute() {
                path
            } else {
                run_dir.path().join(path)
            };
            record_artifact_if_file(store, run_id, "trace-artifact", &resolved);
        }
    }
}

fn record_artifact_if_file(store: &ObservationStore, run_id: &str, kind: &str, path: &Path) {
    if path.is_file() {
        let _ = store.record_artifact(run_id, kind, path);
    }
}

#[cfg(test)]
mod tests;
