use std::io::{IsTerminal, Write};
use std::path::Path;
use std::thread::sleep;
use std::time::{Duration, Instant};
use std::{fs, process};

use clap::{CommandFactory, Parser};
use color_eyre::eyre::{Context, ContextCompat, Result, bail};
use colored_json::to_colored_json_auto;
use comfy_table::{Row, Table};
use daemonize_me::Daemon;
use dir_lock::DirLock;
use humansize::{BINARY, FormatSizeOptions, format_size_i};
use termion::{color, style};
use tokio::runtime::Runtime;
use tracing::{debug, instrument, trace};
use walkdir::WalkDir;

mod cli;
mod qemu;
mod runner;
mod ssh;
mod utils;
mod vm_images;
mod vms;

use crate::cli::{
    CleanCommand, Command, ExecCommand, ImagesCommand, InspectCommand, Interactive, KsmCommand,
    LogsCommand, OsTypeOrImagePath, PruneCommand, PsCommand, PullCommand, ReadyCommand, RunCommand,
    StopCommand,
};
use crate::ssh::{connect_ssh_for_command, create_ssh_connection, ensure_ssh_key};
use crate::vms::{
    DAEMON_LOG_FILENAME, SESSION_LOG_FILENAME, VMEXEC_PID_FILENAME, VmId, VmState, get_vms,
    reap_dead_vm_dirs,
};
use crate::{
    utils::{VmexecDirs, check_ksm_active, find_required_tools, install_tracing, print_ksm_stats},
    vm_images::{OVERLAY_IMAGE_EXTENSION, VmImage, get_images},
};

#[instrument]
fn ps_command(ps_args: PsCommand) -> Result<()> {
    let dirs = VmexecDirs::new()?;

    let lock_dir = &dirs.lock_dir;
    trace!("Trying to lock {lock_dir:?}");
    let lock = DirLock::new_sync(&lock_dir)?;

    let mut table = Table::new();
    table.load_preset(comfy_table::presets::NOTHING);
    table.set_header(["VM ID", "IMAGE", "COMMAND", "CREATED", "STATUS", "PORTS"]);
    table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);

    for column in table.column_iter_mut() {
        column.set_padding((0, 2));
    }

    let vms = get_vms(&dirs.vms_dir)?;
    for vm in vms.values() {
        if let Some(vm_config) = &vm.vm_config {
            // Conditionally also show dead VMs.
            if !ps_args.all && vm.state() == VmState::Exited {
                continue;
            }

            let created_at = timeago::Formatter::new().convert(vm_config.created_at.elapsed()?);

            let vmid = vm_config.vmid.clone();
            let status = vm.state();
            let image = vm_config
                .qemu_launch_opts
                .vm_image
                .image_path
                .file_name()
                .unwrap()
                .to_string_lossy()
                .to_string();
            let command = vm_config.ssh_launch_opts.args.join(" ");
            let ports = vm_config
                .qemu_launch_opts
                .published_ports
                .iter()
                .map(|p| p.to_string())
                .collect::<Vec<String>>()
                .join(", ");

            let mut row = Row::from(&[
                vmid.into_inner(),
                image,
                command,
                created_at,
                status.to_string(),
                ports,
            ]);
            row.max_height(1);
            table.add_row(row);
        } else {
            continue;
        }
    }

    println!("{table}");

    trace!("Unlocking {lock_dir:?}");
    drop(lock);

    Ok(())
}

#[instrument]
fn inspect_command(inspect_args: InspectCommand) -> Result<()> {
    let dirs = VmexecDirs::new()?;

    let vms = get_vms(&dirs.vms_dir)?;
    let vm = &vms[&inspect_args.vmid];

    let vm_config = vm
        .vm_config
        .as_ref()
        .wrap_err("VM doesn't have a config - maybe it's not well-formed?")?;

    println!("{}", to_colored_json_auto(vm_config)?);

    Ok(())
}

