dag-ml-data-cli 0.2.2

use std::path::PathBuf;

use anyhow::{Context, Result};
use clap::{Parser, Subcommand};
use dag_ml_data_core::{
    data_plan_fingerprint, plan_model_input, sample_relation_fingerprint, schema_fingerprint,
    AdapterRegistry, AdapterRegistrySpec, CoordinatorDataMaterializationRequest,
    CoordinatorDataPlanEnvelope, CoordinatorHandleArena, DataPlan, DataPlanRequest, DatasetSchema,
    ModelInputSpec, SampleRelationTable, SourceId,
};

#[derive(Debug, Parser)]
#[command(author, version, about)]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Debug, Subcommand)]
enum Command {
    FingerprintSchema {
        path: PathBuf,
    },
    FingerprintPlan {
        path: PathBuf,
    },
    FingerprintRelations {
        path: PathBuf,
    },
    EnvelopePlan {
        #[arg(long)]
        schema: PathBuf,
        #[arg(long)]
        plan: PathBuf,
        #[arg(long)]
        relations: Option<PathBuf>,
    },
    ValidateEnvelope {
        path: PathBuf,
    },
    MaterializeEnvelope {
        #[arg(long)]
        envelope: PathBuf,
        #[arg(long)]
        request: PathBuf,
        #[arg(long, default_value = "controller:data.provider")]
        owner: String,
    },
    PlanModelInput {
        #[arg(long)]
        schema: PathBuf,
        #[arg(long)]
        model_input: PathBuf,
        #[arg(long)]
        adapters: PathBuf,
        #[arg(long)]
        id: String,
        #[arg(long = "source")]
        sources: Vec<String>,
    },
}

fn main() -> Result<()> {
    let cli = Cli::parse();

    match cli.command {
        Command::FingerprintSchema { path } => {
            let data = std::fs::read(&path)
                .with_context(|| format!("failed to read schema JSON at {}", path.display()))?;
            let schema: DatasetSchema = serde_json::from_slice(&data)
                .with_context(|| format!("failed to parse schema JSON at {}", path.display()))?;
            let fingerprint = schema_fingerprint(&schema)
                .with_context(|| format!("invalid schema at {}", path.display()))?;
            println!("{fingerprint}");
        }
        Command::FingerprintPlan { path } => {
            let plan: DataPlan = read_json(&path, "data plan")?;
            let fingerprint = data_plan_fingerprint(&plan)
                .with_context(|| format!("invalid data plan at {}", path.display()))?;
            println!("{fingerprint}");
        }
        Command::FingerprintRelations { path } => {
            let relations: SampleRelationTable = read_json(&path, "sample relations")?;
            let fingerprint = sample_relation_fingerprint(&relations)
                .with_context(|| format!("invalid sample relations at {}", path.display()))?;
            println!("{fingerprint}");
        }
        Command::EnvelopePlan {
            schema,
            plan,
            relations,
        } => {
            let schema: DatasetSchema = read_json(&schema, "schema")?;
            let plan: DataPlan = read_json(&plan, "data plan")?;
            let relations = relations
                .as_ref()
                .map(|path| read_json(path, "sample relations"))
                .transpose()?;
            let envelope =
                CoordinatorDataPlanEnvelope::from_parts(&schema, plan, relations.as_ref())
                    .context("failed to build coordinator data-plan envelope")?;
            println!("{}", serde_json::to_string_pretty(&envelope)?);
        }
        Command::ValidateEnvelope { path } => {
            let envelope: CoordinatorDataPlanEnvelope =
                read_json(&path, "coordinator data-plan envelope")?;
            envelope
                .validate()
                .with_context(|| format!("invalid data-plan envelope at {}", path.display()))?;
            println!(
                "valid data-plan envelope: plan_fingerprint={}, relations={}",
                envelope.plan_fingerprint,
                envelope
                    .coordinator_relations
                    .as_ref()
                    .map(|relations| relations.records.len())
                    .unwrap_or(0)
            );
        }
        Command::MaterializeEnvelope {
            envelope,
            request,
            owner,
        } => {
            let envelope: CoordinatorDataPlanEnvelope =
                read_json(&envelope, "coordinator data-plan envelope")?;
            let request: CoordinatorDataMaterializationRequest =
                read_json(&request, "coordinator data materialization request")?;
            let arena = CoordinatorHandleArena::new(owner)?;
            let record = arena
                .materialize(&envelope, &request)
                .with_context(|| "failed to materialize coordinator data handle")?;
            println!("{}", serde_json::to_string_pretty(&record)?);
        }
        Command::PlanModelInput {
            schema,
            model_input,
            adapters,
            id,
            sources,
        } => {
            let schema: DatasetSchema = read_json(&schema, "schema")?;
            let model_input: ModelInputSpec = read_json(&model_input, "model input")?;
            let registry_spec: AdapterRegistrySpec = read_json(&adapters, "adapter registry")?;
            let registry = AdapterRegistry::from_spec(registry_spec)
                .with_context(|| format!("invalid adapter registry at {}", adapters.display()))?;
            let plan = plan_model_input(
                &schema,
                &model_input,
                &registry,
                &DataPlanRequest {
                    id,
                    source_ids: (!sources.is_empty())
                        .then(|| {
                            sources
                                .into_iter()
                                .map(SourceId::new)
                                .collect::<dag_ml_data_core::Result<Vec<_>>>()
                        })
                        .transpose()?,
                    planning_policy: Default::default(),
                },
            )
            .context("failed to plan model input")?;
            println!("{}", serde_json::to_string_pretty(&plan)?);
        }
    }

    Ok(())
}

fn read_json<T: serde::de::DeserializeOwned>(path: &PathBuf, label: &str) -> Result<T> {
    let data = std::fs::read(path)
        .with_context(|| format!("failed to read {label} JSON at {}", path.display()))?;
    serde_json::from_slice(&data)
        .with_context(|| format!("failed to parse {label} JSON at {}", path.display()))
}