//! Concurrent execution of shell commands with prefixed-line output.
//!
//! Foreground concurrent groups (from `HookStep::Concurrent`) spawn every
//! command at once and combine their output into a single terminal stream,
//! each line prefixed with its command's colored label. Prefixed lines keep
//! full scrollback intact for debugging failures and work identically under
//! a TTY or pipe (CI logs).
//!
//! TODO: upgrade to a docker-compose-v2-style tailing display — each command
//! gets a fixed terminal region showing its last N lines, overwritten in
//! place as output arrives. Better signal-to-noise for long streams like
//! `cargo test`, but substantially more implementation (cursor tracking,
//! resize handling, TTY/non-TTY fallback, scrollback replay on failure).
//! Prefixed lines suffice until live sections are demonstrated to pay for
//! themselves.
//!
//! ## Execution model
//!
//! For each command:
//! 1. Spawn a shell child with stdout+stderr piped and (on Unix) its own
//!    process group so SIGINT/SIGTERM can be delivered to the whole tree.
//! 2. Pipe `context_json` to stdin if provided, then close.
//! 3. Launch two reader threads that read lines and send labeled lines on
//!    a shared channel. A single consumer writes to stderr — one writer
//!    preserves line atomicity so readers never mix bytes mid-line.
//!
//! The main thread drains the channel. A signal-listener thread blocks
//! on `signal_hook::Signals::forever()` for SIGINT/SIGTERM and forwards
//! the user's signal to every live child's process group; a second user
//! signal escalates every still-live pgroup to SIGKILL. Closing the
//! signal-hook handle on shutdown unblocks the listener.
//!
//! All children always run to completion. Per-child exit status is returned
//! for the caller to fold into a failure, matching alias `thread::scope` and
//! pipeline `run_concurrent_group` semantics.

use std::io::{BufRead, BufReader, Read, Write};
use std::path::Path;
use std::process::{Child, Stdio};
use std::sync::mpsc::{self, Sender};
use std::thread;
use std::time::Instant;

use anyhow::Context;

use worktrunk::command_log::log_command;
use worktrunk::git::WorktrunkError;
use worktrunk::shell_exec::{
    DIRECTIVE_CD_FILE_ENV_VAR, DIRECTIVE_EXEC_FILE_ENV_VAR, DIRECTIVE_FILE_ENV_VAR, ShellConfig,
    scrub_directive_env_vars, scrub_git_discovery_env_vars,
};
#[cfg(unix)]
use worktrunk::signal_forwarder::ForegroundSignals;
use worktrunk::styling::stderr;
use worktrunk::trace::CommandTrace;

use super::handlers::DirectivePassthrough;

/// One command in a concurrent group.
pub struct ConcurrentCommand<'a> {
    /// Short label used as the line prefix (e.g., the command name).
    pub label: &'a str,
    /// Fully-expanded shell command string.
    pub expanded: &'a str,
    /// Child's working directory.
    pub working_dir: &'a Path,
    /// JSON blob written to the child's stdin and closed. Callers that have
    /// no context to pass should supply `"{}"`.
    pub context_json: &'a str,
    /// Optional label for `commands.jsonl` tracing.
    pub log_label: Option<&'a str>,
    /// Directive file env vars to pass through to the child. See
    /// `DirectivePassthrough` for the trust model (CD passthrough, EXEC scrub).
    pub directives: &'a DirectivePassthrough,
    /// Scrub inherited git-discovery vars (`GIT_DIR`/`GIT_WORK_TREE`/…) from the
    /// child. `true` for hooks (they operate on the worktree wt targets), `false`
    /// for aliases (they keep wt's inherited context). See issue #3373.
    pub scrub_git_discovery: bool,
}

