/*
 * Copyright (c) 2021, Luca Fulchir <luker@fenrirproject.org>
 *
 * This file is part of dfim.
 *
 * dfim 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.
 *
 * dfim 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 dfim.  If not, see <https://www.gnu.org/licenses/>.
 */

use super::{CheckResult, CheckStatus, NextTarget};
use crate::config::trgt;

// we could use lifetimes here instead of copying stuff, but
// - we should have realatively low usage anyway
// - lifetimes can get messy
// - we don't have enough trgts to fully know how lifetimes will impact the
//   code.
// so for now, clone
pub struct GPT(trgt::gpt::GPT);

impl GPT {
    pub fn new(cfg: &trgt::gpt::GPT) -> Self {
        GPT(cfg.clone())
    }
}

impl GPT {
    /// generate a standard config that can be checked easily
    /// so that we can compare the config GPT with the device GPT
    fn generate_standard(
        &self,
        dev: &::udev::Device,
    ) -> Result<trgt::gpt::GPT, ::std::io::Error> {
        let mut ret = self.0.clone();
        let (min_lba, max_lba, disk_alignment) = {
            // enclose this in a block just to drop the ::gpt::Disk early

            // we now have to recalculate the partition actual size,
            // and move them all to LBA
            // create (but never write) a new GPT table
            // and work with that to translate to LBA
            let node = dev.devnode().ok_or(::std::io::Error::new(
                ::std::io::ErrorKind::Other,
                "device does not have a node",
            ))?;
            let gpt =
                ::gpt::GptConfig::new().writable(false).initialized(false);
            let mut disk = gpt.open(node)?;
            let mut new_partitions = ::std::collections::BTreeMap::<
                u32,
                ::gpt::partition::Partition,
            >::new();

            disk.update_partitions(new_partitions.clone()).unwrap();
            let disk_alignment = match disk.logical_block_size() {
                ::gpt::disk::LogicalBlockSize::Lb512 => 512,
                ::gpt::disk::LogicalBlockSize::Lb4096 => 4096,
            };
            let (min_lba, max_lba) =
                disk.find_free_sectors().get(0).unwrap().clone();
            (min_lba, max_lba, disk_alignment)
        };

        use trgt::gpt::{
            align_raw, Alignment, Index, PartFrom, PartId, PartTo, PartType,
            PartUUID, Partition,
        };
        // convert all to Index(LBA), Alignment(LBA) for ease of usage
        for p in &mut ret.partitions {
            p.from.lba_align(disk_alignment, min_lba, max_lba);
            p.to.lba_align(disk_alignment, min_lba, max_lba);
        }
        let mut last_lba: u64 = min_lba - 1;
        // now convert everything to
        // PartFrom::Index(Index::LBA(_), Alignment/LBA(0))
        for p in &mut ret.partitions {
            p.from = match &p.from {
                idx @ PartFrom::Index(_, _) => idx.clone(),
                PartFrom::SkipFree(
                    Index::LBA(l),
                    Alignment::Alignment(Index::LBA(al)),
                ) => {
                    let start_from = last_lba + 1 + l;
                    let start_aligned = align_raw(start_from, *al);
                    last_lba = start_aligned;
                    PartFrom::Index(
                        Index::LBA(start_aligned),
                        Alignment::Alignment(Index::LBA(0)),
                    )
                }
                PartFrom::SkipFreePercent(
                    perc,
                    Alignment::Alignment(Index::LBA(al)),
                ) => {
                    let skip_lba: u64 = (((max_lba - last_lba) as f64) / 100.0
                        * perc)
                        .ceil() as u64;
                    let start_from = last_lba + 1 + skip_lba;
                    let start_aligned = align_raw(start_from, *al);
                    last_lba = start_aligned;
                    PartFrom::Index(
                        Index::LBA(start_aligned),
                        Alignment::Alignment(Index::LBA(0)),
                    )
                }
                _ => {
                    assert!(false, "GPT: from: not everything is LBA");
                    return Err(::std::io::Error::new(
                        ::std::io::ErrorKind::Other,
                        "internal error: \"from\" should be LBA",
                    ));
                }
            };
            p.to = match &p.to {
                idx @ PartTo::Index(_, _) => idx.clone(),
                PartTo::Size(
                    Index::LBA(l),
                    Alignment::Alignment(Index::LBA(al)),
                ) => {
                    let up_to = last_lba + l;
                    let up_to_aligned = align_raw(up_to, *al);
                    last_lba = up_to_aligned;
                    PartTo::Index(
                        Index::LBA(up_to_aligned),
                        Alignment::Alignment(Index::LBA(0)),
                    )
                }
                PartTo::Leave(
                    Index::LBA(l),
                    Alignment::Alignment(Index::LBA(al)),
                ) => {
                    let up_to = max_lba - l;
                    let up_to_aligned = align_raw(up_to, *al);
                    last_lba = up_to_aligned;
                    PartTo::Index(
                        Index::LBA(up_to_aligned),
                        Alignment::Alignment(Index::LBA(0)),
                    )
                }
                PartTo::Percent(perc, Alignment::Alignment(Index::LBA(al))) => {
                    let size: u64 = (((max_lba - last_lba) as f64) / 100.0
                        * perc)
                        .ceil() as u64;
                    let up_to = last_lba + size;
                    let up_to_aligned = align_raw(up_to, *al);
                    last_lba = up_to_aligned;
                    PartTo::Index(
                        Index::LBA(up_to_aligned),
                        Alignment::Alignment(Index::LBA(0)),
                    )
                }
                PartTo::LeavePercent(
                    perc,
                    Alignment::Alignment(Index::LBA(al)),
                ) => {
                    let size: u64 = (((max_lba - last_lba) as f64) / 100.0
                        * perc)
                        .ceil() as u64;
                    let up_to = max_lba - size;
                    let up_to_aligned = align_raw(up_to, *al);
                    last_lba = up_to_aligned;
                    PartTo::Index(
                        Index::LBA(up_to_aligned),
                        Alignment::Alignment(Index::LBA(0)),
                    )
                }
                _ => {
                    assert!(false, "GPT: to: not everything is LBA");
                    return Err(::std::io::Error::new(
                        ::std::io::ErrorKind::Other,
                        "internal error: \"to\" should be LBA",
                    ));
                }
            };
            if let PartFrom::Index(Index::LBA(l_f), _) = p.from {
                if let PartTo::Index(Index::LBA(l_t), _) = p.to {
                    if l_f > l_t || l_t > max_lba || l_f < min_lba {
                        return Err(::std::io::Error::new(
                            ::std::io::ErrorKind::InvalidData,
                            "Wrong partition boundaries",
                        ));
                    }
                }
            };
        }
        Ok(ret)
    }
    /// get the current device GPT. make it standard enough that we can
    /// easily check with what comes from generate_standard(dev)
    fn get_current(
        dev: &::udev::Device,
    ) -> Result<trgt::gpt::GPT, ::std::io::Error> {
        let gpt = ::gpt::GptConfig::new().writable(false);
        let node = dev.devnode().ok_or(::std::io::Error::new(
            ::std::io::ErrorKind::Other,
            "device does not have a node",
        ))?;
        let disk = gpt.open(node)?;
        let disk_header =
            disk.primary_header().ok_or(::std::io::Error::new(
                ::std::io::ErrorKind::InvalidData,
                "Failed to retrieve GPT primary header",
            ))?;

        let mut ret = trgt::gpt::GPT {
            id: "DFIM generated for {}".to_owned()
                + dev.syspath().to_str().unwrap(),
            guid: trgt::gpt::GptUuid::UUID(disk.guid().clone()),
            max_partitions: disk_header.part_size,
            partitions: Vec::with_capacity(disk_header.num_parts as usize),
        };

        for (num, disk_p) in disk.partitions() {
            use trgt::gpt::{
                Alignment, Index, PartFrom, PartId, PartTo, PartUUID, Partition,
            };
            let p = Partition {
                label: disk_p.name.clone(),
                number: *num,
                part_type: PartId::UUID(
                    ::uuid::Uuid::parse_str(&disk_p.part_type_guid.guid)
                        .unwrap(),
                ),
                part_uuid: PartUUID::UUID(disk_p.part_guid),
                flags: trgt::gpt::Flag::from_bits(disk_p.flags),
                from: PartFrom::Index(
                    Index::LBA(disk_p.first_lba),
                    Alignment::Alignment(Index::LBA(0)),
                ),
                to: PartTo::Index(
                    Index::LBA(disk_p.last_lba),
                    Alignment::Alignment(Index::LBA(0)),
                ),
                target: trgt::TargetId::Id("DFIM generated None".to_owned()),
            };
            ret.partitions.push(p);
        }
        Ok(ret)
    }
    fn find_children(
        &self,
        cfg: &crate::config::trgt::gpt::GPT,
        dev: &::udev::Device,
    ) -> Result<Vec<NextTarget>, ::std::io::Error> {
        let mut scanner = ::udev::Enumerator::new()?;
        scanner.match_parent(dev)?;
        scanner.match_subsystem(
            crate::config::device::DevType::Partition.to_str(),
        )?;

        let mut ret = Vec::with_capacity(cfg.partitions.len());
        let children = scanner.scan_devices()?;

        for child in children {
            // these should be partitions of the disk
            // ...but how do we distinguish one from the other?
            match child.parent() {
                None => continue,
                Some(parent) => {
                    // I am unsure if we match also the children's children
                    // Therefore make sure we are
                    // really only looking at the first
                    // children of the dev we are given
                    if parent.syspath() != dev.syspath() {
                        continue;
                    }
                }
            }
            if None == child.devnode() {
                assert!(false, "We have a partition without a node?");
                continue;
            }
            let partnum = match child.sysnum() {
                None => {
                    assert!(false, "Partition without sys number?");
                    continue;
                }
                Some(partnum) => partnum as u32,
            };
            let p = match cfg.partitions.iter().find(|x| x.number == partnum) {
                Some(p) => p,
                None => {
                    return Err(::std::io::Error::new(
                        ::std::io::ErrorKind::InvalidData,
                        "We tought we had partitions we don't have",
                    ));
                }
            };
            ret.push(NextTarget {
                id: p.target.id().to_owned(),
                dev: child,
            });
        }

        Ok(ret)
    }
}

