//! 通用独立可执行文件加载器
//!
//! 预编译的 loader，从同目录的 .tsnobj 文件加载编译数据并执行。
//! 用法：tsn-standalone-loader.exe <entry_name>
//! 或：直接运行，自动查找同目录的 .tsnobj 文件

use std::ffi::{c_void, CString};
use std::os::windows::ffi::OsStrExt;
use std::path::{Path, PathBuf};
use std::ptr;

#[link(name = "kernel32")]
extern "system" {
    fn LoadLibraryW(lp_lib_file_name: *const u16) -> *mut c_void;
    fn GetProcAddress(h_module: *mut c_void, lp_proc_name: *const i8) -> *mut c_void;
    fn FreeLibrary(h_lib_module: *mut c_void) -> i32;
    fn VirtualAlloc(lpAddress: *mut c_void, dwSize: usize, flAllocationType: u32, flProtect: u32) -> *mut c_void;
    fn VirtualFree(lpAddress: *mut c_void, dwSize: usize, dwFreeType: u32) -> i32;
}

const MEM_COMMIT: u32 = 0x1000;
const MEM_RESERVE: u32 = 0x2000;
const MEM_RELEASE: u32 = 0x8000;
const PAGE_EXECUTE_READWRITE: u32 = 0x40;
const PAGE_READWRITE: u32 = 0x04;

#[repr(C)]
#[derive(Clone)]
struct Symbol {
    section: u8,
    offset: u32,
    name: [u8; 64],
}

#[repr(C)]
#[derive(Clone)]
struct Relocation {
    section: u8,
    offset: u32,
    type_: u8,
    _padding: u8,
    target_name: [u8; 64],
}

/// 模块数据结构（与 standalone_gen.rs 对应）
#[repr(C)]
struct ModuleData {
    code_len: u32,
    data_len: u32,
    rdata_len: u32,
    symbol_count: u32,
    reloc_count: u32,
    entry_name: [u8; 64],
    // 数据紧跟在结构体后面：
    // code[code_len]
    // data[data_len]
    // rdata[rdata_len]
    // symbols[symbol_count]
    // relocations[reloc_count]
}

fn find_runtime_dll(exe_dir: &Path) -> Option<PathBuf> {
    let local = exe_dir.join("tsn_runtime.dll");
    if local.exists() { return Some(local); }
    let target_release = exe_dir.join("..").join("..").join("target").join("release").join("tsn_runtime.dll");
    if target_release.exists() { return Some(target_release); }
    let current = PathBuf::from("tsn_runtime.dll");
    if current.exists() { return Some(current); }
    None
}

fn find_tsnobj_file(exe_dir: &Path, entry_name: &str) -> Option<PathBuf> {
    // 1. 同目录 <entry_name>.tsnobj
    let direct = exe_dir.join(format!("{}.tsnobj", entry_name));
    if direct.exists() { return Some(direct); }
    // 2. 同目录下与 exe 同名的 .tsnobj
    None
}

fn load_module_data(path: &Path) -> Option<(Vec<u8>, Vec<u8>, Vec<u8>, Vec<Symbol>, Vec<Relocation>, String)> {
    let bytes = std::fs::read(path).ok()?;
    if bytes.len() < std::mem::size_of::<ModuleData>() {
        return None;
    }

    let header = unsafe { &*(bytes.as_ptr() as *const ModuleData) };
    let code_len = header.code_len as usize;
    let data_len = header.data_len as usize;
    let rdata_len = header.rdata_len as usize;
    let symbol_count = header.symbol_count as usize;
    let reloc_count = header.reloc_count as usize;

    let header_size = std::mem::size_of::<ModuleData>();
    let mut offset = header_size;

    let code = bytes.get(offset..offset + code_len)?.to_vec();
    offset += code_len;

    let data = bytes.get(offset..offset + data_len)?.to_vec();
    offset += data_len;

    let rdata = bytes.get(offset..offset + rdata_len)?.to_vec();
    offset += rdata_len;

    let symbols_size = symbol_count * std::mem::size_of::<Symbol>();
    let symbols_bytes = bytes.get(offset..offset + symbols_size)?;
    offset += symbols_size;

    let relocs_size = reloc_count * std::mem::size_of::<Relocation>();
    let relocs_bytes = bytes.get(offset..offset + relocs_size)?;

    let symbols: Vec<Symbol> = unsafe {
        std::slice::from_raw_parts(symbols_bytes.as_ptr() as *const Symbol, symbol_count).to_vec()
    };
    let relocations: Vec<Relocation> = unsafe {
        std::slice::from_raw_parts(relocs_bytes.as_ptr() as *const Relocation, reloc_count).to_vec()
    };

    let entry_name = std::str::from_utf8(&header.entry_name)
        .unwrap_or("")
        .trim_end_matches('\0')
        .to_string();

    Some((code, data, rdata, symbols, relocations, entry_name))
}