/// Run every command concurrently and return each per-child result in input
/// order. `Err(WorktrunkError::ChildProcessExited { .. })` signals a non-zero
/// exit; other errors come from spawn/IO failures.
///
/// When the `WORKTRUNK_TEST_SERIAL_CONCURRENT` env var is set, commands run
/// sequentially in input order — same prefix-line output path, just one child
/// at a time. Tests use this to pin deterministic interleaving for snapshots.
pub fn run_concurrent_commands(
    cmds: &[ConcurrentCommand<'_>],
) -> anyhow::Result<Vec<anyhow::Result<()>>> {
    let prefix_width = cmds.iter().map(|c| c.label.len()).max().unwrap_or(0);
    let shell = ShellConfig::get()?;

    if std::env::var_os("WORKTRUNK_TEST_SERIAL_CONCURRENT").is_some() {
        return Ok(run_serial_with_prefix(shell, cmds, prefix_width));
    }

    // Install the SIGINT/SIGTERM latch BEFORE spawning any children so a
    // signal that arrives mid-spawn is captured rather than default-killing
    // wt (which would orphan already-spawned children, since each runs in
    // its own process group and wouldn't see the tty's Ctrl-C broadcast).
    #[cfg(unix)]
    let signals = ForegroundSignals::install()?;

    // Spawn each child and record its start time for commands.jsonl. If any
    // spawn fails partway, kill and reap every child we already spawned —
    // otherwise they'd outlive wt as unreaped orphans with nobody draining
    // their pipes (and `Child::drop` does not kill the process on Unix).
    let mut children: Vec<SpawnedChild> = Vec::with_capacity(cmds.len());
    for (i, cmd) in cmds.iter().enumerate() {
        match spawn_child(shell, i, cmd) {
            Ok(spawned) => children.push(spawned),
            Err(e) => {
                for mut prior in children {
                    let _ = prior.child.kill();
                    let _ = prior.child.wait();
                    // Spawned but torn down because a sibling failed to spawn —
                    // record it rather than leaving the trace guard unresolved.
                    prior.trace.complete(false);
                }
                return Err(e);
            }
        }
    }

    // Print one prefixed line at a time from a single writer. Each reader
    // sends complete lines through `line_tx`; the consumer drains them.
    let (line_tx, line_rx) = mpsc::channel::<LabeledLine>();

    let mut readers: Vec<thread::JoinHandle<()>> = Vec::new();
    for (i, spawned) in children.iter_mut().enumerate() {
        let label = cmds[i].label.to_string();
        if let Some(stdout) = spawned.child.stdout.take() {
            readers.push(spawn_reader(i, label.clone(), stdout, line_tx.clone()));
        }
        if let Some(stderr) = spawned.child.stderr.take() {
            readers.push(spawn_reader(i, label, stderr, line_tx.clone()));
        }
    }
    // Drop the original sender so the channel closes once every reader exits.
    drop(line_tx);

    // Start the forwarder thread now that we have the pgid list. The
    // `signals` latch was installed up-front, so any signal received during
    // the spawn loop was queued and will be delivered on the thread's first
    // poll.
    #[cfg(unix)]
    let signal_thread = signals.forward_to_pgids(
        children
            .iter()
            .map(|c| c.child.id() as i32)
            .collect::<Vec<_>>(),
    );

    // Consume labeled lines until every reader drops its sender.
    {
        let mut stderr = stderr().lock();
        for labeled in line_rx {
            let prefix = render_prefix(labeled.index, &labeled.label, prefix_width);
            writeln!(stderr, "{}{}", prefix, labeled.line).ok();
        }
    }

    for r in readers {
        let _ = r.join();
    }

    // After the last output line prints, wait on each child.
    let mut outcomes = Vec::with_capacity(children.len());
    for (spawned, cmd) in children.into_iter().zip(cmds) {
        outcomes.push(collect_outcome(spawned, cmd));
    }

    #[cfg(unix)]
    let originating_signal = {
        // All children have exited — shut the signal forwarder down. Returns
        // the first signal the forwarder observed (None if no signal arrived).
        signal_thread.stop()
    };
    #[cfg(not(unix))]
    let originating_signal: Option<i32> = None;

    // If a child died from SIGKILL because the user pressed Ctrl-C twice, the
    // user-visible exit code shouldn't be 137 — they only ever sent SIGINT, and
    // wt's escalation to SIGKILL on the second press is an implementation
    // detail. Override each child's reported signal with the originating signal
    // so wt's `exit_code()` and `interrupt_exit_code()` reflect the user's
    // intent (130 from SIGINT, not 137 from SIGKILL).
    if let Some(orig) = originating_signal {
        for outcome in &mut outcomes {
            override_with_originating_signal(outcome, orig);
        }
    }

    Ok(outcomes)
}

/// Replace a child's signal-derived `ChildProcessExited` error with one whose
/// `signal` / `code` / `message` reflect the originating signal wt received
/// (rather than the SIGKILL the forwarder may have escalated to on a second
/// user press). No-op if the outcome isn't a signal-derived error.
fn override_with_originating_signal(outcome: &mut anyhow::Result<()>, originating: i32) {
    let Err(err) = outcome else { return };
    let Some(WorktrunkError::ChildProcessExited {
        signal: Some(child_sig),
        ..
    }) = err.downcast_ref::<WorktrunkError>()
    else {
        return;
    };
    if *child_sig == originating {
        return;
    }
    *outcome = Err(WorktrunkError::ChildProcessExited {
        code: 128 + originating,
        message: format!("terminated by signal {originating}"),
        signal: Some(originating),
    }
    .into());
}

/// Serial fallback for `WORKTRUNK_TEST_SERIAL_CONCURRENT`.
///
/// Runs each command to completion before starting the next — same prefix-line
/// output, deterministic interleaving. Outcomes come back in input order.
fn run_serial_with_prefix(
    shell: &ShellConfig,
    cmds: &[ConcurrentCommand<'_>],
    prefix_width: usize,
) -> Vec<anyhow::Result<()>> {
    let mut results = Vec::with_capacity(cmds.len());
    for (i, cmd) in cmds.iter().enumerate() {
        let spawned = match spawn_child(shell, i, cmd) {
            Ok(s) => s,
            Err(e) => {
                results.push(Err(e));
                continue;
            }
        };
        let result = drain_and_wait_single(spawned, cmd, i, prefix_width);
        results.push(result);
    }
    results
}

/// Spawn readers for one child, drain to stderr with prefixes, then wait.
fn drain_and_wait_single(
    mut spawned: SpawnedChild,
    cmd: &ConcurrentCommand<'_>,
    index: usize,
    prefix_width: usize,
) -> anyhow::Result<()> {
    let (tx, rx) = mpsc::channel::<LabeledLine>();
    let mut readers: Vec<thread::JoinHandle<()>> = Vec::new();
    let label = cmd.label.to_string();
    if let Some(stdout) = spawned.child.stdout.take() {
        readers.push(spawn_reader(index, label.clone(), stdout, tx.clone()));
    }
    if let Some(stderr) = spawned.child.stderr.take() {
        readers.push(spawn_reader(index, label, stderr, tx.clone()));
    }
    drop(tx);
    {
        let mut out = stderr().lock();
        for labeled in rx {
            let prefix = render_prefix(labeled.index, &labeled.label, prefix_width);
            writeln!(out, "{prefix}{}", labeled.line).ok();
        }
    }
    for r in readers {
        let _ = r.join();
    }
    collect_outcome(spawned, cmd)
}

struct SpawnedChild {
    child: Child,
    cmd_str: String,
    log_label: Option<String>,
    started_at: Instant,
    /// `[wt-trace]` record for this child, captured at spawn time and resolved
    /// in `collect_outcome`. Held across the output-draining window so the
    /// recorded duration is the full spawn → wait span, not just the wait.
    trace: CommandTrace,
}

fn spawn_child(
    shell: &ShellConfig,
    index: usize,
    cmd: &ConcurrentCommand<'_>,
) -> anyhow::Result<SpawnedChild> {
    let mut command = shell.command(cmd.expanded);
    command
        .current_dir(cmd.working_dir)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());

    // User hooks discover their repo from the cwd wt sets, not an inherited
    // GIT_DIR/GIT_WORK_TREE (issue #3373). Aliases keep the inherited context.
    if cmd.scrub_git_discovery {
        scrub_git_discovery_env_vars(&mut command);
    }

    // Scrub all directive env vars, then re-add the passthroughs.
    scrub_directive_env_vars(&mut command);
    if let Some(path) = &cmd.directives.cd_file {
        command.env(DIRECTIVE_CD_FILE_ENV_VAR, path);
    }
    if let Some(path) = &cmd.directives.exec_file {
        command.env(DIRECTIVE_EXEC_FILE_ENV_VAR, path);
    }
    if let Some(path) = &cmd.directives.legacy_file {
        command.env(DIRECTIVE_FILE_ENV_VAR, path);
    }

    #[cfg(unix)]
    {
        use std::os::unix::process::CommandExt;
        command.process_group(0);
    }

    tracing::debug!(
        command = %cmd.expanded,
        shell = %shell.name,
        "$ {} (concurrent #{index}, shell: {})",
        cmd.expanded,
        shell.name
    );

    // Start the trace just before spawning so its duration brackets the real
    // spawn → wait span (the child keeps running while we drain its output).
    // Each child is fed its own `context_json` on stdin, so mark it stdin-reading
    // — the same command across worktrees isn't a duplicate (different input).
    let mut trace = CommandTrace::new(None, cmd.expanded).reads_stdin(true);
    let mut child = match command.spawn() {
        Ok(child) => child,
        Err(e) => {
            trace.fail(&e);
            return Err(e)
                .with_context(|| format!("failed to spawn concurrent command '{}'", cmd.label));
        }
    };

    if let Some(mut stdin) = child.stdin.take() {
        // Ignore BrokenPipe — child may exit or close stdin early.
        let _ = stdin.write_all(cmd.context_json.as_bytes());
    }

    Ok(SpawnedChild {
        child,
        cmd_str: cmd.expanded.to_string(),
        log_label: cmd.log_label.map(str::to_string),
        started_at: Instant::now(),
        trace,
    })
}

fn spawn_reader<R: Read + Send + 'static>(
    index: usize,
    label: String,
    stream: R,
    tx: Sender<LabeledLine>,
) -> thread::JoinHandle<()> {
    thread::spawn(move || {
        // Read by bytes, not by `BufRead::lines()` — `lines()` returns
        // `InvalidData` on non-UTF-8 bytes and terminates the iterator, which
        // would leave the child's pipe un-drained and eventually hang
        // `child.wait()` once the pipe buffer fills. Real-world triggers
        // include `git diff` of binary files, tools that honor non-UTF-8
        // locales, and any raw-byte output. Lossy-decoding preserves the
        // stream and keeps the child unblocked.
        let mut reader = BufReader::new(stream);
        let mut buf = Vec::with_capacity(256);
        loop {
            buf.clear();
            match reader.read_until(b'\n', &mut buf) {
                Ok(0) => return,
                Ok(_) => {
                    // Strip the trailing newline (and optional \r) before
                    // forwarding — the writer re-adds a newline per line.
                    if buf.last() == Some(&b'\n') {
                        buf.pop();
                        if buf.last() == Some(&b'\r') {
                            buf.pop();
                        }
                    }
                    let line = String::from_utf8_lossy(&buf).into_owned();
                    if tx
                        .send(LabeledLine {
                            index,
                            label: label.clone(),
                            line,
                        })
                        .is_err()
                    {
                        return; // consumer dropped
                    }
                }
                Err(_) => return, // IO error on the pipe — give up on this stream
            }
        }
    })
}