impl super::TargetApply for GPT {
    fn check(
        &mut self,
        logger: &::slog::Logger,
        cfg: &crate::state::cfg::Cfg,
        cfg_dev: &crate::config::device::Device,
        dev: &::udev::Device,
    ) -> Result<CheckResult, ::std::io::Error> {
        if let Ok(true) = crate::state::is_empty(dev) {
            return Ok(CheckResult {
                status: CheckStatus::CanBeApplied,
                next: Vec::new(),
            });
        }
        let to_be_applied = self.generate_standard(dev)?;
        let current_gpt = match GPT::get_current(dev) {
            Ok(current_gpt) => current_gpt,
            Err(e) => {
                ::slog::debug!(logger, "GPT load error: {}", e);
                return Ok(CheckResult {
                    status: if cfg_dev.only_if_empty {
                        CheckStatus::DoNotApply
                    } else {
                        CheckStatus::CanBeApplied
                    },
                    next: Vec::new(),
                });
            }
        };

        if to_be_applied != current_gpt {
            ::slog::debug!(
                logger,
                "GPT: device {:?} has different partitions than target {}",
                dev.syspath(),
                self.0.id
            );
            return Ok(CheckResult {
                status: if cfg_dev.only_if_empty {
                    CheckStatus::DoNotApply
                } else {
                    CheckStatus::CanBeApplied
                },
                next: Vec::new(),
            });
        }
        let next_targets = self.find_children(&to_be_applied, dev)?;
        Ok(CheckResult {
            status: CheckStatus::AlreadyApplied,
            next: next_targets,
        })
    }
    fn apply(
        &mut self,
        logger: &::slog::Logger,
        cfg: &crate::state::cfg::Cfg,
        cfg_dev: &crate::config::device::Device,
        dev: &::udev::Device,
    ) -> Result<Vec<NextTarget>, ::std::io::Error> {
        let node = dev.devnode().ok_or(::std::io::Error::new(
            ::std::io::ErrorKind::Other,
            "device does not have a node",
        ))?;
        if cfg_dev.only_if_empty {
            if !crate::state::is_empty(dev)? {
                return Err(::std::io::Error::new(
                    ::std::io::ErrorKind::InvalidData,
                    "GPT: can't apply: not empty disk",
                ));
            }
        }
        let to_be_applied = self.generate_standard(dev)?;

        let gptconf = ::gpt::GptConfig::new()
            .writable(!cfg.main.dry_run)
            .initialized(false);
        let mut disk = gptconf.open(node)?;

        if let trgt::gpt::GptUuid::UUID(disk_uuid) = to_be_applied.guid {
            disk.update_guid(Some(disk_uuid))?;
        }
        let mut disk_parts = ::std::collections::BTreeMap::<
            u32,
            ::gpt::partition::Partition,
        >::new();

        use trgt::gpt::{Index, PartFrom, PartId, PartTo, PartUUID};
        for conf_p in to_be_applied.partitions.iter() {
            let from_lba = match conf_p.from {
                PartFrom::Index(Index::LBA(from_lba), _) => from_lba,
                _ => {
                    return Err(::std::io::Error::new(
                        ::std::io::ErrorKind::Other,
                        "GPT: nonstandard standard partition",
                    ));
                }
            };
            let to_lba = match conf_p.to {
                PartTo::Index(Index::LBA(to_lba), _) => to_lba,
                _ => {
                    return Err(::std::io::Error::new(
                        ::std::io::ErrorKind::Other,
                        "GPT: nonstandard standard partition",
                    ));
                }
            };
            let type_uuid = match conf_p.part_type {
                PartId::UUID(type_uuid) => type_uuid,
                _ => {
                    return Err(::std::io::Error::new(
                        ::std::io::ErrorKind::Other,
                        "GPT: nonstandard standard partition",
                    ));
                }
            };
            let part_uuid = match conf_p.part_uuid {
                PartUUID::UUID(part_uuid) => part_uuid.clone(),
                PartUUID::Random => ::uuid::Uuid::new_v4(),
            };
            let p = ::gpt::partition::Partition {
                part_type_guid: ::gpt::partition_types::Type::from_uuid(
                    &type_uuid,
                )
                .unwrap(),
                part_guid: part_uuid,
                first_lba: from_lba,
                last_lba: to_lba,
                flags: conf_p.flag_bits(),
                name: conf_p.label.clone(),
            };
            disk_parts.insert(conf_p.number, p);
        }

        match disk.update_partitions_embedded(disk_parts, self.0.max_partitions)
        {
            Ok(_) => {}
            Err(e) => {
                ::slog::error!(
                    logger,
                    "GPT: target {}: wrong partition definition?",
                    self.0.id,
                );
                return Err(e);
            }
        }
        match cfg.main.dry_run {
            true => {
                ::slog::info!(
                    logger,
                    "DRY RUN: GPT: refusing to write target {} on drive {:?}",
                    self.0.id,
                    dev.syspath()
                );
                return Ok(Vec::new());
            }
            false => {
                match disk.write() {
                    Err(e) => {
                        ::slog::debug!(
                            logger,
                            "GPT: error writing table: {:?}",
                            e
                        );
                        return Err(e);
                    }
                    Ok(_) => {}
                }
                use std::os::unix::io::AsRawFd;
                let mut disk_file = match ::std::fs::File::open(node) {
                    Ok(disk_file) => disk_file,
                    Err(e) => {
                        ::slog::debug!(
                            logger,
                            "GPT: can't reload partition table: {:?}",
                            e
                        );
                        return Err(e);
                    }
                };
                let fd = unsafe { disk_file.as_raw_fd() };
                // reread partition table
                // include/uapi/linux/fs.h: #define BLKRRPART  _IO(0x12,95)
                ::nix::ioctl_none!(reread_partition_table, 0x12, 95);
                unsafe {
                    reread_partition_table(fd);
                }
            }
        }
        self.find_children(&to_be_applied, dev)
    }
}