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

use clap::{CommandFactory, Parser, crate_name};
use colored_json::to_colored_json_auto;
use comfy_table::{Row, Table};
use daemonize_me::Daemon;
use eyre::{Context, ContextCompat, Result, bail};
use humansize::{BINARY, FormatSizeOptions, format_size_i};
use termion::{color, style};
use tokio::runtime::Runtime;
use tracing::{Level, debug, instrument, warn};
use tracing_error::ErrorLayer;
use tracing_subscriber::prelude::*;
use tracing_subscriber::{EnvFilter, fmt};

mod cli;

use crate::cli::{
    CleanCommand, Command, ExecCommand, ImagesCommand, InspectCommand, KsmCommand, LogsCommand,
    PruneCommand, PsCommand, PullCommand, ReadyCommand, RunCommand, StopCommand,
};
use vmexec::runner::{self, Readiness};
use vmexec::types::{Io, Pull, Tty};
use vmexec::vms::{
    DAEMON_LOG_FILENAME, SESSION_LOG_FILENAME, VMEXEC_PID_FILENAME, VmState, get_vm, get_vms,
    reap_dead_vm_dirs, reserve_vm, stop_vm,
};
use vmexec::{
    utils::{VmexecDirs, find_required_tools},
    vm_images::{get_images, prune_overlays, pull_image},
};

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

    let mut table = Table::new();
    table.load_style(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}");

    Ok(())
}

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

    let vm = get_vm(&dirs.vms_dir, &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_style(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 vm = get_vm(&dirs.vms_dir, &logs_args.vmid)?;
    let session_log_path = vm.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_vm(&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()?;

    // We're not locking here because otherwise we wouldn't be able to abort a booting VM.

    stop_vm(&dirs.vms_dir, &stop_args.vmid)?;

    Ok(())
}

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

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

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

    let rt = Runtime::new()?;
    let exit_code = rt.block_on(runner::exec_in_running_vm(
        &dirs,
        &exec_args.vmid,
        runner::ExecOpts {
            ssh_timeout: exec_args.ssh_timeout,
            tty,
            interactive: exec_args.interactive,
            env_vars: exec_args.env,
            workdir: exec_args.workdir,
            args: exec_args.args,
            stdout: Io::Inherit,
            stderr: Io::Inherit,
        },
    ))?;

    Ok(exit_code.map(|output| output.exit_code))
}

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

    let rt = Runtime::new()?;
    let vm_image = rt.block_on(pull_image(&dirs, pull_args.os_type, Pull::Newer))?;
    println!("Image available at {}", vm_image.image_path.display());

    if pull_args.prepare {
        rt.block_on(runner::prepare_pulled_image(
            &dirs,
            vm_image,
            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,
            },
        ))?;

        println!("Image prepared");
    }

    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();

    // Reserve the state dir for this VM (usually ~/.local/state/vmexec/vms/<random id>/).
    // This also takes the state dir lock.
    let reservation = reserve_vm(&dirs, run_args.vmid)?;

    let pid_file = reservation.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}", vmid = reservation.vmid);

        // Log daemon stdout and stderr into same file.
        let daemon_log = fs::File::create(reservation.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 tty = match run_args.tty {
        Tty::Always => true,
        Tty::Never => false,
        Tty::Auto => std::io::stdout().is_terminal(),
    };

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

    let rt = Runtime::new()?;
    let exit_code = rt.block_on(runner::prepare_and_run(
        &dirs,
        reservation,
        tool_paths,
        runner::RunOptions {
            image_source: run_args.image_source,
            pull: run_args.pull,
            ssh_timeout: run_args.ssh_timeout,
            memory: run_args.memory.unwrap_or_default().gb,
            disable_kvm: run_args.disable_kvm,
            show_vm_window: run_args.show_vm_window,
            tty,
            interactive: run_args.interactive,
            env_vars: run_args.env,
            workdir: run_args.workdir,
            args: run_args.args,
            volumes: run_args.volumes,
            pmems: run_args.pmems,
            published_ports: run_args.published_ports,
            stdout: Io::Inherit,
            stderr: Io::Inherit,
            rm: run_args.rm,
        },
    ))?;

    Ok(exit_code.map(|output| output.exit_code))
}

