// This file is part of Deja-vu.
//
// Deja-vu is free software: you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// Deja-vu is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
// See the GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with Deja-vu. If not, see <https://www.gnu.org/licenses/>.
//
use std::{
    collections::HashMap,
    fs::File,
    io::Write,
    process::{Command, Output, Stdio},
};

use once_cell::sync::Lazy;
use regex::Regex;
use tempfile::tempdir;
use tracing::{info, warn};

use crate::system::disk::{
    get_disk_serial_number, string_to_bytes, Disk, Partition, PartitionTableType,
};

static DEFAULT_MAX_DISKS: usize = 8;

pub fn get_disk_details(disk_id: usize, disk_size: u64, disk_details: Option<&str>) -> Disk {
    static RE_DETAILS: Lazy<Regex> = Lazy::new(|| {
        Regex::new(r"(.*?)\r?\nDisk ID\s*:\s+(.*?)\r?\nType\s*:\s+(.*?)\r?\n").unwrap()
    });
    static RE_UUID: Lazy<Regex> = Lazy::new(|| {
        Regex::new(r"^\{[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}\}$").unwrap()
    });
    let mut disk = Disk {
        id: disk_id,
        size: disk_size,
        ..Default::default()
    };

    // Get details
    let details: String;
    if let Some(details_str) = disk_details {
        details = String::from(details_str);
    } else {
        let script = format!("select disk {disk_id}\r\ndetail disk");
        details = run_script(&script);
    };

    // Parse details
    for (_, [model, part_type, conn_type]) in
        RE_DETAILS.captures_iter(&details).map(|c| c.extract())
    {
        disk.model = String::from(model);
        disk.conn_type = String::from(conn_type);
        if RE_UUID.is_match(part_type) {
            disk.part_type = PartitionTableType::Guid;
        } else {
            disk.part_type = PartitionTableType::Legacy;
        }
        disk.serial = get_disk_serial_number(disk_id);
    }

    // Done
    disk
}

pub fn get_partition_details(
    disk_id: usize,
    disk_details: Option<&str>,
    part_details: Option<&str>,
) -> Vec<Partition> {
    static RE_LIS: Lazy<Regex> =
        Lazy::new(|| Regex::new(r"Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+B)").unwrap());
    let mut parts = Vec::new();

    // List partition
    let contents: String;
    if let Some(details) = disk_details {
        contents = String::from(details);
    } else {
        let script = format!("select disk {disk_id}\r\nlist partition");
        contents = run_script(&script);
    };
    for (_, [number, size]) in RE_LIS.captures_iter(&contents).map(|c| c.extract()) {
        let part_num = number.parse().unwrap();
        if part_num != 0 {
            // part_num == 0 is reserved for extended partition "containers" so we can exclude them
            let part = Partition {
                id: part_num,
                size: string_to_bytes(size),
                ..Default::default()
            };
            parts.push(part);
        }
    }

    // Detail parititon
    let mut script = vec![format!("select disk {}", disk_id)];
    for part in &parts {
        if part_details.is_some() {
            // Currently only used by tests
            break;
        }
        script.push(format!(
            // Remove/Assign included to ensure all (accessible) volumes have letters
            "\r\nselect partition {}\r\nremove noerr\r\nassign noerr\r\ndetail partition",
            part.id
        ));
    }
    let part_contents: String;
    if let Some(details) = part_details {
        part_contents = String::from(details);
    } else {
        part_contents = run_script(script.join("\r\n").as_str());
    };
    parse_partition_details(&mut parts, &part_contents);

    // Done
    parts
}

#[must_use]
pub fn build_dest_format_script(disk_id: usize, part_type: &PartitionTableType) -> String {
    let disk_id = format!("{disk_id}");
    let mut script = vec!["select disk {disk_id}", "clean"];
    match part_type {
        PartitionTableType::Guid => {
            script.push("convert gpt");
            script.push("create partition efi size=260");
            script.push("format fs=fat32 quick label=ESP");
            script.push("create partition msr size=16");
        }
        PartitionTableType::Legacy => {
            script.push("create partition primary size=100");
            script.push("active");
            script.push("format fs=ntfs quick label=System");
        }
    }
    script.join("\r\n").replace("{disk_id}", &disk_id)
}

