use std::path::PathBuf;

use anyhow::Result;

#[derive(Debug, Clone)]
pub struct ModuleBConfig {
    pub memfd_bytes: u64,
    pub socket_path: Option<PathBuf>,
    pub marker: String,
}

#[derive(Debug, Clone)]
pub struct ModuleBReceiverConfig {
    pub socket_path: PathBuf,
    pub marker: String,
    pub memfd_bytes: u64,
}

#[derive(Debug, Clone)]
pub struct ModuleBStats {
    pub memfd_bytes: u64,
    pub handoff_latency_ms: f64,
    pub marker_ok: bool,
    pub socket_path: PathBuf,
}

impl ModuleBStats {
    pub fn to_json(&self) -> String {
        format!(
            "{{\"module\":\"B\",\"memfd_bytes\":{},\"handoff_latency_ms\":{:.3},\"marker_ok\":{},\"socket_path\":\"{}\"}}",
            self.memfd_bytes,
            self.handoff_latency_ms,
            self.marker_ok,
            self.socket_path.display()
        )
    }
}

#[cfg(target_os = "linux")]
mod linux {
    use std::ffi::CString;
    use std::fs;
    use std::io::{IoSlice, IoSliceMut, Read, Write};
    use std::num::NonZeroUsize;
    use std::os::fd::{AsFd, AsRawFd, FromRawFd, OwnedFd, RawFd};
    use std::os::unix::net::{UnixListener, UnixStream};
    use std::path::{Path, PathBuf};
    use std::process::Command;
    use std::ptr::NonNull;
    use std::thread;
    use std::time::{Duration, Instant};

    use anyhow::{Context, Result};
    use nix::cmsg_space;
    use nix::sys::memfd::{memfd_create, MFdFlags};
    use nix::sys::mman::{mmap, munmap, MapFlags, ProtFlags};
    use nix::sys::socket::{recvmsg, sendmsg, ControlMessage, ControlMessageOwned, MsgFlags};
    use nix::unistd::ftruncate;

    use crate::module_b::{ModuleBConfig, ModuleBReceiverConfig, ModuleBStats};

    const ACK_SUCCESS: u8 = 1;

    pub fn run_parent(config: ModuleBConfig) -> Result<ModuleBStats> {
        let memfd_len = usize::try_from(config.memfd_bytes)
            .context("memfd_bytes exceeds usize on this platform")?;
        if memfd_len == 0 {
            anyhow::bail!("memfd_bytes must be > 0");
        }
        if config.marker.is_empty() {
            anyhow::bail!("marker must not be empty");
        }
        if config.marker.len() > memfd_len {
            anyhow::bail!("marker length must be <= memfd_bytes");
        }

        let socket_path = config
            .socket_path
            .unwrap_or_else(|| default_socket_path(std::process::id()));
        if socket_path.exists() {
            fs::remove_file(&socket_path).with_context(|| {
                format!(
                    "failed to remove pre-existing socket path {}",
                    socket_path.display()
                )
            })?;
        }

        let listener = UnixListener::bind(&socket_path)
            .with_context(|| format!("failed to bind UDS listener at {}", socket_path.display()))?;

        let memfd_name = CString::new("tracer-bullet-memfd").context("invalid memfd name")?;
        let memfd = memfd_create(memfd_name.as_c_str(), MFdFlags::MFD_CLOEXEC)
            .context("memfd_create failed")?;
        ftruncate(&memfd, config.memfd_bytes as i64).context("ftruncate(memfd) failed")?;

        let mut writable_map = MmapGuard::map(&memfd, memfd_len, true)
            .context("failed to mmap writable sender memory")?;
        write_marker(writable_map.as_mut_slice(), config.marker.as_bytes())
            .context("failed to write validation marker into memfd")?;

        let exe = std::env::current_exe().context("failed to locate current executable")?;
        let mut child = Command::new(exe)
            .arg("module-b-receiver-internal")
            .arg("--socket-path")
            .arg(&socket_path)
            .arg("--marker")
            .arg(&config.marker)
            .arg("--memfd-bytes")
            .arg(config.memfd_bytes.to_string())
            .spawn()
            .context("failed to spawn module-b receiver child process")?;

        let (mut stream, _) = listener
            .accept()
            .context("receiver failed to connect to UDS")?;

        let start = Instant::now();
        send_fd(stream.as_raw_fd(), memfd.as_raw_fd()).context("sendmsg SCM_RIGHTS failed")?;

        let mut ack = [0_u8; 1];
        stream
            .read_exact(&mut ack)
            .context("failed to read receiver ACK")?;
        let handoff_latency_ms = start.elapsed().as_secs_f64() * 1_000.0;

        let status = child
            .wait()
            .context("failed waiting for receiver child process")?;
        if !status.success() {
            anyhow::bail!("receiver child exited unsuccessfully: {status}");
        }

        drop(writable_map);

        if socket_path.exists() {
            let _ = fs::remove_file(&socket_path);
        }

        if ack[0] != ACK_SUCCESS {
            anyhow::bail!("receiver returned non-success ACK byte: {}", ack[0]);
        }

        Ok(ModuleBStats {
            memfd_bytes: config.memfd_bytes,
            handoff_latency_ms,
            marker_ok: true,
            socket_path,
        })
    }

