start-os/core/startos/src/registry/package/get.rs

use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};

use clap::{Parser, ValueEnum};
use exver::{ExtendedVersion, VersionRange};
use helpers::to_tmp_path;
use imbl_value::{InternedString, json};
use itertools::Itertools;
use models::PackageId;
use serde::{Deserialize, Serialize};
use ts_rs::TS;

use crate::context::CliContext;
use crate::prelude::*;
use crate::progress::{FullProgressTracker, ProgressUnits};
use crate::registry::context::RegistryContext;
use crate::registry::device_info::DeviceInfo;
use crate::registry::package::index::{PackageIndex, PackageVersionInfo};
use crate::s9pk::merkle_archive::source::ArchiveSource;
use crate::s9pk::v2::SIG_CONTEXT;
use crate::util::VersionString;
use crate::util::io::TrackingIO;
use crate::util::serde::{WithIoFormat, display_serializable};
use crate::util::tui::choose;

#[derive(
    Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize, TS, ValueEnum,
)]
#[serde(rename_all = "camelCase")]
#[ts(export)]
pub enum PackageDetailLevel {
    None,
    Short,
    Full,
}
impl Default for PackageDetailLevel {
    fn default() -> Self {
        Self::Short
    }
}

#[derive(Debug, Deserialize, Serialize, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export)]
pub struct PackageInfoShort {
    pub release_notes: String,
}

#[derive(Debug, Deserialize, Serialize, TS, Parser)]
#[serde(rename_all = "camelCase")]
#[command(rename_all = "kebab-case")]
#[ts(export)]
pub struct GetPackageParams {
    pub id: Option<PackageId>,
    #[ts(type = "string | null")]
    #[arg(long, short = 'v')]
    pub target_version: Option<VersionRange>,
    #[arg(long)]
    pub source_version: Option<VersionString>,
    #[ts(skip)]
    #[arg(skip)]
    #[serde(rename = "__DeviceInfo_device_info")]
    pub device_info: Option<DeviceInfo>,
    #[serde(default)]
    #[arg(default_value = "none")]
    pub other_versions: PackageDetailLevel,
}

#[derive(Debug, Deserialize, Serialize, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export)]
pub struct GetPackageResponse {
    #[ts(type = "string[]")]
    pub categories: BTreeSet<InternedString>,
    pub best: BTreeMap<VersionString, PackageVersionInfo>,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[ts(optional)]
    pub other_versions: Option<BTreeMap<VersionString, PackageInfoShort>>,
}
impl GetPackageResponse {
    pub fn tables(&self) -> Vec<prettytable::Table> {
        use prettytable::*;

        let mut res = Vec::with_capacity(self.best.len());

        for (version, info) in &self.best {
            let mut table = info.table(version);

            let lesser_versions: BTreeMap<_, _> = self
                .other_versions
                .as_ref()
                .into_iter()
                .flatten()
                .filter(|(v, _)| ***v < **version)
                .collect();

            if !lesser_versions.is_empty() {
                table.add_row(row![bc => "OLDER VERSIONS"]);
                table.add_row(row![bc => "VERSION", "RELEASE NOTES"]);
                for (version, info) in lesser_versions {
                    table.add_row(row![AsRef::<str>::as_ref(version), &info.release_notes]);
                }
            }

            res.push(table);
        }

        res
    }
}

#[derive(Debug, Deserialize, Serialize, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export)]
pub struct GetPackageResponseFull {
    #[ts(type = "string[]")]
    pub categories: BTreeSet<InternedString>,
    pub best: BTreeMap<VersionString, PackageVersionInfo>,
    pub other_versions: BTreeMap<VersionString, PackageVersionInfo>,
}
impl GetPackageResponseFull {
    pub fn tables(&self) -> Vec<prettytable::Table> {
        let mut res = Vec::with_capacity(self.best.len());

        let all: BTreeMap<_, _> = self.best.iter().chain(self.other_versions.iter()).collect();

        for (version, info) in all {
            res.push(info.table(version));
        }

        res
    }
}

pub type GetPackagesResponse = BTreeMap<PackageId, GetPackageResponse>;
pub type GetPackagesResponseFull = BTreeMap<PackageId, GetPackageResponseFull>;