pub fn run_standalone(exe_path: &Path, obj_path: &Path) -> Result<f64, String> {
    let exe_dir = exe_path.parent().unwrap_or(Path::new("."));

    // 1. 加载 .tsnobj 数据
    let (code, data, rdata, symbols, relocations, entry_name) = load_module_data(obj_path)
        .ok_or("Failed to load .tsnobj file")?;

    println!("[Standalone] Entry: {}", entry_name);
    println!("[Standalone] Code: {} bytes, Data: {} bytes, RData: {} bytes",
        code.len(), data.len(), rdata.len());
    println!("[Standalone] Symbols: {}, Relocations: {}", symbols.len(), relocations.len());

    // 2. 加载 runtime DLL
    let runtime_path = find_runtime_dll(exe_dir)
        .ok_or("tsn_runtime.dll not found")?;
    println!("[Runtime] Loading: {}", runtime_path.display());

    let wide_path: Vec<u16> = runtime_path
        .as_os_str()
        .encode_wide()
        .chain(Some(0))
        .collect();
    let hmod = unsafe { LoadLibraryW(wide_path.as_ptr()) };
    if hmod.is_null() {
        return Err("Failed to load runtime DLL".to_string());
    }

    // 3. 分配可执行内存
    let code_size = code.len();
    if code_size == 0 {
        unsafe { FreeLibrary(hmod); }
        return Err("Code section is empty".to_string());
    }

    let exec_mem = unsafe {
        VirtualAlloc(ptr::null_mut(), code_size, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE)
    };
    if exec_mem.is_null() {
        unsafe { FreeLibrary(hmod); }
        return Err("VirtualAlloc failed for code".to_string());
    }

    unsafe {
        ptr::copy_nonoverlapping(code.as_ptr(), exec_mem as *mut u8, code_size);
    }

    // 4. 分配数据内存
    let total_data_size = data.len() + rdata.len();
    let data_mem = if total_data_size > 0 {
        let mem = unsafe {
            VirtualAlloc(ptr::null_mut(), total_data_size, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE)
        };
        if mem.is_null() {
            unsafe { VirtualFree(exec_mem, 0, MEM_RELEASE); FreeLibrary(hmod); }
            return Err("VirtualAlloc failed for data".to_string());
        }
        if !data.is_empty() {
            unsafe { ptr::copy_nonoverlapping(data.as_ptr(), mem as *mut u8, data.len()); }
        }
        if !rdata.is_empty() {
            unsafe {
                ptr::copy_nonoverlapping(
                    rdata.as_ptr(),
                    (mem as *mut u8).add(data.len()),
                    rdata.len()
                );
            }
        }
        mem
    } else {
        ptr::null_mut()
    };

    // 5. 应用重定位
    let mut reloc_ok = 0usize;
    let mut reloc_fail = 0usize;

    for reloc in &relocations {
        let target_name = std::str::from_utf8(&reloc.target_name)
            .unwrap_or("")
            .trim_end_matches('\0');

        if target_name.is_empty() {
            reloc_fail += 1;
            continue;
        }

        let target_addr = if target_name.starts_with("js_") || target_name.starts_with("__") {
            let c_name = match CString::new(target_name) {
                Ok(s) => s,
                Err(_) => {
                    eprintln!("[WARN] Invalid symbol name: {}", target_name);
                    reloc_fail += 1;
                    continue;
                }
            };
            let addr = unsafe { GetProcAddress(hmod, c_name.as_ptr() as *const i8) };
            if addr.is_null() {
                eprintln!("[WARN] Unresolved external symbol: {}", target_name);
                reloc_fail += 1;
                continue;
            }
            addr
        } else {
            let mut found = None;
            for sym in &symbols {
                let sym_name = std::str::from_utf8(&sym.name)
                    .unwrap_or("")
                    .trim_end_matches('\0');
                if sym_name == target_name {
                    found = Some((sym.section, sym.offset));
                    break;
                }
            }
            match found {
                Some((0, offset)) => (exec_mem as usize + offset as usize) as *mut c_void,
                Some((1, offset)) => {
                    if data_mem.is_null() { reloc_fail += 1; continue; }
                    (data_mem as usize + offset as usize) as *mut c_void
                }
                Some((2, offset)) => {
                    if data_mem.is_null() { reloc_fail += 1; continue; }
                    (data_mem as usize + data.len() + offset as usize) as *mut c_void
                }
                _ => {
                    eprintln!("[WARN] Unresolved internal symbol: {}", target_name);
                    reloc_fail += 1;
                    continue;
                }
            }
        };

        let patch_addr = (exec_mem as usize + reloc.offset as usize) as *mut u8;
        match reloc.type_ {
            0 => {
                let rel = (target_addr as isize) - (patch_addr as isize + 4);
                unsafe { *(patch_addr as *mut i32) = rel as i32; }
            }
            1 => {
                unsafe { *(patch_addr as *mut u64) = target_addr as u64; }
            }
            _ => { reloc_fail += 1; continue; }
        }
        reloc_ok += 1;
    }

    println!("[Reloc] {}/{} ok, {} failed", reloc_ok, relocations.len(), reloc_fail);

    // 6. 查找入口点并执行
    let mut entry_offset = 0u32;
    for sym in &symbols {
        let sym_name = std::str::from_utf8(&sym.name)
            .unwrap_or("")
            .trim_end_matches('\0');
        if sym_name == "main" && sym.section == 0 {
            entry_offset = sym.offset;
            break;
        }
    }

    println!("[Exec] Calling entry point (main @ offset {})...", entry_offset);
    let result = unsafe {
        let entry_addr = (exec_mem as usize + entry_offset as usize) as *const u8;
        let fn_ptr: extern "C" fn() -> f64 = std::mem::transmute(entry_addr);
        fn_ptr()
    };

    // 7. 清理
    unsafe { VirtualFree(exec_mem, 0, MEM_RELEASE); }
    if !data_mem.is_null() {
        unsafe { VirtualFree(data_mem, 0, MEM_RELEASE); }
    }
    unsafe { FreeLibrary(hmod); }

    Ok(result)
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let exe_path = std::env::current_exe().unwrap_or_default();
    let exe_dir = exe_path.parent().unwrap_or(Path::new("."));

    // 确定入口名和 .tsnobj 文件路径
    let (entry_name, obj_path) = if args.len() > 1 {
        let name = &args[1];
        let path = exe_dir.join(format!("{}.tsnobj", name));
        (name.clone(), path)
    } else {
        // 尝试查找同目录下与 exe 同名的 .tsnobj
        let stem = exe_path.file_stem()
            .map(|s| s.to_string_lossy().to_string())
            .unwrap_or_else(|| "main".to_string());
        let path = exe_dir.join(format!("{}.tsnobj", stem));
        (stem, path)
    };

    if !obj_path.exists() {
        eprintln!("[ERROR] .tsnobj file not found: {}", obj_path.display());
        eprintln!("Usage: {} <entry_name>", exe_path.display());
        std::process::exit(1);
    }

    match run_standalone(&exe_path, &obj_path) {
        Ok(result) => {
            println!("\n=== Result: {} ===", result);
        }
        Err(e) => {
            eprintln!("[ERROR] {}", e);
            std::process::exit(1);
        }
    }
}