fn check_ksm_active() {
    let ksm_run = fs::read_to_string("/sys/kernel/mm/ksm/run");

    // If we can't read this for some reason, we'll just consider it unavailable and move on. It's
    // not critical for operation.
    if let Ok(ksm) = ksm_run {
        let ksm_active = ksm.trim() == "1";

        if !ksm_active {
            warn!("Kernel Samepage Merging (KSM) is disabled.");
            warn!("It is strongly recommended to enable it.");
            warn!("You can run `vmexec ksm --enable`");
        }
    }
}

fn print_ksm_stats() -> Result<()> {
    let pages_scanned = fs::read_to_string("/sys/kernel/mm/ksm/pages_scanned")?;
    let pages_sharing = fs::read_to_string("/sys/kernel/mm/ksm/pages_sharing")?;
    let full_scans = fs::read_to_string("/sys/kernel/mm/ksm/full_scans")?;
    let general_profit: i64 = fs::read_to_string("/sys/kernel/mm/ksm/general_profit")?
        .trim()
        .parse()?;
    let format_size_options = FormatSizeOptions::from(BINARY).decimal_places(1);
    let general_profit_human = format_size_i(general_profit, format_size_options);

    println!(
        "Pages scanned: {}{}{:>10}{}{}",
        style::Bold,
        color::Fg(color::Blue),
        pages_scanned.trim(),
        color::Fg(color::Reset),
        style::Reset,
    );
    println!(
        "Pages sharing: {}{}{:>10}{}{}",
        style::Bold,
        color::Fg(color::Blue),
        pages_sharing.trim(),
        color::Fg(color::Reset),
        style::Reset,
    );
    println!(
        "Full scans: {}{}{:>13}{}{}",
        style::Bold,
        color::Fg(color::Blue),
        full_scans.trim(),
        color::Fg(color::Reset),
        style::Reset,
    );
    println!(
        "{}General profit: {:>9}{}{}{}",
        style::Bold,
        general_profit_human,
        color::Fg(color::LightBlue),
        color::Fg(color::Reset),
        style::Reset,
    );

    Ok(())
}

#[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)?;

    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 deleted_files = prune_overlays(&dirs)?;

    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 timeout = if ready_args.wait {
        Some(ready_args.timeout)
    } else {
        None
    };

    let rt = Runtime::new()?;
    let readiness = rt.block_on(runner::check_ready(&dirs, &ready_args.vmid, timeout))?;

    let is_ready = match readiness {
        Readiness::Ready => {
            println!("ready");
            true
        }
        Readiness::NotReady => {
            println!("not ready");
            false
        }
        Readiness::Exited => {
            println!("exited");
            false
        }
    };

    Ok(is_ready)
}

fn install_tracing(log_level: Level) {
    let format = fmt::format::debug_fn(|writer, field, value| {
        if field.name() == "message" {
            write!(writer, "{value:?}")
        } else {
            // We'll format the field name and value separated with a colon.
            write!(writer, "")
        }
    })
    // Separate each field with a comma.
    // This method is provided by an extension trait in the
    // `tracing-subscriber` prelude.
    .delimited("");

    let filter_layer = EnvFilter::try_new(format!("{}={}", crate_name!(), log_level)).unwrap();

    let subscriber = tracing_subscriber::registry()
        .with(filter_layer)
        .with(ErrorLayer::default());

    // Log to stderr so that stdout isn't polluted by tracing output.
    // Don't show fields if debug level is INFO or less to cut back on the output a bit.
    if log_level <= Level::INFO {
        let fmt_layer = fmt::layer()
            .with_target(false)
            .compact()
            .fmt_fields(format)
            .with_writer(std::io::stderr);
        subscriber.with(fmt_layer).init();
    } else {
        let fmt_layer = fmt::layer()
            .with_target(false)
            .compact()
            .with_writer(std::io::stderr);
        subscriber.with(fmt_layer).init();
    };
}

// 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) => {
            let exit_code = exec_command(exec_args)?;
            if let Some(child_code) = exit_code {
                process::exit(child_code.try_into()?);
            }
        }
        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);
}