fn get_matching_models<'a>(
    db: &'a Model<PackageIndex>,
    GetPackageParams {
        id,
        source_version,
        device_info,
        ..
    }: &GetPackageParams,
) -> Result<Vec<(PackageId, ExtendedVersion, &'a Model<PackageVersionInfo>)>, Error> {
    if let Some(id) = id {
        if let Some(pkg) = db.as_packages().as_idx(id) {
            vec![(id.clone(), pkg)]
        } else {
            vec![]
        }
    } else {
        db.as_packages().as_entries()?
    }
    .iter()
    .map(|(k, v)| {
        Ok(v.as_versions()
            .as_entries()?
            .into_iter()
            .map(|(v, info)| {
                Ok::<_, Error>(
                    if source_version.as_ref().map_or(Ok(true), |source_version| {
                        Ok::<_, Error>(
                            source_version.satisfies(
                                &info
                                    .as_source_version()
                                    .de()?
                                    .unwrap_or(VersionRange::any()),
                            ),
                        )
                    })? && device_info
                        .as_ref()
                        .map_or(Ok(true), |device_info| info.works_for_device(device_info))?
                    {
                        Some((k.clone(), ExtendedVersion::from(v), info))
                    } else {
                        None
                    },
                )
            })
            .flatten_ok())
    })
    .flatten_ok()
    .map(|res| res.and_then(|a| a))
    .collect()
}

pub async fn get_package(ctx: RegistryContext, params: GetPackageParams) -> Result<Value, Error> {
    use patch_db::ModelExt;

    let peek = ctx.db.peek().await;
    let mut best: BTreeMap<PackageId, BTreeMap<VersionString, &Model<PackageVersionInfo>>> =
        Default::default();
    let mut other: BTreeMap<PackageId, BTreeMap<VersionString, &Model<PackageVersionInfo>>> =
        Default::default();
    for (id, version, info) in get_matching_models(&peek.as_index().as_package(), &params)? {
        let package_best = best.entry(id.clone()).or_default();
        let package_other = other.entry(id.clone()).or_default();
        if params
            .target_version
            .as_ref()
            .map_or(true, |v| version.satisfies(v))
            && package_best.keys().all(|k| !(**k > version))
        {
            for worse_version in package_best
                .keys()
                .filter(|k| ***k < version)
                .cloned()
                .collect_vec()
            {
                if let Some(info) = package_best.remove(&worse_version) {
                    package_other.insert(worse_version, info);
                }
            }
            package_best.insert(version.into(), info);
        } else {
            package_other.insert(version.into(), info);
        }
    }
    if let Some(id) = params.id {
        let categories = peek
            .as_index()
            .as_package()
            .as_packages()
            .as_idx(&id)
            .map(|p| p.as_categories().de())
            .transpose()?
            .unwrap_or_default();
        let best = best
            .remove(&id)
            .unwrap_or_default()
            .into_iter()
            .map(|(k, v)| v.de().map(|v| (k, v)))
            .try_collect()?;
        let other = other.remove(&id).unwrap_or_default();
        match params.other_versions {
            PackageDetailLevel::None => to_value(&GetPackageResponse {
                categories,
                best,
                other_versions: None,
            }),
            PackageDetailLevel::Short => to_value(&GetPackageResponse {
                categories,
                best,
                other_versions: Some(
                    other
                        .into_iter()
                        .map(|(k, v)| from_value(v.as_value().clone()).map(|v| (k, v)))
                        .try_collect()?,
                ),
            }),
            PackageDetailLevel::Full => to_value(&GetPackageResponseFull {
                categories,
                best,
                other_versions: other
                    .into_iter()
                    .map(|(k, v)| v.de().map(|v| (k, v)))
                    .try_collect()?,
            }),
        }
    } else {
        match params.other_versions {
            PackageDetailLevel::None => to_value(
                &best
                    .into_iter()
                    .map(|(id, best)| {
                        let categories = peek
                            .as_index()
                            .as_package()
                            .as_packages()
                            .as_idx(&id)
                            .map(|p| p.as_categories().de())
                            .transpose()?
                            .unwrap_or_default();
                        Ok::<_, Error>((
                            id,
                            GetPackageResponse {
                                categories,
                                best: best
                                    .into_iter()
                                    .map(|(k, v)| v.de().map(|v| (k, v)))
                                    .try_collect()?,
                                other_versions: None,
                            },
                        ))
                    })
                    .try_collect::<_, GetPackagesResponse, _>()?,
            ),
            PackageDetailLevel::Short => to_value(
                &best
                    .into_iter()
                    .map(|(id, best)| {
                        let categories = peek
                            .as_index()
                            .as_package()
                            .as_packages()
                            .as_idx(&id)
                            .map(|p| p.as_categories().de())
                            .transpose()?
                            .unwrap_or_default();
                        let other = other.remove(&id).unwrap_or_default();
                        Ok::<_, Error>((
                            id,
                            GetPackageResponse {
                                categories,
                                best: best
                                    .into_iter()
                                    .map(|(k, v)| v.de().map(|v| (k, v)))
                                    .try_collect()?,
                                other_versions: Some(
                                    other
                                        .into_iter()
                                        .map(|(k, v)| {
                                            from_value(v.as_value().clone()).map(|v| (k, v))
                                        })
                                        .try_collect()?,
                                ),
                            },
                        ))
                    })
                    .try_collect::<_, GetPackagesResponse, _>()?,
            ),
            PackageDetailLevel::Full => to_value(
                &best
                    .into_iter()
                    .map(|(id, best)| {
                        let categories = peek
                            .as_index()
                            .as_package()
                            .as_packages()
                            .as_idx(&id)
                            .map(|p| p.as_categories().de())
                            .transpose()?
                            .unwrap_or_default();
                        let other = other.remove(&id).unwrap_or_default();
                        Ok::<_, Error>((
                            id,
                            GetPackageResponseFull {
                                categories,
                                best: best
                                    .into_iter()
                                    .map(|(k, v)| v.de().map(|v| (k, v)))
                                    .try_collect()?,
                                other_versions: other
                                    .into_iter()
                                    .map(|(k, v)| v.de().map(|v| (k, v)))
                                    .try_collect()?,
                            },
                        ))
                    })
                    .try_collect::<_, GetPackagesResponseFull, _>()?,
            ),
        }
    }
}