#[must_use]
pub fn build_get_disk_script(disk_nums: Option<Vec<&str>>) -> String {
    let capacity = DEFAULT_MAX_DISKS * 3 + 1;
    let script: String;

    // Get disk and partition details
    if let Some(disks) = disk_nums {
        // (Slower case)
        let mut script_parts = Vec::with_capacity(capacity);

        // Get list of disks
        script_parts.push(String::from("list disk"));

        // Add disks from provided list
        for num in disks {
            script_parts.push(format!("select disk {num}"));
            script_parts.push(String::from("detail disk"));
            script_parts.push(String::from("list partition"));
        }

        // Done
        script = script_parts.join("\n");
    } else {
        // (Best case)
        let mut script_parts = Vec::with_capacity(capacity);

        // Get list of disks
        script_parts.push("list disk");

        // Assuming first disk number is zero
        script_parts.push("select disk 0");
        script_parts.push("detail disk");
        script_parts.push("list partition");

        // Limit to 8 disks (if there's more the manual "worst" case will be used)
        let mut i = 0;
        while i < 8 {
            script_parts.push("select disk next");
            script_parts.push("detail disk");
            script_parts.push("list partition");
            i += 1;
        }

        // Done
        script = script_parts.join("\n");
    }

    script
}

pub fn get_disks() -> Vec<Disk> {
    static RE_DIS_DET: Lazy<Regex> = Lazy::new(|| {
        Regex::new(r"(?s)Disk (\d+) is now the selected disk.*?\r?\n\s*\r?\n(.*)").unwrap()
    });
    static RE_DIS_LIS: Lazy<Regex> =
        Lazy::new(|| Regex::new(r"Disk\s+(\d+)\s+(\w+)\s+(\d+\s+\w+B)").unwrap());
    let mut contents: String;
    let mut output;
    let mut script: String;

    // Run diskpart and check result
    script = build_get_disk_script(None);
    output = run_script_raw(&script);
    contents = String::from_utf8_lossy(&output.stdout).to_string();
    if let Some(return_code) = output.status.code() {
        let disk_nums = parse_disk_numbers(&contents);
        if return_code != 0 && !disk_nums.is_empty() {
            // The base assumptions were correct! skipping fallback method
            //
            // Since the return_code was not zero, and at least one disk was detected, that
            // means that the number of disks was less than DEFAULT_MAX_DISKS
            //
            // If the number of disks is exactly DEFAULT_MAX_DISKS then we may be omitting some
            // disk(s) from the list so a manual check is needed
            // This only adds one additional diskpart call, but avoiding it is preferred
            //
            // If zero disks were detected then either the first disk number wasn't zero or no
            // disks were actually connected
            // (this could be due to missing drivers but we're leaving that issue to the user)
            info!("Disk assumptions correct!");
        } else {
            // The base assumptions were wrong so we need to check which disk numbers are present
            warn!("Disk assumptions incorrect, falling back to manual method.");
            script = build_get_disk_script(Some(disk_nums));
            output = run_script_raw(&script);
            contents = String::from_utf8_lossy(&output.stdout).to_string();
        }
    }

    // Split Diskpart output contents into sections
    let mut dp_sections = split_diskpart_disk_output(&contents);

    // Build Disk structs
    // NOTE: A HashMap is used because it's possible to have gaps in the list of disks
    //       i.e. 0, 1, 3, 4
    //       For instance, this can happen if a drive is disconnected after startup
    let mut disks_map: HashMap<&str, Disk> = HashMap::with_capacity(DEFAULT_MAX_DISKS);
    for (_, [number, _status, size]) in RE_DIS_LIS
        .captures_iter(dp_sections.remove(0)) // This is the "list disk" section
        .map(|c| c.extract())
    {
        disks_map.insert(
            number,
            Disk {
                id: number.parse().unwrap(),
                size: string_to_bytes(size),
                ..Default::default()
            },
        );
    }

    // Add Disk details
    let mut disks_raw: Vec<Disk> = Vec::with_capacity(DEFAULT_MAX_DISKS);
    for section in dp_sections {
        for (_, [id, details]) in RE_DIS_DET.captures_iter(section).map(|c| c.extract()) {
            if let Some(disk) = disks_map.remove(id) {
                // We remove the disk from the HashMap because we're moving it to the Vec
                let mut disk = get_disk_details(disk.id, disk.size, Some(details));
                disk.parts = get_partition_details(disk.id, Some(details), None);
                disk.generate_descriptions();
                disks_raw.push(disk);
            }
        }
    }

    // Return list of Disks
    disks_raw
}