#[instrument]
fn images_command(_images_args: ImagesCommand) -> Result<()> {
    let dirs = VmexecDirs::new()?;

    let mut table = Table::new();
    table.load_preset(comfy_table::presets::NOTHING);
    table.set_header(["OS", "VERSION", "PREPARED", "CREATED", "SIZE"]);
    table.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);

    for column in table.column_iter_mut() {
        column.set_padding((0, 2));
    }

    for (os, versions) in get_images(&dirs.images_dir)? {
        for image in versions {
            let prepared = if image.prepared { "yes" } else { "no" };

            let created_at = timeago::Formatter::new().convert(image.created_at.elapsed()?);

            let format_size_options = FormatSizeOptions::from(BINARY).decimal_places(1);
            let size = format_size_i(image.size as f64, format_size_options);

            table.add_row([
                os.clone(),
                image.version,
                prepared.to_string(),
                created_at,
                size,
            ]);
        }
    }

    println!("{table}");

    Ok(())
}

#[instrument]
fn logs_command(logs_args: LogsCommand) -> Result<()> {
    let dirs = VmexecDirs::new()?;

    let vms = get_vms(&dirs.vms_dir)?;
    let session_log_path = vms[&logs_args.vmid].vm_dir.join(SESSION_LOG_FILENAME);

    let mut session_log = fs::File::open(&session_log_path)
        .wrap_err(format!("Couldn't open session log at {session_log_path:?}"))?;
    let mut stdout = std::io::stdout();

    // Print everything that's currently in the log.
    std::io::copy(&mut session_log, &mut stdout)?;
    stdout.flush()?;

    if logs_args.follow {
        // Keep printing new output as it arrives until the VM exits, like `tail -f`.
        loop {
            // Only stop once a read that happened *after* the VM was already seen dead returns
            // nothing. If we instead checked the state after an empty read, output written
            // between that read and the check would be lost.
            let vm_exited = get_vms(&dirs.vms_dir)?[&logs_args.vmid].state() == VmState::Exited;

            let copied = std::io::copy(&mut session_log, &mut stdout)?;
            stdout.flush()?;

            if copied == 0 {
                if vm_exited {
                    break;
                }
                sleep(Duration::from_millis(250));
            }
        }
    }

    Ok(())
}

fn stop_command(stop_args: StopCommand) -> Result<()> {
    let dirs = VmexecDirs::new()?;

    let vms = get_vms(&dirs.vms_dir)?;
    let vm = &vms[&stop_args.vmid];
    let qemu_pid = vm
        .cid_pids
        .qemu_pid
        .expect("Somehow qemu.pid doesn't refer to a live process");
    let vmexec_pid = vm
        .cid_pids
        .vmexec_pid
        .expect("Somehow vmexec.pid doesn't refer to a live process");

    debug!("Found {vmexec_pid} to be live, sending kill signal");
    nix::sys::signal::kill(
        nix::unistd::Pid::from_raw(vmexec_pid as i32),
        nix::sys::signal::SIGTERM,
    )?;

    debug!("Found {qemu_pid} to be live, sending kill signal");
    nix::sys::signal::kill(
        nix::unistd::Pid::from_raw(qemu_pid as i32),
        nix::sys::signal::SIGTERM,
    )?;

    let rm_flag = vm.vm_config.as_ref().map(|c| c.rm).unwrap_or(false);

    // Clean up after running in case this has been requested.
    if rm_flag {
        debug!("Waiting for processes to exit before cleanup...");
        let mut retries = 0;
        loop {
            let current_vms = get_vms(&dirs.vms_dir)?;
            let current_vm = &current_vms[&stop_args.vmid];
            if current_vm.cid_pids.qemu_pid.is_none() && current_vm.cid_pids.vmexec_pid.is_none() {
                break;
            }
            if retries > 20 {
                debug!("Processes still running after waiting, proceeding with cleanup anyway");
                break;
            }
            sleep(Duration::from_millis(100));
            retries += 1;
        }

        debug!("Removing VM dir at {:?}", &vm.vm_dir);
        fs::remove_dir_all(&vm.vm_dir)?;
    }

    Ok(())
}

