use crate::clibwrap;
use crate::debug_target::{DebugError, PTraceError, RunningTarget, WaitError};
use crate::syscall_info::{syscall_info, SyscallInfo, SyscallInfoError};
use libc::c_long;
use nix::sys::wait::{waitid, Id, WaitPidFlag};
use nix::unistd::Pid;
use std::collections::HashMap;

pub struct StoppedTarget {
	pub pid: Pid,
	pub breakpoints: HashMap<u64, u8>,
}


#[allow(dead_code)]
impl StoppedTarget {
	fn to_running(self) -> RunningTarget {
		RunningTarget { pid: self.pid, breakpoints: self.breakpoints }
	}

	pub fn new(pid: Pid) -> Result<Self, DebugError> {
		waitid(Id::Pid(pid), WaitPidFlag::WSTOPPED).map_err(WaitError)?;
		// Needed for waiting on syscalls and apparently also for getting syscall info (can not get it to work without this)
		clibwrap::ptrace::set_options(pid.as_raw(), clibwrap::ptrace::PTraceOptions::PTRACE_O_TRACESYSGOOD).map_err(PTraceError)?;
		Ok(Self { pid, breakpoints: HashMap::new() })
	}

	pub fn on_breakpoint(&self) -> Result<bool, PTraceError> {
		let rip = self.get_registers()?.rip;
		Ok(self.breakpoints.contains_key(&(rip - 1)))
	}

	fn breakpoint_remove_and_rewind(&mut self) -> Result<(), DebugError> {
		let mut registers = self.get_registers()?;
		self.remove_breakpoint(registers.rip - 1)?;
		registers.rip -= 1;
		clibwrap::ptrace::set_regs(self.pid.as_raw(), registers).map_err(PTraceError)?;
		Ok(())
	}

	pub fn cont(mut self) -> Result<RunningTarget, DebugError> {
		if self.on_breakpoint()? {
			self.breakpoint_remove_and_rewind()?;
		}
		clibwrap::ptrace::cont(self.pid.as_raw()).map_err(PTraceError)?;
		Ok(self.to_running())
	}

	pub fn stepi(mut self) -> Result<RunningTarget, DebugError> {
		if self.on_breakpoint()? {
			self.breakpoint_remove_and_rewind()?;
		}
		clibwrap::ptrace::single_step(self.pid.as_raw()).map_err(PTraceError)?;
		Ok(self.to_running())
	}

	pub fn cont_syscall(mut self) -> Result<RunningTarget, DebugError> {
		if self.on_breakpoint()? {
			self.breakpoint_remove_and_rewind()?;
		}
		clibwrap::ptrace::syscall(self.pid.as_raw()).map_err(PTraceError)?;
		Ok(self.to_running())
	}

	pub fn get_registers(&self) -> Result<clibwrap::ptrace::registers, PTraceError> {
		clibwrap::ptrace::get_regs(self.pid.as_raw()).map_err(PTraceError)
	}

	pub fn get_syscall_info(&self) -> Result<SyscallInfo, SyscallInfoError> {
		Ok(syscall_info(self.pid.as_raw())?)
	}

	#[cfg(target_endian = "little")] // With a bit more work it could be implemented on both architectures, but I only have little-endian computers 🤷
	pub fn add_breakpoint(&mut self, address: u64) -> Result<(), DebugError> {
		if self.breakpoints.contains_key(&address) {
			return Err(DebugError::DuplicateBreakpoint { address });
		}
		let orig_bytes = clibwrap::ptrace::peek_data(self.pid.as_raw(), address).map_err(PTraceError)?;
		let target_byte: u8 = (orig_bytes & 0xFF as c_long) as u8;
		let new_content = (orig_bytes & (!0xff as c_long)) | (0xCC as c_long);
		clibwrap::ptrace::poke_data(self.pid.as_raw(), address, new_content).map_err(PTraceError)?;

		self.breakpoints.insert(address, target_byte);
		Ok(())
	}

	#[cfg(target_endian = "little")] // Same as add_breakpoint
	pub fn remove_breakpoint(&mut self, address: u64) -> Result<(), DebugError> {
		match self.breakpoints.remove(&address) {
			None => Err(DebugError::NonExistingBreakpoint { address }),
			Some(original_byte) => {
				let content = clibwrap::ptrace::peek_data(self.pid.as_raw(), address).map_err(PTraceError)?;
				let new_content = (content & (!0xff as c_long)) | (original_byte as c_long);
				clibwrap::ptrace::poke_data(self.pid.as_raw(), address, new_content).map_err(PTraceError)?;
				Ok(())
			}
		}
	}
}