diskotech 0.0.5

use anyhow::{bail, Result};
use clap::{ArgAction, Parser};
use diskotech::Record;
use prettytable::{format::FormatBuilder, Cell, Row, Table};
use std::{
    env,
    fmt::{Debug, Display},
    path::Path,
};

/// View relationship(s) between disk names (sda, nvme0n1, etc.), sequence, id,
/// partuuid, etc.
#[derive(Clone, Default, Parser)]
#[clap(version)]
struct Config {
    #[clap(short, long, action = ArgAction::Count)]
    verbose: u8,

    /// Display the fully-qualified paths to devices instead of short names.
    #[clap(long)]
    long: bool,

    /// Sort by a field.
    #[clap(short, long, default_value = "name")]
    sort_by: String,

    /// List of fields to include.
    #[clap(
        short,
        long,
        default_value = "diskseq,name,id,label,partlabel,partuuid,path,uuid"
    )]
    fields: String,
}

fn main() -> Result<()> {
    let cfg = Config::parse();

    if cfg.verbose > 0 {
        env::set_var("RUST_BACKTRACE", "1");
    }

    jacklog::from_level(
        2 + cfg.verbose as usize,
        Some(&[env!("CARGO_BIN_NAME")]),
    )?;

    let f: Box<dyn PathFormatter> = if cfg.long {
        Box::new(LongFmt {})
    } else {
        Box::new(ShortFmt {})
    };

    // Construct a sorting function.
    let sort_by = match cfg.sort_by.as_str() {
        "name" => |a: &Record, b: &Record| a.name.cmp(&b.name),
        "diskseq" => |a: &Record, b: &Record| a.diskseq.cmp(&b.diskseq),
        "label" => |a: &Record, b: &Record| a.label.cmp(&b.label),
        "partlabel" => |a: &Record, b: &Record| a.partlabel.cmp(&b.partlabel),
        "partuuid" => |a: &Record, b: &Record| a.partuuid.cmp(&b.partuuid),
        "path" => |a: &Record, b: &Record| a.path.cmp(&b.path),
        "uuid" => |a: &Record, b: &Record| a.uuid.cmp(&b.uuid),
        c => bail!("unrecognized column: {c}"),
    };

    // Collect the data.
    let mut records: Vec<Record> =
        diskotech::collect()?.into_values().collect();
    records.sort_by(sort_by);

    let mut table = Table::new();
    table.set_format(FormatBuilder::new().column_separator('\t').build());

    // Collect the headers we want to use.
    let fields: Vec<&str> = cfg.fields.split(',').collect();
    let headers = fields.iter().map(|f| Cell::new(f.trim())).collect();

    // Add the headers.
    table.add_row(Row::new(headers));
    for rec in records {
        table.add_row(Row::new(
            fields
                .iter()
                .map(|field| match *field {
                    "diskseq" => Cell::new(&f.fmt(rec.diskseq.as_deref())),
                    "name" => Cell::new(&f.fmt(rec.name.as_deref())),
                    "id" => Cell::new(
                        &rec.ids
                            .iter()
                            .map(|id| f.fmt(Some(id)))
                            .collect::<Vec<String>>()
                            .join("\n"),
                    ),
                    "label" => Cell::new(&f.fmt(rec.label.as_deref())),
                    "partlabel" => Cell::new(&f.fmt(rec.partlabel.as_deref())),
                    "partuuid" => Cell::new(&f.fmt(rec.partuuid.as_deref())),
                    "path" => Cell::new(&f.fmt(rec.path.as_deref())),
                    "uuid" => Cell::new(&f.fmt(rec.uuid.as_deref())),
                    f => Cell::new(&format!("unrecognized field: {f}")),
                })
                .collect(),
        ));
    }
    table.printstd();

    Ok(())
}

trait PathFormatter {
    fn fmt(&self, a: Option<&Path>) -> String;
}

struct ShortFmt;

impl PathFormatter for ShortFmt {
    fn fmt(&self, p: Option<&Path>) -> String {
        let Some(p)  = p else {
            return String::new();
        };

        let p: Box<dyn Display> = match p.file_name() {
            Some(p) => match p.to_str() {
                Some(p) => Box::new(p),
                None => Box::new(format!("{p:?}")),
            },
            None => Box::new(format!("{p:?}")),
        };

        format!("{p}")
    }
}

struct LongFmt;

impl PathFormatter for LongFmt {
    fn fmt(&self, p: Option<&Path>) -> String {
        let Some(p)  = p else {
            return String::new();
        };

        let p: Box<dyn Display> = match p.to_str() {
            Some(p) => Box::new(p),
            None => Box::new(format!("{p:?}")),
        };

        format!("{p}")
    }
}