#[instrument]
fn exec_command(exec_args: ExecCommand) -> Result<()> {
    let dirs = VmexecDirs::new()?;

    let ssh_keypair = ensure_ssh_key(&dirs.secrets_dir)?;
    let vms = get_vms(&dirs.vms_dir)?;
    let cid = vms[&exec_args.vmid]
        .cid_pids
        .cid
        .wrap_err("cid didn't exist")?;
    let session_log_path = vms[&exec_args.vmid].vm_dir.join(SESSION_LOG_FILENAME);

    let tty = match exec_args.tty {
        cli::Tty::Always => true,
        cli::Tty::Never => false,
        cli::Tty::Auto => std::io::stdout().is_terminal(),
    };

    debug!(
        "tty: {tty}, interactive: {interactive:?}",
        interactive = exec_args.interactive
    );

    let ssh_launch_opts = ssh::SshLaunchOpts {
        timeout: exec_args.ssh_timeout,
        tty,
        interactive: exec_args.interactive,
        env_vars: exec_args.env,
        workdir: exec_args.workdir,
        args: exec_args.args,
        privkey: ssh_keypair.privkey_str,
        cid,
    };

    let rt = Runtime::new()?;
    rt.block_on(connect_ssh_for_command(
        None,
        ssh_launch_opts,
        session_log_path,
    ))?;
    Ok(())
}

#[instrument]
fn pull_command(pull_args: PullCommand) -> Result<()> {
    let dirs = VmexecDirs::new()?;

    // Hold the lock so that we don't race a concurrent `run`'s download of the same image
    // version as both would write to the same paths in the cache.
    let lock_dir = &dirs.lock_dir;
    trace!("Trying to lock {lock_dir:?}");
    let lock = DirLock::new_sync(lock_dir)?;

    let rt = Runtime::new()?;
    let vm_image = rt.block_on(async {
        match pull_args.os_type {
            cli::OsType::Archlinux => VmImage::archlinux(&dirs.images_dir, cli::Pull::Newer).await,
        }
    })?;
    println!("Image available at {}", vm_image.image_path.display());

    if pull_args.prepare {
        let tool_paths = find_required_tools()?;
        let ssh_keypair = ensure_ssh_key(&dirs.secrets_dir)?;

        // The warmup needs a VM dir just like a regular run. It's thrown away afterwards.
        let vmid = VmId::default();
        let vm_dir = Path::new(&dirs.vms_dir).join(vmid.as_str());
        fs::create_dir(&vm_dir)
            .wrap_err(format!("Couldn't make temp dir in {:?}", dirs.vms_dir))?;

        let prepare_result = rt.block_on(runner::prepare_image(
            &dirs,
            &lock,
            tool_paths,
            &vm_dir,
            vm_image,
            ssh_keypair,
            runner::PrepareOpts {
                ssh_timeout: pull_args.ssh_timeout,
                memory: pull_args.memory.unwrap_or_default().gb,
                disable_kvm: pull_args.disable_kvm,
                show_vm_window: pull_args.show_vm_window,
            },
        ));

        // Clean up the throwaway VM dir whether or not the preparation worked.
        fs::remove_dir_all(&vm_dir)?;
        prepare_result?;

        println!("Image prepared");
    }

    drop(lock);
    Ok(())
}

