use std::{
    io,
    mem::MaybeUninit,
    os::yggdrasil::process::{ChildExt, CommandExt},
    process::{Child, Command},
};

use runtime::{
    abi::{
        arch::SavedFrame,
        debug::{self, TraceEventPayload, TraceFlags},
    },
    rt::{
        debug::debug_control,
        process::ThreadEvent,
        time::{get_monotonic_time, SystemTime},
    },
};

pub struct CommandTracer {
    #[allow(unused)]
    child: Child,
    tid: u32,
    buffer: [u8; 512],
    start: SystemTime,

    last_syscall: Option<EnteredSyscall>,
}

struct EnteredSyscall {
    seq: u64,
    timestamp: SystemTime,
    regs: SavedFrame,
}

#[allow(unused)]
#[derive(Debug)]
pub struct SyscallTrace {
    pub entry: SystemTime,
    pub exit: SystemTime,
    pub function: usize,
    pub args: [usize; 6],
    pub result: usize,
}

impl CommandTracer {
    pub fn spawn(mut command: Command) -> io::Result<Self> {
        let mut buffer = [0; 512];
        unsafe { command.attach_tracing() };
        let start = get_monotonic_time().unwrap();
        let child = command.spawn()?;
        let tid = child.main_thread_id()?;

        debug_control::<debug::SetTraceFlags>(
            tid,
            &mut buffer,
            &(TraceFlags::SYSCALL_ENTRY | TraceFlags::SYSCALL_EXIT),
        )
        .map_err(io::Error::from)?;

        Ok(Self {
            child,
            tid,
            buffer,
            start,

            last_syscall: None,
        })
    }

    pub fn start_timestamp(&self) -> SystemTime {
        self.start
    }

    fn next_event(&self) -> io::Result<ThreadEvent> {
        let mut event = MaybeUninit::uninit();
        unsafe { runtime::rt::sys::wait_thread(self.tid, &mut event) }.map_err(io::Error::from)?;
        Ok(unsafe { event.assume_init() })
    }

    fn resume(&mut self) -> io::Result<()> {
        debug_control::<debug::Resume>(self.tid, &mut self.buffer, &false).map_err(io::Error::from)
    }

    fn get_regs(&mut self) -> io::Result<SavedFrame> {
        debug_control::<debug::GetRegisters>(self.tid, &mut self.buffer, &())
            .map_err(io::Error::from)
    }

    pub fn read_word(&mut self, address: usize) -> io::Result<usize> {
        debug_control::<debug::ReadMemory>(self.tid, &mut self.buffer, &address)
            .map_err(io::Error::from)
    }

    pub fn read_memory(&mut self, mut address: usize, buffer: &mut [u8]) -> io::Result<()> {
        const ALIGN: usize = size_of::<usize>() - 1;

        let mut offset = 0;
        let mut remaining = buffer.len();
        while remaining != 0 {
            let aligned_address = address & !ALIGN;
            let offset_in_word = address & ALIGN;
            let amount = (size_of::<usize>() - offset_in_word).min(remaining);
            let word = self.read_word(aligned_address)?.to_ne_bytes();

            buffer[offset..offset + amount]
                .copy_from_slice(&word[offset_in_word..offset_in_word + amount]);

            offset += amount;
            address += amount;
            remaining -= amount;
        }

        Ok(())
    }

    pub fn run<F: Fn(&mut Self, &SyscallTrace)>(mut self, formatter: F) -> io::Result<()> {
        loop {
            // Wait for event
            let event = self.next_event()?;

            match event {
                ThreadEvent::Exited => {
                    break;
                }
                ThreadEvent::Trace(trace) => {
                    match trace.payload {
                        TraceEventPayload::SyscallEntry(seq) => {
                            let regs = self.get_regs()?;
                            self.last_syscall = Some(EnteredSyscall {
                                seq,
                                timestamp: trace.timestamp,
                                regs,
                            });
                        }
                        TraceEventPayload::SyscallExit(seq) => {
                            let regs = self.get_regs()?;

                            if let Some(last) = self.last_syscall.take()
                                && last.seq == seq
                            {
                                let event = SyscallTrace::from_parts(&last, &regs, trace.timestamp);
                                formatter(&mut self, &event);
                            }
                        }
                        _ => (),
                    }

                    if trace.suspend {
                        // Resume
                        self.resume()?;
                    }
                }
            }
        }

        Ok(())
    }
}

#[cfg(any(target_arch = "x86_64", rust_analyzer))]
impl SyscallTrace {
    fn from_parts(entry: &EnteredSyscall, exit_regs: &SavedFrame, exit: SystemTime) -> Self {
        let function = entry.regs.rax as usize;
        let args = [
            entry.regs.rdi as usize,
            entry.regs.rsi as usize,
            entry.regs.rdx as usize,
            entry.regs.r10 as usize,
            entry.regs.r8 as usize,
            entry.regs.r9 as usize,
        ];
        let result = exit_regs.rax as usize;
        let entry = entry.timestamp;

        Self {
            entry,
            exit,
            function,
            args,
            result,
        }
    }
}

#[cfg(any(target_arch = "riscv64", rust_analyzer))]
impl SyscallTrace {
    fn from_parts(entry: &EnteredSyscall, exit_regs: &SavedFrame, exit: SystemTime) -> Self {
        let function = entry.regs.an[0];
        let args = [
            entry.regs.an[1],
            entry.regs.an[2],
            entry.regs.an[3],
            entry.regs.an[4],
            entry.regs.an[5],
            entry.regs.an[6],
        ];
        let result = exit_regs.an[0];
        let entry = entry.timestamp;

        Self {
            entry,
            exit,
            function,
            args,
            result,
        }
    }
}

#[cfg(any(target_arch = "aarch64", rust_analyzer))]
impl SyscallTrace {
    fn from_parts(entry: &EnteredSyscall, exit_regs: &SavedFrame, exit: SystemTime) -> Self {
        let function = entry.regs.gp_regs[8];
        let args = [
            entry.regs.gp_regs[0],
            entry.regs.gp_regs[1],
            entry.regs.gp_regs[2],
            entry.regs.gp_regs[3],
            entry.regs.gp_regs[4],
            entry.regs.gp_regs[5],
        ];
        let result = exit_regs.gp_regs[0];
        let entry = entry.timestamp;

        Self {
            entry,
            exit,
            function,
            args,
            result,
        }
    }
}
