// SPDX-License-Identifier: Apache-2.0

pub mod cli_handler;
pub mod deploy;
pub mod grpc;
mod injector;
pub mod update;

use crate::config::{self, Config};
use crate::snp;
use crate::utils::{block_on, display_mib_or_gib, shorten_string};
use grpc::proto;
use lazy_static::lazy_static;
use serde::{Deserialize, Serialize};
use tabled::Tabled;

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error(transparent)]
    Config(#[from] config::Error),
    #[error("Could not find a contract with the ID {0}")]
    VmContractNotFound(String),
    #[error("Did not find a SNP node that matches your criteria")]
    NoValidNodeFound,
    #[error("Could not read yaml from disk: {0}")]
    YamlNotFound(#[from] std::io::Error),
    #[error("Failed to parse yaml update file: {0}")]
    YamlFormat(#[from] serde_yaml::Error),
    #[error("Brain returned the following error: {0}")]
    Brain(#[from] grpc::Error),
    #[error("Node returned the following error: {0}")]
    Node(String),
    #[error("Failed to obtain boot measurement for VM")]
    NoMeasurement,
    #[error("Failed to inject secrets: {0}")]
    Injector(#[from] injector::Error),
}

//  TODO: push this out of snp module
#[derive(Serialize, Deserialize, Default)]
pub struct Location {
    pub node_ip: Option<String>,
    pub country: Option<String>,
    pub region: Option<String>,
    pub city: Option<String>,
}

impl From<&str> for Location {
    fn from(s: &str) -> Self {
        match s {
            "Canada" => Self { country: Some("CA".to_string()), ..Default::default() },
            "Montreal" => Self { city: Some("Montréal".to_string()), ..Default::default() },
            "Vancouver" => Self { city: Some("Vancouver".to_string()), ..Default::default() },
            "US" => Self { country: Some("US".to_string()), ..Default::default() },
            "California" => Self { country: Some("US".to_string()), ..Default::default() },
            "France" => Self { country: Some("FR".to_string()), ..Default::default() },
            "GB" => Self { country: Some("GB".to_string()), ..Default::default() },
            "DE" => Self { country: Some("DE".to_string()), ..Default::default() },
            "Any" => Self { ..Default::default() },
            _ => Self { ..Default::default() },
        }
    }
}

#[derive(Serialize, Default)]
pub struct VmSshArgs {
    uuid: String,
    hostname: String,
    ip: String,
    port: String,
    user: String,
    key_path: String,
}

impl crate::HumanOutput for VmSshArgs {
    fn human_cli_print(&self) {
        println!("To SSH into {} (UUID: {}), run the following command:", self.hostname, self.uuid);
        println!("    ssh -i {} -p {} {}@{}", self.key_path, self.port, self.user, self.ip);
    }
}

impl TryFrom<grpc::proto::VmContract> for VmSshArgs {
    type Error = Error;

    fn try_from(contract: grpc::proto::VmContract) -> Result<Self, Self::Error> {
        let mut args = VmSshArgs { ..Default::default() };
        args.uuid = contract.uuid;
        args.hostname = contract.hostname;
        args.user = "root".to_string();
        args.key_path =
            Config::init_config().get_ssh_pubkey()?.trim_end_matches(".pub").to_string();
        if !contract.vm_public_ipv4.is_empty() {
            args.ip = contract.vm_public_ipv4;
            args.port = "22".to_string();
        } else {
            args.port = contract.mapped_ports[0].host_port.to_string();
            log::debug!(
                "This VM does not have a public IP. Using node public IP: {}",
                contract.node_ip
            );
            args.ip = contract.node_ip;
        }
        Ok(args)
    }
}

#[derive(Clone, Serialize, Deserialize, Debug, PartialEq)]
pub struct Dtrfs {
    #[serde(default)]
    name: String,
    #[serde(default)]
    vendor: String,
    dtrfs_url: String,
    dtrfs_sha: String,
    kernel_url: String,
    kernel_sha: String,
}

impl Dtrfs {
    pub fn print_dtrfs_list() -> Vec<Self> {
        vec![DEFAULT_DTRFS.clone(), ALTERNATIVE_INIT[0].clone(), ALTERNATIVE_INIT[1].clone()]
    }

    fn load_from_file(config_path: &str) -> Result<Self, Error> {
        let content = std::fs::read_to_string(config_path)?;
        Ok(serde_yaml::from_str(&content)?)
    }
}

impl super::HumanOutput for Vec<Dtrfs> {
    fn human_cli_print(&self) {
        print!("{}", serde_yaml::to_string(&self).unwrap())
    }
}

#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct Distro {
    name: String,
    vendor: String,
    template_url: String,
    template_sha: String,
}

impl super::HumanOutput for Vec<Distro> {
    fn human_cli_print(&self) {
        print!("{}", serde_yaml::to_string(&self).unwrap())
    }
}

impl Distro {
    pub fn get_template_list() -> Vec<Self> {
        vec![
            DEFAULT_ARCHLINUX.clone(),
            DEFAULT_UBUNTU.clone(),
            DEFAULT_FEDORA.clone(),
            ALTERNATIVE_DISTROS[0].clone(),
            ALTERNATIVE_DISTROS[1].clone(),
            ALTERNATIVE_DISTROS[2].clone(),
        ]
    }

    pub fn from_string(distro: &str) -> Distro {
        match distro {
            "arch" => DEFAULT_ARCHLINUX.clone(),
            "ubuntu" => DEFAULT_UBUNTU.clone(),
            "fedora" => DEFAULT_FEDORA.clone(),
            _ => {
                log::error!("Only arch, fedora and ubuntu are currently available");
                panic!();
            }
        }
    }
}

#[derive(Tabled, Debug, Serialize, Deserialize)]
pub struct VmContract {
    #[tabled(rename = "Location")]
    pub location: String,
    #[tabled(rename = "UUID pfx", display_with = "shorten_string")]
    pub uuid: String,
    pub hostname: String,
    #[tabled(rename = "Cores")]
    pub vcpus: u64,
    #[tabled(rename = "Mem", display_with = "display_mib_or_gib")]
    pub mem: u64,
    #[tabled(rename = "Disk", display_with = "display_mib_or_gib")]
    pub disk: u64,
    #[tabled(rename = "credits/h")]
    pub cost_h: f64,
    #[tabled(rename = "time left", display_with = "display_mins")]
    pub time_left: u64,
}

impl crate::HumanOutput for Vec<VmContract> {
    fn human_cli_print(&self) {
        let style = tabled::settings::Style::rounded();
        let mut table = tabled::Table::new(self);
        table.with(style);
        println!("{table}");
    }
}

fn display_mins(minutes: &u64) -> String {
    let mins = minutes % 60;
    let hours = minutes / 60;

    format!("{hours}h {mins}m")
}

impl From<proto::VmContract> for VmContract {
    fn from(brain_contract: proto::VmContract) -> Self {
        Self {
            uuid: brain_contract.uuid,
            hostname: brain_contract.hostname,
            vcpus: brain_contract.vcpus as u64,
            mem: brain_contract.memory_mb as u64,
            disk: brain_contract.disk_size_gb as u64,
            location: brain_contract.location,
            cost_h: (brain_contract.nano_per_minute * 60) as f64 / 1_000_000_000.0,
            time_left: brain_contract.locked_nano / brain_contract.nano_per_minute,
        }
    }
}

#[derive(Tabled, Debug, Serialize, Deserialize)]
pub struct TabledVmNode {
    #[tabled(rename = "Operator", display_with = "shorten_string")]
    pub operator: String,
    #[tabled(rename = "Main IP")]
    pub main_ip: String,
    #[tabled(rename = "City, Region, Country")]
    pub location: String,
    #[tabled(rename = "Cores")]
    pub vcpus: u64,
    #[tabled(rename = "Mem", display_with = "display_mib_or_gib")]
    pub memory_mib: u64,
    #[tabled(rename = "Disk", display_with = "display_mib_or_gib")]
    pub disk_mib: u64,
    #[tabled(rename = "Extra IPv4", display_with = "display_ip_support")]
    pub public_ipv4: bool,
    #[tabled(rename = "IPv6", display_with = "display_ip_support")]
    pub public_ipv6: bool,
    #[tabled(rename = "Price per unit")]
    pub price: String,
    #[tabled(rename = "Reports")]
    pub reports: usize,
}

impl From<proto::VmNodeListResp> for TabledVmNode {
    fn from(brain_node: proto::VmNodeListResp) -> Self {
        let (mut vcpus, mut memory_mib, mut disk_mib): (u64, u64, u64) = (0, 0, 0);
        let mut price = 0;
        for offer in brain_node.offers.iter() {
            vcpus = vcpus.saturating_add(offer.vcpus);
            memory_mib = memory_mib.saturating_add(offer.memory_mib);
            disk_mib = disk_mib.saturating_add(offer.disk_mib);
            if offer.price > price {
                price = offer.price;
            }
        }
        Self {
            operator: brain_node.operator,
            location: brain_node.city + ", " + &brain_node.region + ", " + &brain_node.country,
            main_ip: brain_node.ip,
            price: format!("{} nano/min", price),
            reports: brain_node.reports.len(),
            public_ipv4: brain_node.public_ipv4,
            public_ipv6: brain_node.public_ipv6,
            vcpus,
            memory_mib,
            disk_mib,
        }
    }
}

pub fn ssh(uuid: &str, just_print: bool) -> Result<VmSshArgs, Error> {
    log::info!("Getting VM information about {uuid} from brain...");
    let req = proto::ListVmContractsReq {
        wallet: Config::get_detee_wallet()?,
        uuid: uuid.to_string(),
        ..Default::default()
    };
    let contracts = block_on(snp::grpc::list_contracts(req))?;
    if contracts.is_empty() {
        return Err(Error::VmContractNotFound(uuid.to_string()));
    }
    let args: VmSshArgs = contracts[0].clone().try_into()?;
    if just_print {
        return Ok(args);
    }
    eprintln!(
        "Running SSH command: ssh -i {} -p {} {}@{}",
        args.key_path, args.port, args.user, args.ip
    );
    use std::os::unix::process::CommandExt;
    let e = std::process::Command::new("ssh")
        .arg("-i")
        .arg(args.key_path)
        .arg("-p")
        .arg(args.port)
        .arg(format!("{}@{}", args.user, args.ip))
        .stdin(std::process::Stdio::inherit())
        .stdout(std::process::Stdio::inherit())
        .stderr(std::process::Stdio::inherit())
        .exec();

    println!("Error: Failed to execute ssh process: {e}");
    std::process::exit(2);
}

pub fn get_one_contract(uuid: &str) -> Result<proto::VmContract, Error> {
    log::debug!("Getting contract {uuid} from brain...");
    let req = proto::ListVmContractsReq {
        wallet: Config::get_detee_wallet()?,
        uuid: uuid.to_string(),
        ..Default::default()
    };
    let contracts = block_on(snp::grpc::list_contracts(req))?;
    if contracts.is_empty() {
        return Err(Error::VmContractNotFound(uuid.to_string()));
    }
    Ok(contracts[0].clone())
}

pub fn get_node_by_contract(uuid: &str) -> Result<proto::VmNodeListResp, Error> {
    let contract = get_one_contract(uuid)?;
    Ok(block_on(snp::grpc::get_one_node(proto::VmNodeFilters {
        node_pubkey: contract.node_pubkey,
        ..Default::default()
    }))?)
}

pub fn delete_contract(uuid: &str) -> Result<(), Error> {
    Ok(block_on(snp::grpc::delete_vm(uuid))?)
}

pub fn list_contracts(as_operator: bool) -> Result<Vec<VmContract>, Error> {
    let req = proto::ListVmContractsReq {
        wallet: Config::get_detee_wallet()?,
        as_operator,
        ..Default::default()
    };
    let contracts: Vec<VmContract> =
        block_on(grpc::list_contracts(req))?.into_iter().map(|n| n.into()).collect();
    let _ = write_uuid_list(&contracts);
    Ok(contracts)
}

fn write_uuid_list(contracts: &[VmContract]) -> Result<(), Error> {
    let vm_uuid_list_path = Config::vm_uuid_list_path()?;
    let mut file = std::fs::File::create(vm_uuid_list_path)?;
    let output: String = contracts
        .iter()
        .map(|vm| format!("{}\t{}", vm.uuid, vm.hostname).to_string())
        .collect::<Vec<String>>()
        .join("\n");
    let output = output + "\n";
    std::io::Write::write_all(&mut file, output.as_bytes())?;
    Ok(())
}

pub fn append_uuid_list(uuid: &str, hostname: &str) -> Result<(), Error> {
    use std::fs::OpenOptions;
    use std::io::prelude::*;
    let mut file =
        OpenOptions::new().create(true).append(true).open(Config::vm_uuid_list_path()?)?;
    writeln!(file, "{uuid}\t{hostname}")?;
    Ok(())
}

impl super::HumanOutput for Vec<proto::VmNodeListResp> {
    fn human_cli_print(&self) {
        let nodes: Vec<TabledVmNode> = self.iter().map(|n| n.clone().into()).collect();
        let style = tabled::settings::Style::rounded();
        let mut table = tabled::Table::new(nodes);
        table.with(style);
        println!("{table}");
    }
}

pub fn search_nodes(location: Location) -> Result<Vec<proto::VmNodeListResp>, Error> {
    log::debug!("This will support flags in the future, but we have only one node atm.");
    let req = proto::VmNodeFilters {
        city: location.city.unwrap_or_default(),
        country: location.country.unwrap_or_default(),
        region: location.region.unwrap_or_default(),
        ..Default::default()
    };
    Ok(block_on(grpc::get_node_list(req))?)
}

#[derive(Tabled, Debug, Serialize, Deserialize)]
pub struct NodeOffer {
    #[tabled(rename = "Location")]
    pub location: String,
    #[tabled(rename = "Cores")]
    pub vcpus: u64,
    #[tabled(rename = "Mem", display_with = "display_mib_or_gib")]
    pub mem: u64,
    #[tabled(rename = "Disk", display_with = "display_mib_or_gib")]
    pub disk: u64,
    #[tabled(rename = "Public IPv4", display_with = "display_ip_support")]
    pub ipv4: bool,
    // #[tabled(rename = "Public IPv6", display_with = "display_ip_support")]
    // pub ipv6: bool,
    #[tabled(rename = "cost/h")]
    pub cost_h: f64,
    #[tabled(rename = "cost/m")]
    pub cost_m: f64,
}

fn display_ip_support(support: &bool) -> String {
    match support {
        true => "Available".to_string(),
        false => "Unavailable".to_string(),
    }
}

impl super::HumanOutput for Vec<NodeOffer> {
    fn human_cli_print(&self) {
        let style = tabled::settings::Style::rounded();
        let mut table = tabled::Table::new(self);
        table.with(style);
        println!("{table}");
    }
}

pub fn print_node_offers(location: Location) -> Result<Vec<NodeOffer>, Error> {
    log::debug!("This will support flags in the future, but we have only one node atm.");
    let req = proto::VmNodeFilters {
        city: location.city.unwrap_or_default(),
        country: location.country.unwrap_or_default(),
        region: location.region.unwrap_or_default(),
        ..Default::default()
    };
    let node_list = block_on(grpc::get_node_list(req))?;
    let mut offers: Vec<NodeOffer> = Vec::new();
    for node in node_list.iter() {
        for offer in node.offers.iter() {
            if offer.vcpus == 0 || offer.memory_mib == 0 || offer.disk_mib == 0 {
                continue;
            }
            let mem_per_cpu = offer.memory_mib / offer.vcpus;
            let disk_per_cpu = offer.disk_mib / offer.vcpus;
            for i in 1..offer.vcpus + 1 {
                let price_per_month = calculate_nanocredits(
                    i as u32,
                    (mem_per_cpu * i) as u32,
                    (disk_per_cpu * i) as u32,
                    false,
                    732,
                    offer.price,
                ) as f64
                    / 1_000_000_000_f64;
                let price_per_hour = price_per_month / 732_f64;
                let price_per_month = (price_per_month * 100.0).round() / 100.0;
                let price_per_hour = (price_per_hour * 1000.0).round() / 1000.0;
                offers.push(NodeOffer {
                    location: node.city.clone() + ", " + &node.region + ", " + &node.country,
                    vcpus: i,
                    mem: i * mem_per_cpu,
                    disk: i * disk_per_cpu,
                    cost_h: price_per_hour,
                    cost_m: price_per_month,
                    ipv4: node.public_ipv4,
                    // ipv6: node.public_ipv6,
                });
            }
        }
    }
    offers.sort_by_key(|n| n.cost_m as u64);
    Ok(offers)
}

pub fn inspect_node(ip: String) -> Result<proto::VmNodeListResp, Error> {
    let req = proto::VmNodeFilters { ip, ..Default::default() };
    Ok(block_on(grpc::get_one_node(req))?)
}

pub fn calculate_nanocredits(
    vcpus: u32,
    memory_mb: u32,
    disk_size_mib: u32,
    public_ipv4: bool,
    hours: u32,
    node_price: u64,
) -> u64 {
    // this calculation needs to match the calculation of the network
    let total_units = (vcpus as u64 * 10)
        + ((memory_mb + 256) as u64 / 200)
        + (disk_size_mib as u64 / 1024 / 10)
        + (public_ipv4 as u64 * 10);
    let locked_nano = hours as u64 * 60 * total_units * node_price;
    locked_nano
}

lazy_static! {
    static ref DEFAULT_DTRFS: Dtrfs = Dtrfs {
        name: "dtrfs-6.15.4-arch2-1".to_string(),
        vendor: "ghe0".to_string(),
        dtrfs_url: "http://registry.detee.ltd/detee-archtop-6.15.4-arch2-1.cpio.gz".to_string(),
        dtrfs_sha: "dfde2c360341d9c7622c0f0e5200bb8ed9343cb9302ea5a06523d41705b0e4f9".to_string(),
        kernel_url: "http://registry.detee.ltd/vmlinuz-linux-6.15.4-arch2-1".to_string(),
        kernel_sha: "01581fba284c237131ee8d6662e3fde4ebbd55c496fcae2979448360ac3f05b0".to_string()
    };
    static ref DEFAULT_ARCHLINUX: Distro = Distro {
        name: "archlinux_2025-07-04".to_string(),
        vendor: "gheorghe".to_string(),
        template_url: "http://registry.detee.ltd/detee_arch_2025-07-04.fsa".to_string(),
        template_sha: "af86b01e71b75328b2df4d7f0fda36f69b4ae68d20ed1ce3351da1f77a4eb260"
            .to_string()
    };
    static ref DEFAULT_UBUNTU: Distro = Distro {
        name: "ubuntu_2025-07-04".to_string(),
        vendor: "gheorghe".to_string(),
        template_url: "http://registry.detee.ltd/detee_ubuntu_2025-07-04.fsa".to_string(),
        template_sha: "291aa82bfee3fd997724cfe8f2b2454c2e73b37120d5008bbdc60a669a13a591"
            .to_string()
    };
    static ref DEFAULT_FEDORA: Distro = Distro {
        name: "fedora_2025-07-04".to_string(),
        vendor: "gheorghe".to_string(),
        template_url: "http://registry.detee.ltd/detee_fedora_2025-07-04.fsa".to_string(),
        template_sha: "43adbf96ab43efd07179bf8c2d2f529870d89a20d173a68915095df4fb632ccf"
            .to_string()
    };
    static ref ALTERNATIVE_INIT: Vec<Dtrfs> = vec![
        Dtrfs {
            name: "dtrfs-6.14.2-arch1-1".to_string(),
            vendor: "ghe0".to_string(),
            dtrfs_url: "http://registry.detee.ltd/detee-archtop-6.14.2-arch1-1.cpio.gz".to_string(),
            dtrfs_sha: "d207644ee60d54009b6ecdfb720e2ec251cde31774dd249fcc7435aca0377990"
                .to_string(),
            kernel_url: "http://registry.detee.ltd/vmlinuz-linux-6.14.2-arch1-1".to_string(),
            kernel_sha: "e765e56166ef321b53399b9638584d1279821dbe3d46191c1f66bbaa075e7919"
                .to_string()
        },
        Dtrfs {
            name: "dtrfs-6.13.7-arch1-1".to_string(),
            vendor: "ghe0".to_string(),
            dtrfs_url: "http://registry.detee.ltd/detee-archtop-6.13.7-arch1-1.cpio.gz".to_string(),
            dtrfs_sha: "dc02e091da80c281fe735a1be86b3fe766f1741d82c32f5dc64344b345827c6d"
                .to_string(),
            kernel_url: "http://registry.detee.ltd/vmlinuz-linux-6.13.7-arch1-1".to_string(),
            kernel_sha: "469a89668d2f5744b3f80417fcf0a4ce0140fcb78f1e8834ef8e3668eecc934c"
                .to_string()
        },
    ];
    static ref ALTERNATIVE_DISTROS: Vec<Distro> = vec![
        Distro {
            name: "archlinux_2025-04-03".to_string(),
            vendor: "gheorghe".to_string(),
            template_url: "http://registry.detee.ltd/detee_arch_2025-04-03.fsa".to_string(),
            template_sha: "7fdb19d9325c63d246140c984dc3764538f6ea329ed877e947993ea7bc8c2067"
                .to_string()
        },
        Distro {
            name: "ubuntu_2025-04-03".to_string(),
            vendor: "gheorghe".to_string(),
            template_url: "http://registry.detee.ltd/detee_ubuntu_2025-04-03.fsa".to_string(),
            template_sha: "324895a7a1788e43253cf9699aa446df1a5519fe072917cedcc4ed356546e34a"
                .to_string()
        },
        Distro {
            name: "fedora_2025-02-21".to_string(),
            vendor: "gheorghe".to_string(),
            template_url: "http://registry.detee.ltd/detee_fedora_2025-02-21.fsa".to_string(),
            template_sha: "c0fdd08d465939077ef8ed746903005fc190af12cdf70917cc8c6f872da85777"
                .to_string()
        }
    ];
}