/* src/disks/mod.rs
 *
 * Copyright 2025 Mission Center Developers
 *
 * This program 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.
 *
 * This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * SPDX-License-Identifier: GPL-3.0-or-later
 */

use std::collections::{HashMap, HashSet};

use phf::phf_set;

use magpie_platform::disks::Disk;

use crate::{async_runtime, sync, system_bus};

use crate::disks::disk_wrapper::DiskWrapper;
pub use manager::DisksManager;
use stats::Stats;

mod disk_wrapper;
mod manager;
mod smart_data;
mod stats;
mod util;

static IGNORED_DISK_PREFIXES: phf::Set<&'static str> = phf_set! {
    "loop",
    "ram",
    "zram",
    "fd",
    "md",
    "dm",
    "zd",
};

pub struct DisksCache {
    disks: Vec<Disk>,
    stats: Vec<Stats>,

    ignored: HashSet<String>,

    udisks2: Option<udisks2::Client>,

    disk_wrappers: HashMap<String, DiskWrapper>,
}

impl magpie_platform::disks::DisksCache for DisksCache {
    fn new() -> Self
    where
        Self: Sized,
    {
        let bus = match system_bus() {
            Some(bus) => bus.clone(),
            None => {
                log::warn!("Failed to connect to system bus");
                return Self {
                    disks: Vec::new(),
                    stats: Vec::new(),
                    ignored: HashSet::new(),
                    udisks2: None,
                    disk_wrappers: HashMap::new(),
                };
            }
        };

        let rt = async_runtime();
        let udisks2 = match sync!(rt, udisks2::Client::new_for_connection(bus)) {
            Ok(udisks2) => Some(udisks2),
            Err(e) => {
                log::warn!("Failed to connect to udisks2: {}", e);
                None
            }
        };

        Self {
            disks: Vec::new(),
            stats: Vec::new(),
            ignored: HashSet::new(),
            udisks2,
            disk_wrappers: HashMap::new(),
        }
    }

    fn refresh(&mut self) {
        let udisks2 = self.udisks2.as_ref();
        let rt = async_runtime();

        let mut prev_disks = std::mem::take(&mut self.disks);
        let mut prev_stats = std::mem::take(&mut self.stats);

        let dir = match std::fs::read_dir("/sys/block") {
            Ok(dir) => dir,
            Err(e) => {
                log::warn!("Failed to read `/sys/block`: {e}");
                return;
            }
        };

        'outer: for entry in dir.filter_map(Result::ok) {
            let file_name = entry.file_name();
            let disk_id = file_name.to_string_lossy();
            let disk_id = disk_id.as_ref();

            if self.ignored.contains(disk_id) {
                continue;
            }

            for i in 2..=disk_id.len().min(4) {
                if IGNORED_DISK_PREFIXES.contains(&disk_id[..i]) {
                    self.ignored.insert(disk_id.to_string());
                    continue 'outer;
                }
            }

            let mut prev_disk_index = None;
            for (i, disk) in prev_disks.iter().enumerate() {
                if disk.id == disk_id {
                    prev_disk_index = Some(i);
                    break;
                }
            }

            if !self.disk_wrappers.contains_key(disk_id) {
                if let Some((udisks2, object)) = udisks2
                    .and_then(|client| util::object(client, disk_id).map(|obj| (client, obj)))
                {
                    self.disk_wrappers.insert(
                        disk_id.to_string(),
                        sync!(rt, DiskWrapper::new(disk_id, udisks2, object)),
                    );
                }
            }

            let Some(disk_wrapper) = self.disk_wrappers.get_mut(disk_id) else {
                continue;
            };

            let (mut disk, stats) = match prev_disk_index {
                Some(i) => (prev_disks.swap_remove(i), prev_stats.swap_remove(i)),
                None => (
                    sync!(rt, disk_wrapper.create_disk_obj()),
                    disk_wrapper.create_stats(),
                ),
            };

            let mut new_stats = disk_wrapper.create_stats();
            new_stats.update(&stats, &mut disk);

            sync!(rt, disk_wrapper.update_disk_obj(&mut disk, udisks2));

            self.disks.push(disk);
            self.stats.push(new_stats);
        }
    }

    fn cached_entries(&self) -> &[Disk] {
        &self.disks
    }
}

#[cfg(test)]
mod tests {
    use magpie_platform::disks::DisksCache;

    #[test]
    fn test_disks_cache() {
        let mut cache = super::DisksCache::new();
        cache.refresh();
    }
}