#[instrument]
fn run_command(run_args: RunCommand) -> Result<Option<u32>> {
    let dirs = VmexecDirs::new()?;

    // Make sure the tools we need are actually installed.
    let tool_paths = find_required_tools()?;

    // Check whether KSM is active.
    check_ksm_active();

    // Lock the VMs dir here. We have to keep the lock around until QEMU has launched for this
    // run. If we unlock too early (before there's a live qemu.pid) then a second launch that
    // happens right after this one might try to clean up the "dead" dir. We don't want this race
    // condition.
    let lock_dir = &dirs.lock_dir;
    trace!("Trying to lock {lock_dir:?}");
    let lock = DirLock::new_sync(&lock_dir)?;

    // Dir for this VM (usually ~/.local/state/vmexec/vms/<random id>/)
    // Supposed to get cleaned up after QEMU stops running.
    let vmid = run_args.vmid.unwrap_or_default();

    // Check for collision with existing VMs
    let existing_vms = get_vms(&dirs.vms_dir)?;
    if existing_vms.contains_key(&vmid) {
        bail!("vmid '{}' already exists", vmid.as_str());
    }

    let vm_dir = Path::new(&dirs.vms_dir).join(vmid.as_str());
    fs::create_dir(&vm_dir).wrap_err(format!("Couldn't make temp dir in {:?}", dirs.vms_dir))?;
    debug!("vm dir is: {:?}", vm_dir);

    let pid_file = vm_dir.join(VMEXEC_PID_FILENAME);
    if run_args.detach {
        // In case the user chooses to detach this, we'll daemonize at this point. That means the main
        // process will exit while the child process will continue running the commands.

        debug!("Detach was requested so we're now daemonizing");
        println!("{vmid}");

        // Log daemon stdout and stderr into same file.
        let daemon_log = fs::File::create(vm_dir.join(DAEMON_LOG_FILENAME))?;
        Daemon::new()
            .pid_file(pid_file, None)
            .stdout(daemon_log.try_clone()?)
            .stderr(daemon_log)
            .setup_post_fork_parent_hook(post_fork_parent)
            .start()?;
    } else {
        // If we're not daemonizing, we're just writing our pid out and that's it.

        fs::write(pid_file, process::id().to_string())?;
    }

    let rt = Runtime::new()?;
    let vm_image = rt.block_on(async {
        match run_args.image_source {
            OsTypeOrImagePath::OsType(os_type) => match os_type {
                cli::OsType::Archlinux => VmImage::archlinux(&dirs.images_dir, run_args.pull).await,
            },
            OsTypeOrImagePath::ImagePath(image_path) => VmImage::generic(&image_path),
        }
    })?;

    let ssh_keypair = ensure_ssh_key(&dirs.secrets_dir)?;

    let memory = run_args.memory.unwrap_or_default().gb;
    let prepared = rt.block_on(runner::prepare_image(
        &dirs,
        &lock,
        tool_paths.clone(),
        &vm_dir,
        vm_image.clone(),
        ssh_keypair.clone(),
        runner::PrepareOpts {
            ssh_timeout: run_args.ssh_timeout,
            memory,
            disable_kvm: run_args.disable_kvm,
            show_vm_window: run_args.show_vm_window,
        },
    ))?;

    // We launch QEMU from the overlay image instead of the source image.
    let vm_image = VmImage {
        image_path: prepared.overlay_image_path,
        ..vm_image
    };
    let qemu_launch_opts = qemu::QemuLaunchOpts {
        volumes: run_args.volumes,
        pmems: run_args.pmems,
        published_ports: run_args.published_ports,
        vm_image,
        ovmf_uefi_vars_path: prepared.ovmf_vars,
        show_vm_window: run_args.show_vm_window,
        pubkey: ssh_keypair.pubkey_str,
        cid: prepared.cid,
        is_warmup: false,
        disable_kvm: run_args.disable_kvm,
        memory,
    };

    let tty = match run_args.tty {
        cli::Tty::Always => true,
        cli::Tty::Never => false,
        cli::Tty::Auto => std::io::stdout().is_terminal(),
    };

    debug!(
        "tty: {tty}, interactive: {interactive:?}",
        interactive = run_args.interactive
    );

    let ssh_launch_opts = ssh::SshLaunchOpts {
        timeout: run_args.ssh_timeout,
        tty,
        interactive: run_args.interactive,
        env_vars: run_args.env,
        workdir: run_args.workdir,
        args: run_args.args,
        privkey: ssh_keypair.privkey_str,
        cid: prepared.cid,
    };

    debug!("SSH command for manual debugging:");
    debug!(
        "ssh root@vsock/{cid} -i {privkey_path:?}",
        cid = prepared.cid,
        privkey_path = ssh_keypair.privkey_path
    );

    let exit_code = rt.block_on(runner::run(
        vmid.clone(),
        &vm_dir,
        lock,
        tool_paths,
        qemu_launch_opts,
        ssh_launch_opts,
        run_args.rm,
    ))?;

    // Clean up after running in case this has been requested.
    if run_args.rm {
        debug!("Removing VM dir at {vm_dir:?}");
        fs::remove_dir_all(&vm_dir)?;
    }

    Ok(exit_code)
}