    pub fn run_receiver_internal(config: ModuleBReceiverConfig) -> Result<()> {
        let memfd_len = usize::try_from(config.memfd_bytes)
            .context("memfd_bytes exceeds usize on this platform")?;
        if memfd_len == 0 {
            anyhow::bail!("memfd_bytes must be > 0");
        }

        let mut stream = connect_with_retry(&config.socket_path, 200, Duration::from_millis(10))
            .with_context(|| {
                format!(
                    "failed to connect to sender socket {}",
                    config.socket_path.display()
                )
            })?;

        let received_fd = recv_fd(stream.as_raw_fd()).context("recvmsg SCM_RIGHTS failed")?;
        let mapped = MmapGuard::map(&received_fd, memfd_len, false)
            .context("receiver failed to mmap read-only memfd")?;

        validate_marker(mapped.as_slice(), config.marker.as_bytes())
            .context("receiver marker validation failed")?;

        stream
            .write_all(&[ACK_SUCCESS])
            .context("failed to write ACK to sender")?;
        stream.flush().context("failed to flush ACK to sender")?;

        Ok(())
    }

    fn default_socket_path(pid: u32) -> PathBuf {
        PathBuf::from(format!("/tmp/tracer-bullet-{}.sock", pid))
    }

    fn send_fd(socket_fd: RawFd, fd_to_send: RawFd) -> Result<()> {
        let payload = [0xAB_u8];
        let iov = [IoSlice::new(&payload)];
        let fds = [fd_to_send];
        let cmsg = [ControlMessage::ScmRights(&fds)];

        let bytes = sendmsg::<()>(socket_fd, &iov, &cmsg, MsgFlags::empty(), None)
            .context("sendmsg failed")?;
        if bytes != payload.len() {
            anyhow::bail!("sendmsg wrote {} bytes instead of {}", bytes, payload.len());
        }

        Ok(())
    }

    fn recv_fd(socket_fd: RawFd) -> Result<OwnedFd> {
        let mut payload = [0_u8; 1];
        let mut iov = [IoSliceMut::new(&mut payload)];
        let mut cmsgspace = cmsg_space!([RawFd; 1]);

        let msg = recvmsg::<()>(socket_fd, &mut iov, Some(&mut cmsgspace), MsgFlags::empty())
            .context("recvmsg failed")?;

        if msg.bytes == 0 {
            anyhow::bail!("recvmsg received no payload bytes");
        }
        if msg
            .flags
            .intersects(MsgFlags::MSG_CTRUNC | MsgFlags::MSG_TRUNC)
        {
            anyhow::bail!("recvmsg control/payload data was truncated");
        }

        for cmsg in msg
            .cmsgs()
            .context("failed to parse recvmsg control messages")?
        {
            if let ControlMessageOwned::ScmRights(fds) = cmsg {
                if let Some(fd) = fds.first().copied() {
                    let owned = unsafe { OwnedFd::from_raw_fd(fd) };
                    return Ok(owned);
                }
            }
        }

        anyhow::bail!("no SCM_RIGHTS file descriptor received")
    }

    fn connect_with_retry(path: &Path, attempts: usize, sleep: Duration) -> Result<UnixStream> {
        let mut last_err = None;
        for _ in 0..attempts {
            match UnixStream::connect(path) {
                Ok(stream) => return Ok(stream),
                Err(err) => {
                    last_err = Some(err);
                    thread::sleep(sleep);
                }
            }
        }

        match last_err {
            Some(err) => {
                Err(err).with_context(|| format!("unable to connect to {}", path.display()))
            }
            None => anyhow::bail!("unable to connect to {}", path.display()),
        }
    }

    fn write_marker(mapped: &mut [u8], marker: &[u8]) -> Result<()> {
        if marker.len() > mapped.len() {
            anyhow::bail!("marker longer than mapped memory");
        }
        let start = mapped.len() - marker.len();
        mapped[start..].copy_from_slice(marker);
        Ok(())
    }

    fn validate_marker(mapped: &[u8], marker: &[u8]) -> Result<()> {
        if marker.len() > mapped.len() {
            anyhow::bail!("marker longer than mapped memory");
        }
        let start = mapped.len() - marker.len();
        if mapped[start..] != *marker {
            anyhow::bail!("marker mismatch at memfd tail");
        }
        Ok(())
    }

