use std::{
    collections::{HashMap, VecDeque},
    os::unix::process::ExitStatusExt,
    path::PathBuf,
    process::{Command, ExitStatus},
};

pub fn get_target_directory() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("target")
}

pub fn get_hook_lib() -> PathBuf {
    let target = get_target_directory();
    let path = target.join("debug").join("deps").join("libtest_hooklib.so");

    if !path.exists() {
        panic!(
            "The libtest_hooklib.so does not seem to be available at: {}",
            path.display()
        );
    }

    path
}

pub fn get_test_binary(name: &str) -> PathBuf {
    let target = get_target_directory();
    let path = target.join("debug").join("victims").join(name);
    if !path.exists() {
        panic!(
            "Test binary {} does not seem to exist at {}",
            name,
            path.display()
        );
    }

    path
}

pub struct HookCall {
    pub function: String,
    pub args: HashMap<String, String>,
}

pub struct TestHarness {
    pub events: HashMap<String, VecDeque<String>>,
    pub hook_calls: VecDeque<HookCall>,
    pub status: ExitStatus,
    pub raw_stdout: String,
    pub binary: PathBuf,
}

impl TestHarness {
    /// Executes a victim binary with LD_PRELOAD and parses the protocol output.
    pub fn new(bin_name: &str) -> Self {
        let binary = get_test_binary(bin_name);
        let lib = get_hook_lib();

        let output = Command::new(&binary)
            .env("LD_PRELOAD", &lib)
            .output()
            .unwrap_or_else(|_| panic!("Failed to execute victim binary: {}", binary.display()));

        let stdout = String::from_utf8_lossy(&output.stdout);
        let mut events: HashMap<String, VecDeque<String>> = HashMap::new();
        let mut hook_calls = VecDeque::new();

        for line in stdout.lines() {
            if let Some(data) = line.strip_prefix("__HOOK_CALL__:") {
                let mut parts = data.splitn(2, '|');
                let fn_name = parts.next().unwrap().to_string();
                let args_raw = parts.next().unwrap_or("");

                let args = args_raw
                    .split(',')
                    .filter(|s| !s.is_empty())
                    .filter_map(|s| {
                        let mut kv = s.splitn(2, '=');
                        Some((kv.next()?.to_string(), kv.next()?.to_string()))
                    })
                    .collect();

                hook_calls.push_back(HookCall {
                    function: fn_name,
                    args,
                });
            } else if let Some(data) = line.strip_prefix("__TEST_EVENT__:") {
                let mut parts = data.splitn(2, '=');
                if let (Some(k), Some(v)) = (parts.next(), parts.next()) {
                    events
                        .entry(k.to_string())
                        .or_default()
                        .push_back(v.to_string());
                }
            }
        }

        Self {
            events,
            hook_calls,
            status: output.status,
            raw_stdout: stdout.to_string(),
            binary,
        }
    }

    /// Returns objdump of the binary if available.
    fn get_binary_objdump(&self) -> Option<String> {
        // run object dump on the binary
        let objdump = Command::new("objdump")
            .arg("-d")
            .arg(&self.binary)
            .output()
            .ok()?;
        if objdump.status.success() {
            Some(String::from_utf8_lossy(&objdump.stdout).to_string())
        } else {
            None
        }
    }

    // --- Assertions ---

    pub fn assert_success(&mut self) -> &mut Self {
        // if program crashed with segfault, include objdump for debugging
        match self.status.signal() {
            Some(code) if code == libc::SIGSEGV => {
                if let Some(objdump) = self.get_binary_objdump() {
                    panic!(
                        "Program crashed with SIGSEGV!\n\n\
                        Stdout:\n{}\n\n\
                        Objdump:\n{}\n\n\
                        ",
                        self.raw_stdout, objdump
                    );
                }
            }
            _ => {}
        }

        assert!(
            self.status.success(),
            "Process failed({})! Stdout:\n{}",
            self.status,
            self.raw_stdout
        );
        self
    }