fn collect_outcome(spawned: SpawnedChild, cmd: &ConcurrentCommand<'_>) -> anyhow::Result<()> {
    let SpawnedChild {
        mut child,
        cmd_str,
        log_label,
        started_at,
        mut trace,
    } = spawned;

    let status = match child.wait() {
        Ok(status) => status,
        Err(e) => {
            trace.fail(&e);
            return Err(e)
                .with_context(|| format!("failed to wait for concurrent command '{}'", cmd.label));
        }
    };
    trace.complete(status.success());

    let duration = started_at.elapsed();
    let exit_code = status.code();

    #[cfg(unix)]
    let signal = std::os::unix::process::ExitStatusExt::signal(&status);
    #[cfg(not(unix))]
    let signal: Option<i32> = None;

    let normalized_code = exit_code.or_else(|| signal.map(|s| 128 + s));

    if let Some(label) = log_label {
        log_command(&label, &cmd_str, normalized_code, Some(duration));
    }

    if status.success() {
        Ok(())
    } else if let Some(sig) = signal {
        Err(WorktrunkError::ChildProcessExited {
            code: 128 + sig,
            message: format!("terminated by signal {sig}"),
            signal: Some(sig),
        }
        .into())
    } else {
        let code = exit_code.unwrap_or(1);
        Err(WorktrunkError::ChildProcessExited {
            code,
            message: format!("exit status: {code}"),
            signal: None,
        }
        .into())
    }
}

