use crate::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use ml_core::device::MlDevice;

/// Initialize tracing subscriber once for all smoke tests.
/// Uses `RUST_LOG` env var (defaults to `info`). `try_init` is idempotent.
static TRACING_INIT: std::sync::Once = std::sync::Once::new();
/// Set CUDA deterministic mode env vars BEFORE the first cuBLAS init.
/// Run once per process via OnceLock so all tests in a single `cargo test`
/// run see the same setting.
static DETERMINISTIC_INIT: std::sync::Once = std::sync::Once::new();

fn init_tracing() {
    TRACING_INIT.call_once(|| {
        let filter = tracing_subscriber::EnvFilter::try_from_default_env()
            .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
        tracing_subscriber::fmt()
            .with_env_filter(filter)
            .with_test_writer() // captures output per-test in `cargo test`
            .try_init()
            .ok(); // ignore if another subscriber is already set
    });
}

/// Make CUDA reductions reproducible for smoke-test runs.
///
/// Sets two env vars BEFORE the first cuBLAS handle is created:
///
/// - `CUBLAS_WORKSPACE_CONFIG=:4096:8`
///     Pins cuBLAS to deterministic algorithm selection (the standard knob
///     PyTorch uses for `torch.use_deterministic_algorithms(True)`). Limits
///     workspace size to 4096 bytes × 8 buffers, which forces the same
///     algorithm choice on every launch instead of heuristic-picking based
///     on shape/timing.
///
/// - `NVIDIA_TF32_OVERRIDE=0`
///     Disables TF32 acceleration globally. The codebase explicitly requests
///     `CUBLAS_COMPUTE_32F_FAST_TF32` in cuBLAS calls, so this env var is
///     mostly belt-and-braces — but for tests we want bit-stable F32, and
///     this guarantees no silent TF32 fallback in cuDNN-style algorithms.
///
/// Trade-off: training is 1.5-3× slower in deterministic mode. We pay this
/// only for tests so we can catch real regressions vs. RNG-noise variance
/// across runs. Production code does NOT call this.
///
/// Setting env vars from inside the test process is safe IF done before
/// the first cuBLAS handle is created — cuBLAS reads the env at handle
/// creation, not on every kernel launch. Tests use a fresh CudaContext per
/// run, so first-test-in-the-process setting is sufficient.
fn init_deterministic_cuda() {
    DETERMINISTIC_INIT.call_once(|| {
        // SAFETY: setting env vars before any thread reads them is safe in
        // Rust 1.85+. Smoke tests use #[ignore] and run sequentially via
        // `cargo test --test-threads=1` (default for single-binary), so no
        // race with concurrent CUDA inits.
        unsafe {
            std::env::set_var("CUBLAS_WORKSPACE_CONFIG", ":4096:8");
            std::env::set_var("NVIDIA_TF32_OVERRIDE", "0");
        }
        tracing::info!(
            "CUDA deterministic mode enabled for smoke tests \
             (CUBLAS_WORKSPACE_CONFIG=:4096:8, NVIDIA_TF32_OVERRIDE=0)"
        );
    });
}

/// Production-config smoke test params — all features on, small network for fast tests.
///
/// Matches production exactly: Rainbow DQN + branching + IQN + curiosity + CQL +
/// Kelly + action masking + circuit breaker + n-step + noisy nets.
/// Sizing values (hidden_dim_base=32, batch_size=16, buffer=1024, epochs=1, etc.)
/// are loaded from the `dqn-smoketest` TOML profile.
pub(super) fn smoke_params() -> DQNHyperparameters {
    let mut p = DQNHyperparameters::conservative();

    // ── Single source of truth: config/training/dqn-smoketest.toml ──
    // All training hyperparameters come from the TOML. No inline overrides.
    let profile = crate::training_profile::DqnTrainingProfile::load("dqn-smoketest");
    profile.apply_to(&mut p);

    // ── Test-only flags (not supported by TOML profile yet) ──
    p.replay_buffer_vram_fraction = 0.0; // Explicit buffer_size from TOML
    p.checkpoint_frequency = 100;

    p
}

/// Create a fresh CUDA device for each test.
///
/// NOT cached in a static — each test gets its own CudaContext + CudaStream.
/// A static OnceLock<MlDevice> caused sequential test failures: CudaSlice Drop
/// records errors on the shared context's error_state, poisoning it for the next test.
///
/// Calls `init_deterministic_cuda()` BEFORE creating the CudaContext so
/// CUBLAS_WORKSPACE_CONFIG / NVIDIA_TF32_OVERRIDE are visible to cuBLAS
/// at handle-creation time. This makes smoke-test results bit-stable across
/// re-runs (slower, but trustworthy for A/B regression detection).
pub(super) fn cuda_device() -> MlDevice {
    init_tracing();
    init_deterministic_cuda();
    MlDevice::new_cuda(0)
        .expect("CUDA device required — no CPU fallback in smoke tests")
}

/// Create trainer on CUDA device (no CPU fallback).
pub(super) fn smoke_trainer() -> anyhow::Result<DQNTrainer> {
    Ok(DQNTrainer::new_with_device(smoke_params(), cuda_device())?
        .with_feature_cache(feature_cache_dir()))
}

/// Create trainer with custom params on CUDA device (no CPU fallback).
/// Always sets the fxcache directory for consistent test performance.
pub(super) fn smoke_trainer_with(params: DQNHyperparameters) -> anyhow::Result<DQNTrainer> {
    Ok(DQNTrainer::new_with_device(params, cuda_device())?
        .with_feature_cache(feature_cache_dir()))
}

