//! Integration test for the IPC worker architecture.
//!
//! Starts cp2k_rs_worker via mpirun, performs a full cycle
//! (InitForceEnv → CalcEnergyForce → GetPotentialEnergy → GetForces → Shutdown),
//! and asserts the results are sane.
//!
//! Run with:
//!   cargo run --example test_worker_ipc --features mpi,extended,build-cp2k
//!
//! The worker binary must already be built:
//!   cargo build --release --bin cp2k_rs_worker --features mpi,extended,build-cp2k

use std::io::{Read, Write};
use std::os::unix::net::UnixStream;
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};

use cp2k_rs::worker_protocol::{Command as Cmd, Cp2kInputSpec, Payload, Request, Response, Status};

// ─── IPC helpers ─────────────────────────────────────────────────────────────

fn write_msg(stream: &mut UnixStream, data: &[u8]) {
    let len = (data.len() as u32).to_le_bytes();
    stream.write_all(&len).unwrap();
    stream.write_all(data).unwrap();
    stream.flush().unwrap();
}

fn read_msg(stream: &mut UnixStream) -> Vec<u8> {
    let mut len_buf = [0u8; 4];
    stream.read_exact(&mut len_buf).unwrap();
    let len = u32::from_le_bytes(len_buf) as usize;
    let mut buf = vec![0u8; len];
    stream.read_exact(&mut buf).unwrap();
    buf
}

fn read_msg_with_timeout(
    stream: &mut UnixStream,
    timeout: Duration,
) -> Result<Vec<u8>, std::io::Error> {
    stream.set_read_timeout(Some(timeout))?;
    let mut len_buf = [0u8; 4];
    if let Err(e) = stream.read_exact(&mut len_buf) {
        let _ = stream.set_read_timeout(None);
        return Err(e);
    }
    let len = u32::from_le_bytes(len_buf) as usize;
    let mut buf = vec![0u8; len];
    let res = stream.read_exact(&mut buf);
    let _ = stream.set_read_timeout(None);
    res.map(|_| buf)
}

fn send_cmd(stream: &mut UnixStream, id: u64, cmd: Cmd) -> Payload {
    let req = Request {
        request_id: id,
        command: cmd,
    };
    let bytes = bincode::serialize(&req).unwrap();
    write_msg(stream, &bytes);

    let raw = read_msg(stream);
    let resp: Response = bincode::deserialize(&raw).unwrap();
    assert_eq!(resp.request_id, id, "Mismatched request_id");
    match resp.status {
        Status::Ok => resp.payload,
        Status::Error(e) => panic!("Worker returned error: {e}"),
    }
}

fn send_cmd_with_timeout(
    stream: &mut UnixStream,
    id: u64,
    cmd: Cmd,
    timeout: Duration,
    output_path: &str,
) -> Payload {
    let req = Request {
        request_id: id,
        command: cmd,
    };
    let bytes = bincode::serialize(&req).unwrap();
    write_msg(stream, &bytes);

    let raw = match read_msg_with_timeout(stream, timeout) {
        Ok(v) => v,
        Err(e) => {
            if matches!(
                e.kind(),
                std::io::ErrorKind::TimedOut | std::io::ErrorKind::WouldBlock
            ) {
                if let Ok(out) = std::fs::read_to_string(output_path) {
                    eprintln!("CP2K output (partial):\n{out}");
                }
                if let Ok(dir) = std::env::var("CP2K_WORKER_LOG_DIR") {
                    if let Ok(entries) = std::fs::read_dir(&dir) {
                        for entry in entries.flatten() {
                            let path = entry.path();
                            if let Ok(content) = std::fs::read_to_string(&path) {
                                eprintln!("Worker log {}:\n{}", path.to_string_lossy(), content);
                            }
                        }
                    }
                }
                if let Ok(entries) = std::fs::read_dir("/tmp") {
                    let mut newest: Option<(std::time::SystemTime, std::path::PathBuf)> = None;
                    for entry in entries.flatten() {
                        let path = entry.path();
                        let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
                        if name.starts_with("cp2k_worker_patched_") && name.ends_with(".inp") {
                            if let Ok(meta) = entry.metadata() {
                                if let Ok(mtime) = meta.modified() {
                                    let replace = match &newest {
                                        Some((t, _)) => mtime > *t,
                                        None => true,
                                    };
                                    if replace {
                                        newest = Some((mtime, path));
                                    }
                                }
                            }
                        }
                    }
                    if let Some((_, path)) = newest {
                        if let Ok(content) = std::fs::read_to_string(&path) {
                            eprintln!("Patched input {}:\n{}", path.to_string_lossy(), content);
                        }
                    }
                }
                panic!("Timed out waiting for worker response to request {id}");
            }
            panic!("IPC read failed: {e}");
        }
    };
    let resp: Response = bincode::deserialize(&raw).unwrap();
    assert_eq!(resp.request_id, id, "Mismatched request_id");
    match resp.status {
        Status::Ok => resp.payload,
        Status::Error(e) => panic!("Worker returned error: {e}"),
    }
}