struct LabeledLine {
    index: usize,
    label: String,
    line: String,
}

fn render_prefix(index: usize, label: &str, width: usize) -> String {
    use anstyle::{AnsiColor, Color, Style};
    let palette = [
        AnsiColor::Cyan,
        AnsiColor::Magenta,
        AnsiColor::Yellow,
        AnsiColor::Green,
        AnsiColor::Blue,
        AnsiColor::BrightCyan,
        AnsiColor::BrightMagenta,
        AnsiColor::BrightYellow,
    ];
    let style = Style::new()
        .fg_color(Some(Color::Ansi(palette[index % palette.len()])))
        .bold();
    format!("{style}{label:<width$}{style:#} │ ")
}

#[cfg(test)]
mod tests {
    //! Unit tests that exercise the executor's option-bearing code paths which
    //! aren't reachable through the alias integration tests today: every child
    //! currently has `log_label=None` (aliases skip per-child logging) and the
    //! CD/legacy directive env vars are usually unset. Driving these branches
    //! with a direct call proves they behave correctly when a future caller
    //! (concurrent foreground hooks, once their deprecation completes) uses them.
    use super::*;

    fn run_one_with_directives(
        label: &str,
        script: &str,
        log_label: Option<&str>,
        directives: &DirectivePassthrough,
    ) -> Vec<anyhow::Result<()>> {
        let wd = std::env::temp_dir();
        let specs = vec![ConcurrentCommand {
            label,
            expanded: script,
            working_dir: &wd,
            context_json: "{}",
            log_label,
            directives,
            scrub_git_discovery: false,
        }];
        run_concurrent_commands(&specs).expect("spawn failed")
    }