/// Resolve the workspace root directory.
///
/// Tries (in order):
/// 1. CARGO_MANIFEST_DIR → parent → parent (crates/ml → crates → repo root)
/// 2. Walk up from cwd until `test_data/` or `Cargo.toml` with [workspace] is found
pub(super) fn workspace_root() -> std::path::PathBuf {
    // Try CARGO_MANIFEST_DIR first (set by cargo test/run)
    if let Ok(manifest) = std::env::var("CARGO_MANIFEST_DIR") {
        let p = std::path::Path::new(&manifest);
        if let Some(root) = p.parent().and_then(|p| p.parent()) {
            if root.join("test_data").exists() || root.join("Cargo.toml").exists() {
                return root.to_path_buf();
            }
        }
    }
    // Walk up from cwd
    let mut d = std::env::current_dir().unwrap_or_default();
    loop {
        if d.join("test_data").exists() && d.join("Cargo.toml").exists() {
            return d;
        }
        if !d.pop() { break; }
    }
    // Last resort: cwd
    std::env::current_dir().unwrap_or_default()
}

/// Resolve the feature-cache directory (sibling of test data).
pub(super) fn feature_cache_dir() -> std::path::PathBuf {
    if let Ok(d) = std::env::var("FOXHUNT_FEATURE_CACHE_DIR") {
        let p = std::path::PathBuf::from(&d);
        if p.is_absolute() && p.exists() { return p; }
        let abs = workspace_root().join(&p);
        if abs.exists() { return abs; }
    }
    workspace_root().join("test_data").join("feature-cache")
}

/// Load fxcache data for smoke tests. Returns None if no cache available.
///
/// This is the canonical way to get training data in smoke tests — always use
/// this instead of `trainer.train(data_dir, ...)` which goes through slow DBN
/// parsing when the cache key doesn't match.
pub(super) fn load_smoke_fxcache() -> Option<crate::fxcache::FxCacheData> {
    let cache_dir = feature_cache_dir();
    if !cache_dir.exists() { return None; }
    let entry = std::fs::read_dir(&cache_dir).ok()?
        .filter_map(|e| e.ok())
        .find(|e| e.path().extension().and_then(|s| s.to_str()) == Some("fxcache"))?;
    crate::fxcache::load_fxcache(&entry.path()).ok()
}

/// Initialize a trainer from fxcache data with train/val split.
///
/// Uploads features to GPU and sets up the training/validation ranges.
/// Use with `load_smoke_fxcache()` for consistent, fast smoke tests.
pub(super) fn init_trainer_from_fxcache(
    trainer: &mut super::super::DQNTrainer,
    data: &crate::fxcache::FxCacheData,
    max_bars: usize,
) -> anyhow::Result<()> {
    let n = if max_bars > 0 { data.bar_count.min(max_bars) } else { data.bar_count };
    let train_end = (n * 80) / 100;
    let features = &data.features[..n];
    let targets = &data.targets[..n];
    let ofi = &data.ofi[..n.min(data.ofi.len())];

    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?;
    rt.block_on(trainer.init_from_fxcache(features, targets, ofi))?;
    trainer.set_training_range(0, train_end, train_end, n);
    trainer.set_val_data_from_slices(
        &data.features[train_end..n],
        &data.targets[train_end..n],
        train_end,
    );
    Ok(())
}

/// Assert a value is finite (not NaN or Inf)
pub(super) fn assert_finite(val: f64, name: &str) {
    assert!(val.is_finite(), "{name} is not finite: {val}");
}

/// Resolve base data directory containing symbol subdirectories.
///
/// Returns the PARENT of the symbol directory (e.g. `test_data/futures-baseline`),
/// not the symbol-specific path. Callers pass the symbol separately to
/// `trainer.train(data_dir, "ES.FUT", ...)` which enables fxcache key matching
/// (the cache key hashes all files at data_dir, not just one symbol).
///
/// Uses `FOXHUNT_TEST_DATA` env var if set, otherwise auto-detects
/// relative to the workspace root. Returns `None` if no data found.
pub(super) fn test_data_dir() -> Option<String> {
    // Check env vars: FOXHUNT_TEST_DATA (explicit), TEST_DATA_DIR (CI PVC root)
    for var in &["FOXHUNT_TEST_DATA", "TEST_DATA_DIR"] {
        if let Ok(dir) = std::env::var(var) {
            let path = std::path::Path::new(&dir);
            // If the env var points directly at a symbol dir, return its parent
            if path.file_name().and_then(|n| n.to_str()) == Some("ES.FUT") {
                if let Some(parent) = path.parent() {
                    return Some(parent.to_string_lossy().into_owned());
                }
            }
            // Check if ES.FUT exists as a subdirectory
            for sub in &["ES.FUT", "ohlcv/ES.FUT"] {
                if path.join(sub).exists() {
                    return Some(path.to_string_lossy().into_owned());
                }
            }
        }
    }
    // Auto-detect from workspace root
    let manifest = std::env::var("CARGO_MANIFEST_DIR").unwrap_or_default();
    let workspace = std::path::Path::new(&manifest)
        .parent() // crates/
        .and_then(|p| p.parent()) // repo root
        .unwrap_or(std::path::Path::new("../.."));
    let test_data = workspace.join("test_data");
    // Return the base dir that contains ES.FUT as a subdirectory
    for base in &["futures-baseline", ".", "ohlcv"] {
        let candidate = test_data.join(base);
        if candidate.join("ES.FUT").exists() {
            return Some(candidate.to_string_lossy().into_owned());
        }
    }
    None
}
