//! Queue-based spawn logic for parallel test execution.
//!
//! Instead of zipping sandboxes with batches 1:1, workers pull batches
//! from a shared queue so N sandboxes can process M batches (M >= N).

use std::io::Write;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use indicatif::ProgressBar;
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, info, warn};

use crate::config::Config;
use crate::framework::TestFramework;
use crate::provider::{OutputLine, Sandbox};
use crate::report::MasterJunitReport;

use super::completion::CompletionTracker;
use super::runner::{ArtifactConfig, BatchOutcome, OutputCallback, RunnerConfig, TestRunner};
use super::scheduler::Scheduler;

/// Configuration for a queue-based spawn worker.
///
/// Each worker loops pulling batches from the shared `queue` and
/// executing them with a [`TestRunner`]. Batch indices are assigned
/// atomically via `batch_counter`.
pub(crate) struct SpawnConfig<'a, F: TestFramework, S: Sandbox> {
    pub config: &'a Config,
    pub framework: &'a F,
    pub scheduler: &'a Scheduler,
    pub progress: &'a ProgressBar,
    pub total_tests_to_run: usize,
    pub all_complete: Arc<AtomicBool>,
    pub cancellation_token: CancellationToken,
    pub sandboxes_for_cleanup: Arc<Mutex<Vec<S>>>,
    pub junit_report: Arc<Mutex<MasterJunitReport>>,
    pub logs_dir: PathBuf,
    pub batch_counter: Arc<AtomicUsize>,
    pub verbose: bool,
    pub tracer: crate::trace::Tracer,
    pub sandbox_index: usize,
    pub fail_fast: bool,
    pub tracker: Arc<Mutex<CompletionTracker>>,
    pub ci: bool,
    pub ci_last_decided: Arc<AtomicUsize>,
}