    pub fn assert_failed(&mut self) -> &mut Self {
        // On Unix, a crash usually results in no exit code but a signal
        assert!(
            !self.status.success(),
            "Process was expected to crash but succeeded!"
        );

        #[cfg(unix)]
        {
            use std::os::unix::process::ExitStatusExt;
            assert!(
                self.status.signal().is_some(),
                "Process exited normally but was expected to be killed by a signal (crash)."
            );
        }
        self
    }

    pub fn assert_event(&mut self, key: &str, expected: &str) -> &mut Self {
        let vals = self
            .events
            .get_mut(key)
            .unwrap_or_else(|| panic!("Event '{}' not found in output!\n{}", key, self.raw_stdout));
        let val = vals
            .pop_front()
            .unwrap_or_else(|| panic!("Event '{}' had no more values to consume!", key));
        assert_eq!(val, expected, "Event value mismatch for '{}'", key);
        self
    }

    pub fn assert_hook_call(&mut self, fn_name: &str, expected_args: &[(&str, &str)]) -> &mut Self {
        let call = self
            .hook_calls
            .pop_front()
            .unwrap_or_else(|| panic!("Expected hook call to '{}', but none recorded!", fn_name));

        assert_eq!(call.function, fn_name);
        for (key, val) in expected_args {
            let actual = call
                .args
                .get(*key)
                .unwrap_or_else(|| panic!("Arg '{}' missing in hook '{}'", key, fn_name));
            assert_eq!(actual, *val);
        }
        self
    }
}

#[ignore]
#[test]
fn interposing_test() {
    let mut harness = TestHarness::new("test_interpose");

    harness
        .assert_success()
        .assert_hook_call("getenv", &[("name", "YOU_ARE_MY_FOO")])
        .assert_event("foo", "BAR")
        .assert_hook_call("getenv", &[("name", "LD_PRELOAD")]);

    let lib = get_hook_lib();
    harness.assert_event("path", &lib.display().to_string());
}

#[ignore]
#[test]
fn retn_hook_test() {
    let mut harness = TestHarness::new("test_basic");

    harness
        .assert_success()
        // Verify the hook captured the inputs
        .assert_hook_call("add", &[("a", "10"), ("b", "20")])
        // Verify the original logic executed correctly
        .assert_event("original_add_a", "10")
        .assert_event("original_add_b", "20")
        .assert_hook_call("add", &[("result", "30")])
        // Verify the hook successfully modified the return value
        .assert_event("add_result", "35");
}

#[ignore]
#[test]
fn hook_guard_test() {
    let mut harness = TestHarness::new("test_hook_guard");

    harness
        .assert_success()
        .assert_hook_call("hook_guard", &[("flag", "1000")])
        .assert_event("value", "-1337");
}

#[ignore]
#[test]
fn jmp_back_hook_test() {
    let mut harness = TestHarness::new("test_jmp_back");
    harness
        .assert_success()
        .assert_hook_call("compute", &[("x", "5"), ("y", "3")])
        .assert_event("compute_result", "999");
}

#[ignore]
#[test]
fn jmp_to_addr_hook_test() {
    let mut harness = TestHarness::new("test_jmp_to_addr");
    harness
        .assert_success()
        .assert_event("hook_security_gate", "called")
        .assert_event("result", "100")
        .assert_event("security_gate", "access granted");
}

#[ignore]
#[test]
fn jmp_to_ret_hook_test() {
    let mut harness = TestHarness::new("test_jmp_to_ret");
    harness
        .assert_success()
        .assert_event("player_id", "4660") // 0x1234
        .assert_hook_call("damage_hook", &[("player_id", "4660")])
        .assert_event("hit_sound", "oof.mp3")
        .assert_event("update_ui_id", "4660")
        .assert_event("update_ui_hp", "0")
        .assert_event("hit_points", "0")
        .assert_event("player_id", "17767") // 0x4567
        .assert_hook_call("damage_hook", &[("player_id", "17767")])
        .assert_event("cheat", "skip damage")
        .assert_event("hit_sound", "oof.mp3")
        .assert_event("update_ui_id", "17767")
        .assert_event("update_ui_hp", "100")
        .assert_event("hit_points", "100");
}