    /// A command with a `log_label` exercises the `log_command` branch in
    /// `collect_outcome` — only hook-origin children take this path today.
    #[test]
    fn test_log_label_is_recorded() {
        let outcomes = run_one_with_directives(
            "job",
            "true",
            Some("test-label"),
            &DirectivePassthrough::default(),
        );
        assert_eq!(outcomes.len(), 1);
        assert!(outcomes[0].is_ok(), "`true` should exit 0");
        // The logger is a global OnceCell not initialised in tests, so
        // `log_command` is effectively a no-op — what we're testing is that
        // passing a log_label doesn't panic and the command still runs.
    }

    /// `DirectivePassthrough` with `cd_file`, `exec_file`, and `legacy_file`
    /// set must propagate all three env vars to the child. The child script
    /// echoes each env var; we assert every value is delivered.
    ///
    /// Unix-only: the script uses POSIX `sh` redirect syntax and relies on
    /// native temp paths that don't need escaping. Git Bash on Windows would
    /// need the paths converted to `/c/...` form first.
    #[test]
    #[cfg(unix)]
    fn test_directive_env_vars_passed_through() {
        use tempfile::NamedTempFile;
        let cd = NamedTempFile::new().unwrap();
        let exec = NamedTempFile::new().unwrap();
        let legacy = NamedTempFile::new().unwrap();
        let directives = DirectivePassthrough {
            cd_file: Some(cd.path().to_path_buf()),
            exec_file: Some(exec.path().to_path_buf()),
            legacy_file: Some(legacy.path().to_path_buf()),
        };
        // Write each env var's value to its matching temp file. If the child
        // didn't receive the env var, the redirect would fail or write an
        // empty file.
        let script = format!(
            "printf CD > {} && printf EXEC > {} && printf LEGACY > {}",
            cd.path().display(),
            exec.path().display(),
            legacy.path().display(),
        );
        let outcomes = run_one_with_directives("job", &script, None, &directives);
        assert!(outcomes[0].is_ok(), "child should exit 0");
        let cd_contents = std::fs::read_to_string(cd.path()).unwrap();
        let exec_contents = std::fs::read_to_string(exec.path()).unwrap();
        let legacy_contents = std::fs::read_to_string(legacy.path()).unwrap();
        assert_eq!(cd_contents, "CD");
        assert_eq!(exec_contents, "EXEC");
        assert_eq!(legacy_contents, "LEGACY");
    }