/// Runs a worker that pulls batches from a shared queue until empty.
///
/// The worker assigns batch indices atomically, sets up per-batch log
/// files, and drives the [`TestRunner`] for each batch.
pub(crate) async fn spawn_task<'a, F: TestFramework, S: Sandbox>(
    cfg: SpawnConfig<'a, F, S>,
    mut sandbox: S,
) {
    loop {
        let batch = tokio::select! {
            batch = cfg.scheduler.pop() => batch,
            () = cfg.cancellation_token.cancelled() => None,
        };
        let Some(batch) = batch else {
            break;
        };

        let batch_idx = cfg.batch_counter.fetch_add(1, Ordering::SeqCst);
        info!(
            "Worker {} picked up batch {} ({} tests)",
            cfg.sandbox_index,
            batch_idx,
            batch.tests.len()
        );

        // Early exit if all tests have completed or fail-fast triggered
        if cfg.all_complete.load(Ordering::SeqCst) || cfg.cancellation_token.is_cancelled() {
            let test_ids: Vec<_> = batch.tests.iter().map(|t| t.id()).collect();
            info!(
                "EARLY STOP: Skipping batch {} ({} tests) - cancelled",
                batch_idx,
                batch.tests.len()
            );
            debug!("Skipped tests: {:?}", test_ids);

            for suffix in ["stdout", "stderr"] {
                let log_src = cfg.logs_dir.join(format!("batch-{}.{}", batch_idx, suffix));
                if log_src.exists() {
                    let log_dst = cfg
                        .logs_dir
                        .join(format!("batch-{}.{}.cancelled", batch_idx, suffix));
                    if let Err(e) = std::fs::rename(&log_src, &log_dst) {
                        warn!("Failed to rename batch log: {}", e);
                    }
                }
            }

            if let Ok(mut cleanups) = cfg.sandboxes_for_cleanup.lock() {
                cleanups.push(sandbox);
            } else {
                error!("sandbox cleanup mutex poisoned during early stop");
            }
            return;
        }

        // Skip batches where all tests have a decided outcome (passed/flaky or retries exhausted)
        if let Ok(tracker) = cfg.tracker.lock()
            && batch.tests.iter().all(|t| tracker.is_decided(t.id()))
        {
            let test_ids: Vec<_> = batch.tests.iter().map(|t| t.id()).collect();
            info!(
                "SKIP: Batch {} ({} tests) all already decided, skipping",
                batch_idx,
                batch.tests.len()
            );
            debug!("Skipped tests: {:?}", test_ids);
            continue;
        }

        let sandbox_pid = crate::trace::sandbox_pid(cfg.sandbox_index);
        let _batch_span = cfg
            .tracer
            .span(
                &format!("batch_{}", batch_idx),
                "exec",
                sandbox_pid,
                crate::trace::TID_EXEC,
            )
            .with_args(serde_json::json!({
                "batch_index": batch_idx,
                "test_count": batch.tests.len(),
                "sandbox_index": cfg.sandbox_index,
            }));

        // Set up runner
        let parts_dir = cfg.config.report.output_dir.join("junit-parts");
        let runner_config = RunnerConfig {
            fail_fast: cfg.fail_fast,
            parts_dir,
            junit_report: Arc::clone(&cfg.junit_report),
            tracker: Arc::clone(&cfg.tracker),
            cancellation_token: cfg.cancellation_token.clone(),
            artifacts: ArtifactConfig {
                globs: cfg.config.report.download_globs.clone(),
                output_dir: cfg.config.report.output_dir.clone(),
                on_failure_only: cfg.config.report.download_globs_failure_only,
            },
        };
        let mut runner = TestRunner::new(
            sandbox,
            cfg.framework,
            Duration::from_secs(cfg.config.offload.test_timeout_secs),
            cfg.tracer.clone(),
            sandbox_pid,
            batch_idx,
            runner_config,
        );

        // Per-runner log files (separate stdout and stderr)
        {
            let stdout_path = cfg.logs_dir.join(format!("batch-{}.stdout", batch_idx));
            let stderr_path = cfg.logs_dir.join(format!("batch-{}.stderr", batch_idx));
            match (
                std::fs::File::create(&stdout_path),
                std::fs::File::create(&stderr_path),
            ) {
                (Ok(mut stdout_file), Ok(mut stderr_file)) => {
                    let callback: OutputCallback = Box::new(move |_test_id, line| {
                        let (file, msg): (&mut std::fs::File, _) = match line {
                            OutputLine::Stdout(s) => (&mut stdout_file, format!("{}\n", s)),
                            OutputLine::Stderr(s) => (&mut stderr_file, format!("{}\n", s)),
                            OutputLine::ExitCode(_) => return,
                        };
                        if let Err(e) = file.write_all(msg.as_bytes()) {
                            warn!("Failed to write to batch log: {}", e);
                        }
                    });
                    runner.set_output_callback(callback);
                }
                (Err(e), _) | (_, Err(e)) => {
                    warn!("Failed to create batch log files: {}", e);
                }
            }
        }

        // Verbose logging
        if cfg.verbose {
            for test in &batch.tests {
                println!("Running: {}", test.id());
            }
        }

        // Run tests, racing against cancellation
        let stdout_src = cfg.logs_dir.join(format!("batch-{}.stdout", batch_idx));
        let stderr_src = cfg.logs_dir.join(format!("batch-{}.stderr", batch_idx));
        let outcome = tokio::select! {
            result = runner.run_tests(&batch.tests) => result,
            () = cfg.cancellation_token.cancelled() => Ok(BatchOutcome::Cancelled),
        };

        info!(
            "Worker {} finished batch {} ({} tests): {:?}",
            cfg.sandbox_index,
            batch_idx,
            batch.tests.len(),
            outcome
                .as_ref()
                .map(|o| *o)
                .unwrap_or(BatchOutcome::Cancelled)
        );

        // Rename log files based on outcome
        let extension = match &outcome {
            Ok(BatchOutcome::Success) => "success",
            Ok(BatchOutcome::Failure) => "failure",
            Ok(BatchOutcome::Cancelled) => "cancelled",
            Err(_) => "error",
        };
        for src in [&stdout_src, &stderr_src] {
            if src.exists() {
                let dst = cfg.logs_dir.join(format!(
                    "{}.{}",
                    src.file_name().unwrap_or_default().to_string_lossy(),
                    extension
                ));
                if let Err(e) = std::fs::rename(src, &dst) {
                    warn!("Failed to rename batch log: {}", e);
                }
            }
        }

        // Handle outcome
        match &outcome {
            Ok(BatchOutcome::Success) | Ok(BatchOutcome::Failure) => {
                // Fail-fast: cancel on first failure
                if cfg.fail_fast
                    && matches!(&outcome, Ok(BatchOutcome::Failure))
                    && !cfg.cancellation_token.is_cancelled()
                {
                    info!(
                        "FAIL-FAST: Batch {} had failures. Cancelling remaining batches.",
                        batch_idx
                    );
                    cfg.cancellation_token.cancel();
                }
            }
            Ok(BatchOutcome::Cancelled) => {
                debug!("Batch {} was cancelled", batch_idx);
            }
            Err(e) => {
                error!("Batch execution error: {}", e);
            }
        }

        // Update progress from tracker (record_batch already called inside runner)
        if let (Ok(report), Ok(tracker)) = (cfg.junit_report.lock(), cfg.tracker.lock()) {
            let decided = tracker.decided_count();
            cfg.progress.set_position(decided as u64);
            let passed = report.passed_count();
            let failed = report.failed_count();
            let flaky = report.flaky_count();
            let awaiting = cfg.total_tests_to_run - decided;
            if cfg.ci {
                let prev = cfg.ci_last_decided.fetch_max(decided, Ordering::SeqCst);
                if decided > prev {
                    let pct = (decided * 100)
                        .checked_div(cfg.total_tests_to_run)
                        .unwrap_or(100);
                    eprintln!(
                        "[ci] {pct}% | passed: {passed}, failed: {failed}, flaky: {flaky}, awaiting: {awaiting}"
                    );
                }
            } else {
                cfg.progress.set_message(format!(
                    "{}\n{}\n{}\n{}",
                    console::style(format!("passed: {passed}")).green(),
                    console::style(format!("failed: {failed}")).red(),
                    console::style(format!("flaky: {flaky}")).yellow(),
                    console::style(format!("awaiting: {awaiting}")).dim(),
                ));
            }

            // Early stop: all tests have a decided outcome
            if tracker.all_complete() && !cfg.all_complete.load(Ordering::SeqCst) {
                info!(
                    "EARLY STOP TRIGGERED: All {} tests have completed after batch {} completed. Cancelling remaining batches.",
                    cfg.total_tests_to_run, batch_idx
                );
                cfg.all_complete.store(true, Ordering::SeqCst);
                cfg.cancellation_token.cancel();
            }
        }
        drop(_batch_span);
        sandbox = runner.into_sandbox();
    }

    // Return sandbox for cleanup
    if let Ok(mut cleanups) = cfg.sandboxes_for_cleanup.lock() {
        cleanups.push(sandbox);
    } else {
        error!("sandbox cleanup mutex poisoned during cleanup");
    }
}