    struct MmapGuard {
        ptr: NonNull<std::ffi::c_void>,
        len: usize,
    }

    impl MmapGuard {
        fn map<Fd: AsFd>(fd: &Fd, len: usize, writable: bool) -> Result<Self> {
            let nz_len = NonZeroUsize::new(len).context("mmap length must be > 0")?;
            let prot = if writable {
                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE
            } else {
                ProtFlags::PROT_READ
            };

            let ptr = unsafe { mmap(None, nz_len, prot, MapFlags::MAP_SHARED, fd, 0) }
                .context("mmap syscall failed")?;

            Ok(Self { ptr, len })
        }

        fn as_slice(&self) -> &[u8] {
            unsafe { std::slice::from_raw_parts(self.ptr.as_ptr() as *const u8, self.len) }
        }

        fn as_mut_slice(&mut self) -> &mut [u8] {
            unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr() as *mut u8, self.len) }
        }
    }

    impl Drop for MmapGuard {
        fn drop(&mut self) {
            let _ = unsafe { munmap(self.ptr, self.len) };
        }
    }

    #[cfg(test)]
    mod tests {
        use std::io::{Read, Write};
        use std::os::fd::AsRawFd;
        use std::os::unix::net::UnixStream;
        use std::thread;

        use super::*;

        #[test]
        fn transfers_fd_and_validates_marker() {
            let memfd_name = CString::new("module-b-test").expect("valid test memfd name");
            let memfd = memfd_create(memfd_name.as_c_str(), MFdFlags::MFD_CLOEXEC)
                .expect("memfd_create should succeed in test");

            let len = 8 * 1024 * 1024;
            let marker = b"TRACER_BULLET_SUCCESS";
            ftruncate(&memfd, len as i64).expect("ftruncate should succeed in test");

            let mut map = MmapGuard::map(&memfd, len, true).expect("mmap should succeed in test");
            write_marker(map.as_mut_slice(), marker).expect("writing marker should succeed");

            let (mut send_stream, recv_stream) =
                UnixStream::pair().expect("socketpair should succeed");

            let marker_vec = marker.to_vec();
            let join = thread::spawn(move || {
                let fd =
                    recv_fd(recv_stream.as_raw_fd()).expect("recv_fd should succeed in thread");
                let mapped = MmapGuard::map(&fd, len, false).expect("receiver mmap should succeed");
                validate_marker(mapped.as_slice(), &marker_vec)
                    .expect("marker should validate in receiver");
                let mut stream = recv_stream;
                stream
                    .write_all(&[ACK_SUCCESS])
                    .expect("thread ACK write should succeed");
            });

            send_fd(send_stream.as_raw_fd(), memfd.as_raw_fd()).expect("send_fd should succeed");
            let mut ack = [0_u8; 1];
            send_stream
                .read_exact(&mut ack)
                .expect("ACK read should succeed in sender");
            assert_eq!(ack, [ACK_SUCCESS]);

            join.join()
                .expect("receiver thread should join successfully");
        }

        #[test]
        #[ignore = "heavy memory mapping validation"]
        fn validates_full_5gb_mapping_without_copy() {
            let memfd_name = CString::new("module-b-test-5gb").expect("valid test memfd name");
            let memfd = memfd_create(memfd_name.as_c_str(), MFdFlags::MFD_CLOEXEC)
                .expect("memfd_create should succeed in test");

            let len = 5 * 1024 * 1024 * 1024usize;
            let marker = b"TRACER_BULLET_SUCCESS";
            ftruncate(&memfd, len as i64).expect("ftruncate should succeed in test");

            let mut map = MmapGuard::map(&memfd, len, true).expect("mmap should succeed in test");
            write_marker(map.as_mut_slice(), marker).expect("write marker should succeed");
            validate_marker(map.as_slice(), marker).expect("validate marker should succeed");
        }
    }
}

#[cfg(target_os = "linux")]
pub fn run_parent(config: ModuleBConfig) -> Result<ModuleBStats> {
    linux::run_parent(config)
}

#[cfg(target_os = "linux")]
pub fn run_receiver_internal(config: ModuleBReceiverConfig) -> Result<()> {
    linux::run_receiver_internal(config)
}

#[cfg(not(target_os = "linux"))]
pub fn run_parent(_config: ModuleBConfig) -> Result<ModuleBStats> {
    anyhow::bail!("module-b requires Linux for memfd + SCM_RIGHTS")
}

#[cfg(not(target_os = "linux"))]
pub fn run_receiver_internal(_config: ModuleBReceiverConfig) -> Result<()> {
    anyhow::bail!("module-b receiver requires Linux")
}