    /// An empty `cmds` slice must return an empty outcomes vec without
    /// spawning anything or erroring. Integration tests never hit this shape
    /// (parser guarantees non-empty Concurrent groups) so we cover it here.
    #[test]
    fn test_empty_cmds_returns_empty() {
        let outcomes = run_concurrent_commands(&[]).expect("no spawn should happen");
        assert!(outcomes.is_empty());
    }

    /// `override_with_originating_signal` table-test: covers the four shapes
    /// of outcome — `Ok`, non-`WorktrunkError` `Err`, exit-code `Err` (no
    /// signal), and signal-derived `Err` matching vs differing from the
    /// originating signal. Only the differing-signal case mutates.
    #[test]
    fn test_override_with_originating_signal() {
        let originating = 2; // SIGINT

        // Ok: untouched.
        let mut ok: anyhow::Result<()> = Ok(());
        override_with_originating_signal(&mut ok, originating);
        assert!(ok.is_ok());

        // Err but not WorktrunkError: untouched.
        let mut other: anyhow::Result<()> = Err(anyhow::anyhow!("spawn failed"));
        override_with_originating_signal(&mut other, originating);
        assert_eq!(other.unwrap_err().to_string(), "spawn failed");

        // ChildProcessExited with no signal (plain non-zero exit): untouched.
        let mut exit_err: anyhow::Result<()> = Err(WorktrunkError::ChildProcessExited {
            code: 1,
            message: "exit status: 1".into(),
            signal: None,
        }
        .into());
        override_with_originating_signal(&mut exit_err, originating);
        let err = exit_err.unwrap_err();
        let we = err.downcast_ref::<WorktrunkError>().unwrap();
        assert!(matches!(
            we,
            WorktrunkError::ChildProcessExited {
                signal: None,
                code: 1,
                ..
            }
        ));

        // ChildProcessExited with same signal as originating: untouched.
        let mut same: anyhow::Result<()> = Err(WorktrunkError::ChildProcessExited {
            code: 130,
            message: "terminated by signal 2".into(),
            signal: Some(2),
        }
        .into());
        override_with_originating_signal(&mut same, originating);
        let err = same.unwrap_err();
        let we = err.downcast_ref::<WorktrunkError>().unwrap();
        assert!(matches!(
            we,
            WorktrunkError::ChildProcessExited {
                signal: Some(2),
                code: 130,
                ..
            }
        ));

        // ChildProcessExited with different signal (escalation case):
        // overridden to the originating signal's code/message/signal.
        let mut escalated: anyhow::Result<()> = Err(WorktrunkError::ChildProcessExited {
            code: 143,
            message: "terminated by signal 15".into(),
            signal: Some(15),
        }
        .into());
        override_with_originating_signal(&mut escalated, originating);
        let err = escalated.unwrap_err();
        let we = err.downcast_ref::<WorktrunkError>().unwrap();
        match we {
            WorktrunkError::ChildProcessExited {
                code,
                message,
                signal,
            } => {
                assert_eq!(*code, 130);
                assert_eq!(*signal, Some(2));
                assert_eq!(message, "terminated by signal 2");
            }
            _ => panic!("expected ChildProcessExited, got {we:?}"),
        }
    }
}
