use std::collections::HashMap;
use std::fs;

use nix::sys::uio::{process_vm_readv, RemoteIoVec, process_vm_writev};
use nix::unistd::Pid;
use std::error::Error;
use std::io::{IoSliceMut, IoSlice};
use nix::sys::ptrace;
use std::mem;
fn list_processes() -> Result<HashMap<String, i32>, std::io::Error>
{
    let mut processes = HashMap::<String, i32>::new();

    let entries = fs::read_dir("/proc")?;
    let dirs = entries.filter_map(|e| {
        let e = e.ok()?;
        let path = e.path();
        if path.is_dir() {
            return path.file_name()?.to_str().map(|s| s.to_string());
        }
        None::<String>
    }).collect::<Vec<String>>();

    for dir in dirs {
        let Ok(pid) = dir.parse::<i32>() else { continue };
        let pid_path = format!("/proc/{}/exe", dir);
        let Ok(name_path) = fs::read_link(&pid_path) else { continue };
        let name_path = name_path.to_string_lossy().to_string();

        if let Some(name) = name_path.split("/").last() {
            processes.insert(name.to_string(), pid);
            println!("{} -> {}", name, dir);
        }
    }

    Ok(processes)
}

pub fn get_pid_by_name(name: &str) -> Result<i32, std::io::Error>
{
    let ps = list_processes()?;
    Ok(ps[name])
}

pub fn read_memory_vm(pid: Pid, start_addr: usize, size: usize) -> Result<Vec<u8>, Box<dyn Error>> {
    let mut buffer = vec![0u8; size];

    let mut local_iov = [IoSliceMut::new(&mut buffer)];

    let remote_iov = [RemoteIoVec {
        base: start_addr,
        len: size,
    }];

    let bytes_read = process_vm_readv(pid, &mut local_iov, &remote_iov)?;

    if bytes_read == size {
        Ok(buffer)
    } else {
        buffer.truncate(bytes_read);
        Ok(buffer)
    }
}

pub fn write_memory_vm(pid: Pid, mut start_addr: usize, mut data: &[u8]) -> Result<usize, Box<dyn Error>> {
    let mut total_written = 0usize;
    while !data.is_empty() {
        let local_iov = [IoSlice::new(data)];
        let remote_iov = [RemoteIoVec {
            base: start_addr,
            len: data.len(),
        }];

        let written = process_vm_writev(pid, &local_iov, &remote_iov)?;

        if written == 0 {
            return Err(format!("process_vm_writev returned 0 (no progress) after writing {} bytes", total_written).into());
        }

        total_written += written;
        start_addr = start_addr.wrapping_add(written);
        data = &data[written..];
    }

    Ok(total_written)
}

pub fn write_memory_ptrace(pid: Pid, start_addr: usize, data: &[u8]) -> Result<usize, Box<dyn Error>> {
    let word_size = mem::size_of::<libc::c_long>();
    if word_size == 0 {
        return Err("invalid word size".into());
    }

    let mut addr = start_addr;
    let mut remaining = data;
    let mut written = 0usize;

    if addr % word_size != 0 && !remaining.is_empty() {
        let n = write_unaligned_head(pid, addr, remaining, word_size)?;
        addr += n;
        remaining = &remaining[n..];
        written += n;
    }

    while remaining.len() >= word_size {
        let n = write_full_word(pid, addr, &remaining[..word_size])?;
        addr += n;
        remaining = &remaining[n..];
        written += n;
    }

    if !remaining.is_empty() {
        let n = write_unaligned_tail(pid, addr, remaining, word_size)?;
        written += n;
    }

    Ok(written)
}

fn write_unaligned_head(
    pid: Pid,
    addr: usize,
    data: &[u8],
    word_size: usize,
) -> Result<usize, Box<dyn Error>> {
    let head_offset = addr % word_size;
    let aligned_addr = addr - head_offset;
    let orig_word = ptrace::read(pid, aligned_addr as *mut libc::c_void)?;
    let mut bytes = orig_word.to_ne_bytes();

    let copy_len = usize::min(word_size - head_offset, data.len());
    bytes[head_offset..head_offset + copy_len].copy_from_slice(&data[..copy_len]);
    let new_word = libc::c_long::from_ne_bytes(bytes);

    ptrace::write(pid, aligned_addr as *mut libc::c_void, new_word)?;
    Ok(copy_len)
}

fn write_full_word(pid: Pid, addr: usize, data: &[u8]) -> Result<usize, Box<dyn Error>> {
    let mut arr = [0u8; mem::size_of::<libc::c_long>()];
    arr.copy_from_slice(data);
    let val = libc::c_long::from_ne_bytes(arr);
    ptrace::write(pid, addr as *mut libc::c_void, val)?;
    Ok(mem::size_of::<libc::c_long>())
}

fn write_unaligned_tail(
    pid: Pid,
    addr: usize,
    data: &[u8],
    _word_size: usize,
) -> Result<usize, Box<dyn Error>> {
    let orig_word = ptrace::read(pid, addr as *mut libc::c_void)?;
    let mut bytes = orig_word.to_ne_bytes();
    bytes[..data.len()].copy_from_slice(data);
    let new_word = libc::c_long::from_ne_bytes(bytes);

    ptrace::write(pid, addr as *mut libc::c_void, new_word)?;
    Ok(data.len())
}