// ─── main ────────────────────────────────────────────────────────────────────

fn main() {
    // ── locate binaries / files ───────────────────────────────────────────
    let worker_bin = std::env::var("CP2K_WORKER_BIN").unwrap_or_else(|_| {
        // Try release build first, then debug.
        let release = "target/release/cp2k_rs_worker";
        let debug = "target/debug/cp2k_rs_worker";
        if std::path::Path::new(release).exists() {
            release.to_string()
        } else if std::path::Path::new(debug).exists() {
            debug.to_string()
        } else {
            panic!(
                "cp2k_rs_worker not found. Build it first, e.g.:\n  cargo build --release --bin cp2k_rs_worker --all-features\nor set CP2K_WORKER_BIN to the worker path."
            );
        }
    });

    // Write the minimal H2 FIST input inline so the test is self-contained.
    let input_path = "/tmp/test_worker_h2.inp";
    let output_path = "/tmp/test_worker_h2.out";
    std::fs::write(input_path, H2_MINIMAL_INPUT).expect("write input");

    let socket_path = format!("/tmp/cp2k_ipc_test_{}.sock", std::process::id());
    let socket_file = format!("{socket_path}.path");
    let ready_path = format!("{socket_path}.ready");
    let log_dir = format!("/tmp/cp2k_rs_worker_logs_{}", std::process::id());
    let _ = std::fs::remove_file(&socket_path);
    let _ = std::fs::remove_file(&socket_file);
    let _ = std::fs::remove_file(&ready_path);
    let _ = std::fs::create_dir_all(&log_dir);
    std::fs::write(&socket_file, &socket_path).expect("write socket file");

    // CP2K_DATA_DIR
    let data_dir = find_cp2k_data_dir();

    // ── start worker ──────────────────────────────────────────────────────
    // Forward OMPI_MCA transport-selection vars explicitly so that OpenMPI
    // uses TCP/self and does not try to load UCX (which crashes with SIGILL
    // on plain VMs / CI runners that lack InfiniBand / PSM3 hardware).
    let ompi_pml = std::env::var("OMPI_MCA_pml").unwrap_or_else(|_| "ob1".to_string());
    let ompi_btl = std::env::var("OMPI_MCA_btl").unwrap_or_else(|_| "tcp,self".to_string());
    let ompi_osc = std::env::var("OMPI_MCA_osc").unwrap_or_else(|_| "pt2pt".to_string());
    let ompi_mtl = std::env::var("OMPI_MCA_mtl").unwrap_or_else(|_| "^ofi,psm,psm2".to_string());
    let ucx_error_signals = std::env::var("UCX_ERROR_SIGNALS").unwrap_or_default();
    println!(
        "Starting worker: mpirun -n 2 -x OMPI_MCA_pml={ompi_pml} -x OMPI_MCA_btl={ompi_btl} -x OMPI_MCA_osc={ompi_osc} {worker_bin}"
    );
    let mut child = Command::new("mpirun")
        .args([
            "-n",
            "2",
            "-x",
            &format!("OMPI_MCA_pml={ompi_pml}"),
            "-x",
            &format!("OMPI_MCA_btl={ompi_btl}"),
            "-x",
            &format!("OMPI_MCA_osc={ompi_osc}"),
            "-x",
            &format!("OMPI_MCA_mtl={ompi_mtl}"),
            "-x",
            &format!("UCX_ERROR_SIGNALS={ucx_error_signals}"),
            "-x",
            &format!("CP2K_WORKER_SOCKET_FILE={socket_file}"),
            "-x",
            &format!("CP2K_WORKER_READY_FILE={ready_path}"),
            "-x",
            &format!("CP2K_DATA_DIR={data_dir}"),
            "-x",
            &format!("CP2K_WORKER_LOG_DIR={log_dir}"),
            "-x",
            "OMPI_ALLOW_RUN_AS_ROOT=1",
            "-x",
            "OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1",
            &worker_bin,
        ])
        .env("CP2K_WORKER_SOCKET_FILE", &socket_file)
        .env("CP2K_WORKER_READY_FILE", &ready_path)
        .env("CP2K_DATA_DIR", &data_dir)
        .env("CP2K_WORKER_LOG_DIR", &log_dir)
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit())
        .spawn()
        .expect("spawn mpirun");

    // ── wait for ready file ───────────────────────────────────────────────
    println!("Waiting for worker to be ready...");
    let deadline = Instant::now() + Duration::from_secs(10);
    loop {
        if std::path::Path::new(&ready_path).exists() {
            break;
        }
        if Instant::now() > deadline {
            child.kill().ok();
            panic!("Timed out waiting for worker ready file: {ready_path}");
        }
        std::thread::sleep(Duration::from_millis(100));
    }
    println!("Worker is ready.");

    let socket_path = std::fs::read_to_string(&ready_path)
        .unwrap_or_else(|_| socket_path.clone())
        .trim()
        .to_string();

    // ── connect ───────────────────────────────────────────────────────────
    let mut stream = UnixStream::connect(&socket_path).expect("connect to socket");

    let mut id = 0u64;

    // ── GetVersion ────────────────────────────────────────────────────────
    match send_cmd(&mut stream, id, Cmd::GetVersion) {
        Payload::String(v) => println!("CP2K version: {v}"),
        p => panic!("Unexpected payload for GetVersion: {p:?}"),
    }
    id += 1;

    // ── InitForceEnvWithGeometry ──────────────────────────────────────────
    println!("InitForceEnvWithGeometry...");
    let symbols = vec!["H".to_string(), "H".to_string()];
    let positions_angstrom = vec![0.0, 0.0, 0.0, 0.0, 0.0, 1.0];
    let cell_angstrom = vec![10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0];
    send_cmd_with_timeout(
        &mut stream,
        id,
        Cmd::InitForceEnvWithGeometry {
            input: Cp2kInputSpec::FromFile {
                path: input_path.to_string(),
            },
            output: output_path.to_string(),
            symbols,
            positions_angstrom,
            cell_angstrom,
            periodic: "NONE".to_string(),
            project_label: None,
            print_level: None,
            run_type: None,
            charge: None,
            multiplicity: None,
            walltime: None,
            seed: None,
            echo_input: None,
            uks: None,
            roks: None,
            wfn_restart_file: None,
            kpoint_config: None,
            dos_config: None,
        },
        Duration::from_secs(10),
        output_path,
    );
    id += 1;

    // ── GetNatom ──────────────────────────────────────────────────────────
    let natom = match send_cmd(&mut stream, id, Cmd::GetNatom) {
        Payload::UInt(n) => n as usize,
        Payload::Int(n) if n >= 0 => n as usize,
        p => panic!("GetNatom: unexpected {p:?}"),
    };
    println!("natom = {natom}");
    assert_eq!(natom, 2, "Expected 2 H atoms");
    id += 1;

    // ── CalcEnergyForce ───────────────────────────────────────────────────
    println!("CalcEnergyForce...");
    send_cmd(&mut stream, id, Cmd::CalcEnergyForce);
    id += 1;

    // ── GetPotentialEnergy ────────────────────────────────────────────────
    let energy = match send_cmd(&mut stream, id, Cmd::GetPotentialEnergy) {
        Payload::Float(e) => e,
        p => panic!("GetPotentialEnergy: unexpected {p:?}"),
    };
    println!("potential energy = {energy:.8} Ha");
    // FIST with LJ ε=0 → energy should be 0.0
    assert!(
        energy.abs() < 1e-6,
        "Expected near-zero FIST energy, got {energy}"
    );
    id += 1;

    // ── GetForces ─────────────────────────────────────────────────────────
    let forces = match send_cmd(&mut stream, id, Cmd::GetForces) {
        Payload::Array1(v) => v,
        p => panic!("GetForces: unexpected {p:?}"),
    };
    println!("forces = {forces:?}");
    assert_eq!(forces.len(), 6, "Expected 2 atoms × 3 components");
    id += 1;

    // ── GetPositions ──────────────────────────────────────────────────────
    let positions = match send_cmd(&mut stream, id, Cmd::GetPositions) {
        Payload::Array1(v) => v,
        p => panic!("GetPositions: unexpected {p:?}"),
    };
    println!("positions = {positions:?}");
    assert_eq!(positions.len(), 6, "Expected 2 atoms × 3 coordinates");
    id += 1;

    // ── GetCell ───────────────────────────────────────────────────────────
    let cell = match send_cmd(&mut stream, id, Cmd::GetCell) {
        Payload::Array2 { rows, cols, data } => {
            assert_eq!((rows, cols), (3, 3));
            data
        }
        p => panic!("GetCell: unexpected {p:?}"),
    };
    println!(
        "cell diagonal = ({:.3}, {:.3}, {:.3})",
        cell[0], cell[4], cell[8]
    );
    id += 1;

    // ── Shutdown ──────────────────────────────────────────────────────────
    println!("Sending Shutdown...");
    send_cmd(&mut stream, id, Cmd::Shutdown);

    // ── wait for worker to exit ───────────────────────────────────────────
    let grace = Instant::now() + Duration::from_secs(10);
    loop {
        match child.try_wait() {
            Ok(Some(status)) => {
                println!("Worker exited with status: {status}");
                break;
            }
            Ok(None) if Instant::now() < grace => {
                std::thread::sleep(Duration::from_millis(100));
            }
            _ => {
                child.kill().ok();
                println!("Worker killed (did not exit within grace period)");
                break;
            }
        }
    }

    // Cleanup
    let _ = std::fs::remove_file(&socket_path);
    let _ = std::fs::remove_file(&socket_file);
    let _ = std::fs::remove_file(&ready_path);
    let _ = std::fs::remove_file(input_path);
    let _ = std::fs::remove_dir_all(&log_dir);

    println!("\n=== ALL IPC TESTS PASSED ===");
}