pub fn display_package_info(
    params: WithIoFormat<GetPackageParams>,
    info: Value,
) -> Result<(), Error> {
    if let Some(format) = params.format {
        return display_serializable(format, info);
    }

    if let Some(_) = params.rest.id {
        if params.rest.other_versions == PackageDetailLevel::Full {
            for table in from_value::<GetPackageResponseFull>(info)?.tables() {
                table.print_tty(false)?;
                println!();
            }
        } else {
            for table in from_value::<GetPackageResponse>(info)?.tables() {
                table.print_tty(false)?;
                println!();
            }
        }
    } else {
        if params.rest.other_versions == PackageDetailLevel::Full {
            for (_, package) in from_value::<GetPackagesResponseFull>(info)? {
                for table in package.tables() {
                    table.print_tty(false)?;
                    println!();
                }
            }
        } else {
            for (_, package) in from_value::<GetPackagesResponse>(info)? {
                for table in package.tables() {
                    table.print_tty(false)?;
                    println!();
                }
            }
        }
    }
    Ok(())
}

#[derive(Debug, Deserialize, Serialize, TS, Parser)]
#[serde(rename_all = "camelCase")]
pub struct CliDownloadParams {
    pub id: PackageId,
    #[arg(long, short = 'v')]
    #[ts(type = "string | null")]
    pub target_version: Option<VersionRange>,
    #[arg(short, long)]
    pub dest: Option<PathBuf>,
}

pub async fn cli_download(
    ctx: CliContext,
    CliDownloadParams {
        ref id,
        target_version,
        dest,
    }: CliDownloadParams,
) -> Result<(), Error> {
    let progress_tracker = FullProgressTracker::new();
    let mut fetching_progress = progress_tracker.add_phase("Fetching".into(), Some(1));
    let download_progress = progress_tracker.add_phase("Downloading".into(), Some(100));
    let mut verify_progress = progress_tracker.add_phase("Verifying".into(), Some(10));

    let progress = progress_tracker.progress_bar_task("Downloading S9PK...");

    fetching_progress.start();
    let mut res: GetPackageResponse = from_value(
        ctx.call_remote::<RegistryContext>(
            "package.get",
            json!({
                "id": &id,
                "targetVersion": &target_version,
            }),
        )
        .await?,
    )?;
    let PackageVersionInfo { s9pk, .. } = match res.best.len() {
        0 => {
            return Err(Error::new(
                eyre!(
                    "Could not find a version of {id} that satisfies {}",
                    target_version.unwrap_or(VersionRange::Any)
                ),
                ErrorKind::NotFound,
            ));
        }
        1 => res.best.pop_first().unwrap().1,
        _ => {
            let choices = res.best.keys().cloned().collect::<Vec<_>>();
            let version = choose(
                &format!("Multiple flavors of {id} available. Choose a version to download:"),
                &choices,
            )
            .await?;
            res.best.remove(version).unwrap()
        }
    };
    s9pk.validate(SIG_CONTEXT, s9pk.all_signers())?;
    fetching_progress.complete();

    let dest = dest.unwrap_or_else(|| Path::new(".").join(id).with_extension("s9pk"));
    let dest_tmp = to_tmp_path(&dest).with_kind(ErrorKind::Filesystem)?;
    let (mut parsed, source) = s9pk
        .download_to(&dest_tmp, ctx.client.clone(), download_progress)
        .await?;
    if let Some(size) = source.size().await {
        verify_progress.set_total(size);
    }
    verify_progress.set_units(Some(ProgressUnits::Bytes));
    let mut progress_sink = verify_progress.writer(tokio::io::sink());
    parsed
        .serialize(&mut TrackingIO::new(0, &mut progress_sink), true)
        .await?;
    progress_sink.into_inner().1.complete();

    source.wait_for_buffered().await?;
    tokio::fs::rename(dest_tmp, dest).await?;

    progress_tracker.complete();
    progress.await.unwrap();

    println!("Download Complete");

    Ok(())
}