Adom Chipsmith

use crate::cli::{eval, err, ok, DEFAULT_PORT};
use anyhow::Result;
use clap::Parser;
use std::path::PathBuf;

#[derive(Parser)]
pub struct Args {
    /// Output PNG path. Defaults to /tmp/adom-step-shot.png.
    #[arg(long)]
    pub out: Option<PathBuf>,

    #[arg(long, default_value_t = 1600)]
    pub width: u32,

    #[arg(long, default_value_t = 1000)]
    pub height: u32,

    #[arg(long, default_value_t = 10)]
    pub timeout_secs: u64,

    #[arg(long, default_value_t = DEFAULT_PORT)]
    pub port: u16,
}

pub fn run(args: Args) -> Result<()> {
    let out = args
        .out
        .unwrap_or_else(|| PathBuf::from("/tmp/adom-step-shot.png"));

    let code = format!(
        "return window.adomStep ? await window.adomStep.screenshot({{width: {}, height: {}}}) : {{ error: 'adomStep not loaded' }};",
        args.width, args.height
    );
    let result = eval::eval_snippet(&code, args.timeout_secs, args.port, true)?;
    if let Some(e) = result.get("error").and_then(|v| v.as_str()) {
        err(format!("screenshot failed: {e}"));
        return Err(anyhow::anyhow!("{e}"));
    }
    let data_url = result.as_str().unwrap_or("");
    let prefix = "data:image/png;base64,";
    if !data_url.starts_with(prefix) {
        err(format!("unexpected screenshot payload: {data_url}"));
        return Err(anyhow::anyhow!("bad payload"));
    }
    let b64 = &data_url[prefix.len()..];
    let bytes = base64_decode(b64)?;
    std::fs::write(&out, &bytes)?;
    ok(format!(
        "wrote {} ({} KB)",
        out.display(),
        bytes.len() / 1024
    ));
    Ok(())
}

/// Lean base64 decoder so we don't pull in the `base64` crate for one use.
fn base64_decode(s: &str) -> Result<Vec<u8>> {
    let s: String = s.chars().filter(|c| !c.is_whitespace()).collect();
    let map = |c: char| -> Result<u32> {
        Ok(match c {
            'A'..='Z' => c as u32 - 'A' as u32,
            'a'..='z' => c as u32 - 'a' as u32 + 26,
            '0'..='9' => c as u32 - '0' as u32 + 52,
            '+' => 62,
            '/' => 63,
            '=' => 0,
            _ => return Err(anyhow::anyhow!("bad base64 char: {c}")),
        })
    };
    let bytes = s.as_bytes();
    let mut out = Vec::with_capacity(bytes.len() * 3 / 4);
    let mut i = 0;
    while i + 3 < bytes.len() {
        let a = map(bytes[i] as char)?;
        let b = map(bytes[i + 1] as char)?;
        let c = map(bytes[i + 2] as char)?;
        let d = map(bytes[i + 3] as char)?;
        let v = (a << 18) | (b << 12) | (c << 6) | d;
        out.push((v >> 16) as u8);
        if bytes[i + 2] != b'=' {
            out.push(((v >> 8) & 0xFF) as u8);
        }
        if bytes[i + 3] != b'=' {
            out.push((v & 0xFF) as u8);
        }
        i += 4;
    }
    Ok(out)
}