datalab_cli/commands/

fill.rs

use crate::cache::Cache;
use crate::client::{add_form_field, build_form_with_file, DatalabClient};
use crate::error::{DatalabError, Result};
use crate::output::Progress;
use base64::Engine;
use clap::Args;
use reqwest::multipart::Form;
use serde_json::json;
use std::fs;
use std::path::PathBuf;

#[derive(Args, Debug)]
pub struct FillArgs {
    /// File path or URL of form to fill
    #[arg(value_name = "FILE|URL")]
    pub input: String,

    /// JSON file or string mapping field names to values
    #[arg(long, value_name = "JSON")]
    pub fields: String,

    /// Additional context for field matching
    #[arg(long, value_name = "TEXT", help_heading = "Matching Options")]
    pub context: Option<String>,

    /// Field matching strictness (0.0-1.0)
    #[arg(
        long,
        default_value = "0.5",
        value_name = "THRESHOLD",
        help_heading = "Matching Options"
    )]
    pub confidence_threshold: f32,

    /// Maximum pages to process
    #[arg(long, value_name = "N", help_heading = "Processing Options")]
    pub max_pages: Option<u32>,

    /// Page range (e.g., "0-5,10")
    #[arg(long, value_name = "RANGE", help_heading = "Processing Options")]
    pub page_range: Option<String>,

    /// Skip local cache lookup
    #[arg(long, help_heading = "Cache Options")]
    pub skip_cache: bool,

    /// Write filled form to file (binary output)
    #[arg(long, short, value_name = "FILE", help_heading = "Output Options")]
    pub output: Option<PathBuf>,

    /// Request timeout in seconds
    #[arg(
        long,
        default_value = "300",
        value_name = "SECS",
        help_heading = "Advanced Options"
    )]
    pub timeout: u64,
}

impl FillArgs {
    fn to_cache_params(&self) -> serde_json::Value {
        json!({
            "fields": self.fields,
            "context": self.context,
            "confidence_threshold": self.confidence_threshold,
            "max_pages": self.max_pages,
            "page_range": self.page_range,
        })
    }

    fn get_fields(&self) -> Result<String> {
        let fields_path = PathBuf::from(&self.fields);
        if fields_path.exists() {
            Ok(fs::read_to_string(&fields_path)?)
        } else {
            serde_json::from_str::<serde_json::Value>(&self.fields).map_err(|_| {
                DatalabError::InvalidInput(
                    "Fields must be valid JSON or a path to a JSON file".to_string(),
                )
            })?;
            Ok(self.fields.clone())
        }
    }

    fn add_to_form(&self, mut form: Form, fields: &str) -> Form {
        form = add_form_field(form, "field_data", fields);
        form = add_form_field(
            form,
            "confidence_threshold",
            &self.confidence_threshold.to_string(),
        );

        if let Some(ref context) = self.context {
            form = add_form_field(form, "context", context);
        }
        if let Some(max_pages) = self.max_pages {
            form = add_form_field(form, "max_pages", &max_pages.to_string());
        }
        if let Some(ref page_range) = self.page_range {
            form = add_form_field(form, "page_range", page_range);
        }

        form
    }
}

pub async fn execute(args: FillArgs, progress: &Progress) -> Result<()> {
    let client = DatalabClient::new(Some(args.timeout))?;
    let cache = Cache::new()?;

    let fields = args.get_fields()?;

    let is_url = args.input.starts_with("http://") || args.input.starts_with("https://");
    let file_path = if is_url {
        None
    } else {
        Some(PathBuf::from(&args.input))
    };

    let file_str = file_path.as_ref().map(|p| p.to_string_lossy().to_string());
    progress.start("fill", file_str.as_deref());

    let file_hash = if let Some(ref path) = file_path {
        if !path.exists() {
            return Err(DatalabError::FileNotFound(path.clone()));
        }
        Some(Cache::hash_file(path)?)
    } else {
        None
    };

    let cache_params = args.to_cache_params();
    let cache_key = Cache::generate_key(
        file_hash.as_deref(),
        if is_url { Some(&args.input) } else { None },
        "fill",
        &cache_params,
    );

    if !args.skip_cache {
        if let Some(cached) = cache.get(&cache_key) {
            progress.cache_hit(&cache_key);
            output_result(&cached, args.output.as_ref())?;
            return Ok(());
        }
    }

    let form = if let Some(ref path) = file_path {
        let (form, _) = build_form_with_file(path)?;
        args.add_to_form(form, &fields)
    } else {
        let form = Form::new().text("file_url", args.input.clone());
        args.add_to_form(form, &fields)
    };

    let result = client.submit_and_poll("fill", form, progress).await?;

    let file_path_str = file_path.as_ref().map(|p| p.to_string_lossy().to_string());
    cache.set(
        &cache_key,
        &result,
        "fill",
        file_hash.as_deref(),
        file_path_str.as_deref(),
    )?;

    output_result(&result, args.output.as_ref())?;

    Ok(())
}

fn output_result(result: &serde_json::Value, output_file: Option<&PathBuf>) -> Result<()> {
    if let Some(path) = output_file {
        if let Some(base64_data) = result.get("output_base64").and_then(|v| v.as_str()) {
            let decoded = base64::engine::general_purpose::STANDARD
                .decode(base64_data)
                .map_err(|e| DatalabError::InvalidInput(format!("Invalid base64: {}", e)))?;
            fs::write(path, &decoded)?;

            let mut meta = result.clone();
            meta.as_object_mut().map(|o| o.remove("output_base64"));
            println!("{}", serde_json::to_string_pretty(&meta)?);
        } else {
            println!("{}", serde_json::to_string_pretty(result)?);
        }
    } else {
        println!("{}", serde_json::to_string_pretty(result)?);
    }

    Ok(())
}