pub fn parse_disk_numbers(contents: &str) -> Vec<&str> {
    //Disk 0 is now the selected disk.
    //
    //Red Hat VirtIO SCSI Disk Device
    //Disk ID: {E9CE8DFA-46B2-43C1-99BB-850C661CEE6B}
    static RE: Lazy<Regex> =
        Lazy::new(|| Regex::new(r"\s+Disk\s+(\d+).*\n.*\n.*\nDisk ID:").unwrap());

    let mut disk_nums = Vec::new();
    for (_, [number]) in RE.captures_iter(contents).map(|c| c.extract()) {
        disk_nums.push(number);
    }
    disk_nums
}

pub fn parse_partition_details(parts: &mut [Partition], contents: &str) {
    static RE_PAR: Lazy<Regex> = Lazy::new(|| {
        Regex::new(
            r"Partition (\d+)\r?\nType\s*: (\S+)(\r?\n.*){5}\s*(Volume.*\r?\n.*\r?\n|There is no volume)(.*)",
        )
        .unwrap()
    });
    static RE_VOL: Lazy<Regex> = Lazy::new(|| {
        //   Volume ###  Ltr  Label        Fs     Type        Size     Status     Info
        //   ----------  ---  -----------  -----  ----------  -------  ---------  --------
        // * Volume 1     S   ESP          FAT32  Partition    100 MB  Healthy    Hidden
        Regex::new(r"..Volume (\d.{2})  (.{3})  (.{11})  (.{5})").unwrap()
    });

    for (part_index, (_, [_part_id, part_type, _, _vol_header, vol_line])) in RE_PAR
        .captures_iter(contents)
        .map(|c| c.extract())
        .enumerate()
    {
        if let Some(part) = parts.get_mut(part_index) {
            // Partition info
            part.part_type = String::from(part_type.trim());

            // Volume info
            for (_, [_id, letter, label, fs_type]) in
                RE_VOL.captures_iter(vol_line).map(|c| c.extract())
            {
                part.label = String::from(label.trim());
                part.letter = String::from(letter.trim());
                part.fs_type = String::from(fs_type.trim());
            }
        }
    }
}

pub fn refresh_disk_info(disk: &Disk) -> Disk {
    info!("Refresh disk info");
    let mut disk = get_disk_details(disk.id, disk.size, None);
    disk.parts = get_partition_details(disk.id, None, None);
    disk.generate_descriptions();
    disk
}

/// # Panics
///
/// Will panic if diskpart fails to run the script or if the script if malformed
pub fn run_script_raw(script: &str) -> Output {
    info!("Running Diskpart: {:?}", &script);
    let temp_dir = tempdir().expect("Failed to create temp dir");
    let script_path = temp_dir.path().join("diskpart.script");
    let mut script_file = File::create(&script_path).expect("Failed to create temp file");
    script_file
        .write_all(script.as_bytes())
        .expect("Failed to write script to disk");
    let output = Command::new("diskpart")
        .args(["/s", format!("{}", script_path.display()).as_str()])
        .stdout(Stdio::piped())
        .output()
        .expect("Failed to execute Diskpart script");
    output
}

#[must_use]
pub fn run_script(script: &str) -> String {
    let output = run_script_raw(script);
    String::from_utf8_lossy(&output.stdout).to_string()
}

pub fn split_diskpart_disk_output(contents: &str) -> Vec<&str> {
    // NOTE: A simple split isn't helpful since we want to include the matching lines
    static RE: Lazy<Regex> =
        Lazy::new(|| Regex::new(r"Disk \d+ is now the selected disk").unwrap());
    let mut sections = Vec::new();
    let mut starts: Vec<usize> = vec![0];
    let mut ends: Vec<usize> = Vec::new();
    let _: Vec<_> = RE
        .find_iter(contents)
        .map(|m| {
            ends.push(m.start() - 1);
            starts.push(m.start());
        })
        .collect();
    ends.push(contents.len());
    let ranges: Vec<(&usize, &usize)> = starts.iter().zip(ends.iter()).collect();
    for range in ranges {
        let start = *range.0;
        let end = *range.1;
        sections.push(&contents[start..end]);
    }
    sections
}