fn find_cp2k_data_dir() -> String {
    // The CP2K source tree is checked out under target/release/cp2k-stable/cp2k/
    // and its data/ directory lives directly inside that checkout.
    // Try the canonical location first, then fall back to the bundled data,
    // then search under target/release/ for any layout variation.
    let canonical = std::path::Path::new("target/release/cp2k-stable/cp2k/data");
    if canonical.is_dir() {
        return canonical.to_string_lossy().into_owned();
    }

    let bundled = std::path::Path::new("cp2k_rs/data");
    if bundled.is_dir() {
        return bundled.to_string_lossy().into_owned();
    }

    // Fallback: walk one level under target/release/cp2k-stable looking for data/.
    let base = "target/release/cp2k-stable";
    if let Ok(entries) = std::fs::read_dir(base) {
        for entry in entries.flatten() {
            let candidate = entry.path().join("data");
            if candidate.is_dir() {
                return candidate.to_string_lossy().into_owned();
            }
        }
    }

    // Last resort: honour the environment variable set by the CI / user.
    std::env::var("CP2K_DATA_DIR").unwrap_or_default()
}

const H2_MINIMAL_INPUT: &str = r#"&GLOBAL
  PROJECT_NAME minimal_test
  RUN_TYPE ENERGY
  PRINT_LEVEL LOW
&END GLOBAL

&FORCE_EVAL
  METHOD FIST
  &MM
    &FORCEFIELD
      &CHARGE
        ATOM H
        CHARGE 0.0
      &END CHARGE
      &NONBONDED
        &LENNARD-JONES
          ATOMS H H
          EPSILON 0.0
          SIGMA 1.0
        &END LENNARD-JONES
      &END NONBONDED
    &END FORCEFIELD
    &POISSON
      &EWALD
        EWALD_TYPE NONE
      &END EWALD
    &END POISSON
  &END MM
  &SUBSYS
    &CELL
      ABC 10.0 10.0 10.0
      PERIODIC NONE
    &END CELL
    &COORD
      H   0.000000    0.000000    0.000000
      H   0.000000    0.000000    1.000000
    &END COORD
    &TOPOLOGY
    &END TOPOLOGY
  &END SUBSYS
&END FORCE_EVAL
"#;