#[instrument]
fn ksm_command(ksm_args: KsmCommand) -> Result<()> {
    if let Some(enable_disable) = ksm_args.ksm_enable_disable {
        if !nix::unistd::getuid().is_root() {
            bail!("You need to run this particular subcommand as root");
        }

        if enable_disable.enable {
            println!("Writing KSM config to /etc/tmpfiles.d/ksm.conf and reloading systemd");
            let ksm_conf = "\
w /sys/kernel/mm/ksm/run - - - - 1
w /sys/kernel/mm/ksm/advisor_mode - - - - scan-time
";
            fs::write("/etc/tmpfiles.d/ksm.conf", ksm_conf)?;
            process::Command::new("systemd-tmpfiles")
                .arg("--create")
                .output()?;
        } else if enable_disable.disable {
            println!("Removing KSM config at /etc/tmpfiles.d/ksm.conf and reloading systemd");
            fs::write("/sys/kernel/mm/ksm/run", "0")?;
            fs::write("/sys/kernel/mm/ksm/advisor_mode", "none")?;
            fs::remove_file("/etc/tmpfiles.d/ksm.conf")?;
        }
    } else {
        let ksm_enabled = fs::read_to_string("/sys/kernel/mm/ksm/run")?.trim() == "1";
        if ksm_enabled {
            println!(
                "{}KSM status: {}enabled{}",
                style::Bold,
                color::Fg(color::LightGreen),
                style::Reset,
            );
            print_ksm_stats()?;
        } else {
            println!(
                "{}KSM status: {}disabled{}",
                style::Bold,
                color::Fg(color::Yellow),
                style::Reset
            );
        }
    }
    Ok(())
}

#[instrument]
fn clean_command(_clean_args: CleanCommand) -> Result<()> {
    let dirs = VmexecDirs::new()?;
    let cleaned_vm_dirs = reap_dead_vm_dirs(&dirs.vms_dir)?;
    if !cleaned_vm_dirs.is_empty() {
        println!(
            "Removed old VMs:\n  {}",
            cleaned_vm_dirs
                .iter()
                .map(|dir| dir.to_string_lossy())
                .collect::<Vec<_>>()
                .join("\n  ")
        );
    } else {
        println!("Nothing to do!");
    }

    Ok(())
}

#[instrument]
fn prune_command(_prune_args: PruneCommand) -> Result<()> {
    let dirs = VmexecDirs::new()?;

    let mut deleted_files = Vec::new();

    for entry in WalkDir::new(&dirs.cache_dir) {
        let entry = entry?;
        let path = entry.path();
        if path.is_file()
            && let Some(name) = path.file_name()
            && name
                .to_string_lossy()
                .ends_with(&format!(".{OVERLAY_IMAGE_EXTENSION}"))
        {
            fs::remove_file(path)?;
            deleted_files.push(path.to_path_buf());
        }
    }

    if !deleted_files.is_empty() {
        println!(
            "Removed overlay images:\n  {}",
            deleted_files
                .iter()
                .map(|dir| dir.to_string_lossy())
                .collect::<Vec<_>>()
                .join("\n  ")
        );
    } else {
        println!("Nothing to do!");
    }

    Ok(())
}

