#![feature(
    yggdrasil_os,
    never_type,
    map_try_insert,
    slice_ptr_get,
    iter_chain,
    naked_functions,
    let_chains
)]
#![allow(
    clippy::new_without_default,
    clippy::missing_transmute_annotations,
    clippy::identity_op
)]

use std::{fs, mem::MaybeUninit, os::yggdrasil::real_binary_path, path::Path, process::ExitCode};

use config::Config;
use error::Error;
use object::Object;
use state::ObjectSet;
use yggdrasil_rt::process::{auxv, AuxValue};

pub mod builtins;
pub mod config;
pub mod env;
pub mod error;
pub mod mapping;
pub mod object;
pub mod relocation;
pub mod search;
pub mod state;
pub mod thread_local;

#[derive(Debug)]
pub struct Args {
    pub this: String,
    pub verbose: bool,
    pub args: Vec<String>,
}

static mut OBJECTS: MaybeUninit<ObjectSet> = MaybeUninit::uninit();

fn run(binary: &Path, args: &Args) -> Result<!, Error> {
    let binary_str = binary.to_str().unwrap();
    let mut config = Config::load_or_default()?;
    config.extend_from_env();

    if args.verbose {
        log::info!("Search paths:");
        for path in config.search_paths.iter() {
            log::info!("* {path:?}");
        }
    }

    // Open root and its dependencies
    let root = Object::open(binary, true)?;
    let mut objects = ObjectSet::new(root);
    objects.open_root_dependencies(&config)?;
    let tls_image = objects.load_all()?;
    objects.state.add_magic_symbols();
    objects.resolve_symbols()?;

    if let Err(undefined) = objects.state.no_undefined_symbols() {
        eprintln!("Undefined symbols:");
        for (path, syms) in undefined {
            eprintln!("  in {path:?}");
            for sym in syms {
                eprintln!("    * {sym:?}");
            }
        }
        return Err(Error::UnresolvedSymbols);
    }

    objects.relocate()?;

    let mut auxv = vec![];
    if let Some(image) = tls_image {
        let layout = objects.state.tls_layout.as_ref().unwrap();

        let size = image.size;
        let base = image.data.leak();

        auxv.push(AuxValue {
            tag: auxv::TLS_MASTER_COPY,
            val: base as _,
        });
        auxv.push(AuxValue {
            tag: auxv::TLS_ALIGN,
            val: image.align as _,
        });
        auxv.push(AuxValue {
            tag: auxv::TLS_DATA_SIZE,
            val: size as _,
        });
        auxv.push(AuxValue {
            tag: auxv::TLS_FULL_SIZE,
            val: size as _,
        });

        for module in layout.segments.iter() {
            auxv.push(AuxValue {
                tag: auxv::TLS_MODULE_ID,
                val: module.module_id as u64,
            });
            auxv.push(AuxValue {
                tag: auxv::TLS_MODULE_OFFSET,
                val: module.offset as u64,
            });
        }
    }

    // Set up TLS for the main thread. This needs to be done here, because
    // we still need to call initializers before handing off control to the executable
    // TODO use an ELF note to indicate the extra size required in the TCB
    yggdrasil_rt::process::thread_local::init_tls_from_auxv(auxv.iter(), true)
        .expect("Could not initialize main thread TLS for the executable");

    auxv.push(AuxValue {
        tag: auxv::TLS_ALREADY_INITIALIZED,
        val: 0,
    });

    // Call initializers
    objects.initialize();

    // Setup arguments and enter the main object
    let entry = objects.root.entry().ok_or(Error::NoEntrypoint)?;
    log::info!("entry = {:p}", entry);

    #[allow(static_mut_refs)]
    unsafe {
        OBJECTS.write(objects)
    };

    let argument = env::build_argument(&args.args, binary_str, &auxv)?;
    unsafe { enter(entry, argument) };
}

unsafe fn enter(entry: extern "C" fn(usize), argument: usize) -> ! {
    #[cfg(any(target_arch = "x86", rust_analyzer))]
    {
        core::arch::asm!(
            r#"
                xor %ebp, %ebp
                and $~0xF, %esp
                sub $8, %esp
                pushl {0}
                call *{1}
            "#,
            in(reg) argument,
            in(reg) entry,
            options(att_syntax, noreturn)
        );
    }
    #[cfg(any(target_arch = "riscv64", rust_analyzer))]
    {
        entry(argument);
        unreachable!()
    }
    #[cfg(any(target_arch = "aarch64", target_arch = "x86_64", rust_analyzer))]
    {
        entry(argument);
        unreachable!()
    }
}

fn real_binary(arg0: &str) -> &Path {
    if fs::exists(arg0).unwrap_or(false) {
        return arg0.as_ref();
    }
    real_binary_path()
}

fn usage_no_args(this: &str) {
    eprintln!("{this}: no arguments provided");
    eprintln!("  {this} is an utility to load dynamically-linked programs");
    eprintln!("  and is not intended to be invoked as a standalone program.");
    eprintln!("  If you need to run a program explicitly, the command can be");
    eprintln!("  used liked follows:");
    eprintln!();
    usage(this);
}

fn usage(this: &str) {
    eprintln!("{this} [-v] PROGRAM [ARGS...]");
}

fn get_args() -> Result<Args, (String, String)> {
    let mut args = std::env::args();
    let this = args.next().unwrap();
    let mut verbose = false;
    let mut pargs = vec![];

    let mut collect_flags = true;
    for arg in args {
        if collect_flags && let Some(flag) = arg.strip_prefix("-") {
            match flag {
                "v" => verbose = true,
                _ => return Err((this, format!("unknown flag '-{flag}'"))),
            }
        } else {
            collect_flags = false;
            pargs.push(arg);
        }
    }

    Ok(Args {
        this,
        verbose,
        args: pargs,
    })
}

fn main() -> ExitCode {
    let args = match get_args() {
        Ok(args) => args,
        Err((this, message)) => {
            eprintln!("{this}: {message}");
            eprintln!("Usage:");
            usage(&this);
            return ExitCode::FAILURE;
        }
    };
    if args.args.is_empty() {
        usage_no_args(&args.this);
        return ExitCode::FAILURE;
    }
    logsink::setup_logging(args.verbose);
    let real_binary = real_binary(&args.args[0]);

    let Err(error) = run(real_binary, &args);
    eprintln!("Error: {error}");
    ExitCode::FAILURE
}