#[instrument]
fn ready_command(ready_args: ReadyCommand) -> Result<bool> {
    let dirs = VmexecDirs::new()?;

    let lock_dir = &dirs.lock_dir;
    trace!("Trying to lock {lock_dir:?}");
    let lock = DirLock::new_sync(&lock_dir)?;

    let ssh_keypair = ensure_ssh_key(&dirs.secrets_dir)?;

    let mut vms = get_vms(&dirs.vms_dir)?;

    // If the VM has already exited, we should exit early.
    if vms[&ready_args.vmid].state() == VmState::Exited {
        println!("exited");
        return Ok(false);
    }

    if ready_args.wait {
        let now = Instant::now();

        while now.elapsed() < ready_args.timeout {
            vms = get_vms(&dirs.vms_dir)?;
            if vms[&ready_args.vmid].state() == VmState::Running {
                break;
            }
            sleep(Duration::from_secs(1));
        }
    }

    let cid = vms[&ready_args.vmid]
        .cid_pids
        .cid
        .wrap_err("cid didn't exist")?;

    let ssh_launch_opts = ssh::SshLaunchOpts {
        timeout: if ready_args.wait {
            ready_args.timeout
        } else {
            Duration::from_secs(0)
        },
        tty: false,
        interactive: Interactive::Never,
        env_vars: vec![],
        workdir: None,
        args: vec!["echo".to_string(), "ok".to_string()],
        privkey: ssh_keypair.privkey_str,
        cid,
    };

    let rt = Runtime::new()?;
    let is_ready = rt.block_on(create_ssh_connection(&ssh_launch_opts)).is_ok();

    if is_ready {
        println!("ready");
    } else {
        println!("not ready");
    }

    trace!("Unlocking {lock_dir:?}");
    drop(lock);

    Ok(is_ready)
}

// As a general note, the reason this program uses mixed sync and async is because we optionally daemonize.
// The program will misbehave if the tokio runtime is already running at the time of daemonizing.
// See for instance https://github.com/xadaemon/daemonize-me/discussions/16
// This requires that we run the program sync until after daemonizing. This is fine since clap
// argument parsing can't really be async anyway.
fn main() -> Result<()> {
    let cli = cli::Cli::parse();

    install_tracing(cli.log_level);
    color_eyre::install()?;

    match cli.command {
        Command::Ps(ps_args) => ps_command(ps_args)?,
        Command::Inspect(inspect_args) => inspect_command(inspect_args)?,
        Command::Images(images_args) => images_command(images_args)?,
        Command::Logs(logs_args) => logs_command(logs_args)?,
        Command::Stop(stop_args) => stop_command(stop_args)?,
        Command::Exec(exec_args) => exec_command(exec_args)?,
        Command::Pull(pull_args) => pull_command(pull_args)?,
        Command::Run(run_args) => {
            let exit_code = run_command(run_args)?;
            if let Some(child_code) = exit_code {
                process::exit(child_code.try_into()?);
            }
        }
        Command::Ksm(ksm_args) => ksm_command(ksm_args)?,
        Command::Clean(clean_args) => clean_command(clean_args)?,
        Command::Prune(prune_args) => prune_command(prune_args)?,
        Command::Ready(ready_args) => {
            let is_ready = ready_command(ready_args)?;
            process::exit(if is_ready { 0 } else { 1 });
        }
        Command::Completions { shell } => {
            let mut clap_app = cli::Cli::command();
            let app_name = clap_app.get_name().to_string();
            clap_complete::generate(shell, &mut clap_app, app_name, &mut std::io::stdout());
        }
        Command::Manpage { out_dir } => {
            let clap_app = cli::Cli::command();
            clap_mangen::generate_to(clap_app, out_dir)?;
        }
    }
    Ok(())
}

fn post_fork_parent(_ppid: i32, _cpid: i32) -> ! {
    process::exit(0);
}
