ff-rdp-cli 0.1.0

use std::path::PathBuf;

use anyhow::Context as _;
use base64::Engine as _;
use ff_rdp_core::{Grip, LongStringActor, ScreenshotActor, ScreenshotContentActor};
use serde_json::json;

use crate::cli::args::Cli;
use crate::error::AppError;
use crate::hints::{HintContext, HintSource};
use crate::output;
use crate::output_pipeline::OutputPipeline;

use super::connect_tab::connect_direct;
use super::js_helpers::eval_or_bail;

/// Build the JavaScript injected into the page to capture a screenshot.
///
/// Uses `CanvasRenderingContext2D.drawWindow`, a Firefox-specific privileged
/// API that renders a region starting at `(0, 0)` into an offscreen canvas.
/// Returns a `data:image/png;base64,...` string on success, or `null` when
/// `drawWindow` is not available (removed in some Firefox configurations).
///
/// `height_override` selects the capture height:
/// - `None` (viewport, default): `window.innerHeight`.
/// - `Some(HeightOverride::FullPage)`: `document.scrollingElement.scrollHeight`.
/// - `Some(HeightOverride::Explicit(n))`: the literal value `n`.
fn build_screenshot_js(height_override: Option<HeightOverride>) -> String {
    let height_expr = match height_override {
        None => "window.innerHeight || document.documentElement.clientHeight || 600".to_owned(),
        Some(HeightOverride::FullPage) => {
            // Take the max of scrollHeight candidates — different browsers
            // populate different properties.
            "Math.max(\
              document.documentElement.scrollHeight,\
              document.body ? document.body.scrollHeight : 0,\
              (document.scrollingElement && document.scrollingElement.scrollHeight) || 0,\
              window.innerHeight || 0\
            )"
            .to_owned()
        }
        Some(HeightOverride::Explicit(n)) => n.to_string(),
    };
    format!(
        "(function() {{\n  var w = window.innerWidth || document.documentElement.clientWidth || 800;\n  var h = {height_expr};\n  var canvas = document.createElement('canvas');\n  canvas.width = w;\n  canvas.height = h;\n  var ctx = canvas.getContext('2d');\n  if (!ctx || typeof ctx.drawWindow !== 'function') {{ return null; }}\n  ctx.drawWindow(window, 0, 0, w, h, 'rgb(255,255,255)');\n  return canvas.toDataURL('image/png');\n}})()"
    )
}

/// Height mode selected by the user.
#[derive(Copy, Clone, Debug)]
enum HeightOverride {
    FullPage,
    Explicit(u32),
}

/// Options accepted by [`run`].
pub(crate) struct ScreenshotOpts<'a> {
    pub(crate) output_path: Option<&'a str>,
    pub(crate) base64_mode: bool,
    pub(crate) full_page: bool,
    pub(crate) viewport_height: Option<u32>,
}

/// Data URL prefix returned by `canvas.toDataURL('image/png')`.
const PNG_DATA_URL_PREFIX: &str = "data:image/png;base64,";

pub fn run(cli: &Cli, opts: &ScreenshotOpts<'_>) -> Result<(), AppError> {
    let height_override = match (opts.full_page, opts.viewport_height) {
        (true, Some(_)) => {
            return Err(AppError::User(
                "screenshot: --full-page and --viewport-height are mutually exclusive".to_owned(),
            ));
        }
        (true, None) => Some(HeightOverride::FullPage),
        (false, Some(n)) => Some(HeightOverride::Explicit(n)),
        (false, None) => None,
    };
    let output_path = opts.output_path;
    let base64_mode = opts.base64_mode;
    let screenshot_js = build_screenshot_js(height_override);
    let screenshot_js = screenshot_js.as_str();
    // Screenshot always connects directly to Firefox, bypassing the daemon.
    // The daemon's watcher subscription interferes with the two-step screenshot
    // protocol, causing Firefox-side timeouts.
    let mut ctx = connect_direct(cli)?;
    let console_actor = ctx.target.console_actor.clone();

    let eval_result = eval_or_bail(
        &mut ctx,
        &console_actor,
        screenshot_js,
        "screenshot JS threw an exception",
    )?;

    // Resolve the result — the data URL may come back as a LongString when the
    // PNG is large enough to exceed Firefox's inline-string threshold.
    let data_url = if let Some(url) = resolve_string(&mut ctx, &eval_result.result)? {
        url
    } else {
        // drawWindow unavailable.  Try fallbacks in order:
        //   1. Legacy single-step screenshotContentActor (Firefox < 149)
        //   2. Two-step protocol: screenshotContentActor.prepareCapture +
        //      root screenshotActor.capture (Firefox 149+)
        //
        // Clone actor IDs to release the borrow on `ctx.target` before
        // taking a mutable borrow on `ctx` for transport calls.
        let sc_actor = ctx.target.screenshot_content_actor.clone();
        let browsing_ctx_id = ctx.target.browsing_context_id;

        if let Some(actor) = sc_actor {
            // Try legacy method first.
            match ScreenshotContentActor::capture(
                ctx.transport_mut(),
                actor.as_ref(),
                height_override.is_some_and(|h| matches!(h, HeightOverride::FullPage)),
            ) {
                Ok(capture) => capture.data,
                Err(legacy_err) if legacy_err.is_unrecognized_packet_type() => {
                    // Legacy method not available — try the Firefox 149+ two-step protocol.
                    try_two_step_screenshot(
                        &mut ctx,
                        &actor,
                        browsing_ctx_id,
                        height_override.is_some_and(|h| matches!(h, HeightOverride::FullPage)),
                    )?
                }
                Err(e) => {
                    return Err(AppError::User(format!(
                        "screenshot: screenshotContentActor capture failed ({e}) — \
                         screenshots require headless mode; relaunch with: ff-rdp launch --headless"
                    )));
                }
            }
        } else {
            return Err(AppError::User(
                "screenshot: drawWindow not available and no screenshotContentActor found — \
                 screenshots require headless mode; relaunch with: ff-rdp launch --headless"
                    .to_owned(),
            ));
        }
    };

    let b64 = data_url.strip_prefix(PNG_DATA_URL_PREFIX).ok_or_else(|| {
        AppError::User(format!(
            "screenshot: unexpected data URL format (expected prefix '{PNG_DATA_URL_PREFIX}')"
        ))
    })?;

    // Decode bytes only for dimension extraction and file output.
    // For base64 mode the raw `b64` string is used directly — no need to
    // decode and re-encode.
    let png_bytes = base64::engine::general_purpose::STANDARD
        .decode(b64)
        .map_err(|e| AppError::from(anyhow::anyhow!("screenshot: base64 decode failed: {e}")))?;

    // Infer dimensions from PNG header: width at bytes 16–19, height at 20–23.
    let (width, height) = png_dimensions(&png_bytes).unwrap_or((0, 0));

    let results = if base64_mode {
        // Return the base64 string directly — strip the data-URL prefix but
        // do not decode+re-encode; the `b64` slice is already valid
        // standard-base64.
        json!({
            "base64": b64,
            "width": width,
            "height": height,
            "bytes": png_bytes.len(),
        })
    } else {
        let dest = resolve_output_path(output_path)
            .map_err(|e| AppError::from(anyhow::anyhow!("screenshot: {e}")))?;

        std::fs::write(&dest, &png_bytes)
            .with_context(|| format!("screenshot: could not write to '{}'", dest.display()))
            .map_err(AppError::from)?;

        let abs_path = dest
            .canonicalize()
            .unwrap_or(dest)
            .to_string_lossy()
            .into_owned();

        json!({
            "path": abs_path,
            "width": width,
            "height": height,
            "bytes": png_bytes.len(),
        })
    };

    let meta = json!({"host": cli.host, "port": cli.port});
    let envelope = output::envelope(&results, 1, &meta);

    let hint_ctx = HintContext::new(HintSource::Screenshot);
    OutputPipeline::from_cli(cli)?
        .finalize_with_hints(&envelope, Some(&hint_ctx))
        .map_err(AppError::from)
}

/// Firefox 149+ two-step screenshot protocol.
///
/// Step 1: `screenshotContentActor.prepareCapture` → viewport DPR/zoom/rect
/// Step 2: `screenshotActor.capture` (root actor) → PNG data URL
///
/// `full_page` is forwarded to both steps so Firefox captures the full scroll
/// height rather than just the visible viewport.
///
/// Returns the `data:image/png;base64,...` string on success.
fn try_two_step_screenshot(
    ctx: &mut super::connect_tab::ConnectedTab,
    sc_actor: &ff_rdp_core::ActorId,
    browsing_ctx_id: Option<u64>,
    full_page: bool,
) -> Result<String, AppError> {
    let browsing_ctx_id = browsing_ctx_id.ok_or_else(|| {
        AppError::User(
            "screenshot: Firefox 149+ screenshot requires a browsing context ID \
             which was not found in the target response. \
             Try upgrading ff-rdp or filing a bug with your Firefox version."
                .to_owned(),
        )
    })?;

    // Step 1: prepare — collect viewport DPR/zoom from the content process actor.
    let prep =
        ScreenshotContentActor::prepare_capture(ctx.transport_mut(), sc_actor.as_ref(), full_page)
            .map_err(|e| {
                AppError::User(format!(
                    "screenshot: screenshotContentActor.prepareCapture failed ({e})"
                ))
            })?;

    // Step 2: capture — call the root-level screenshotActor.
    let screenshot_actor = ScreenshotActor::get_actor_id(ctx.transport_mut()).map_err(|e| {
        AppError::User(format!(
            "screenshot: could not find root screenshotActor ({e}) — \
             this Firefox version may not support the two-step screenshot protocol"
        ))
    })?;

    let data = ScreenshotActor::capture(
        ctx.transport_mut(),
        &screenshot_actor,
        browsing_ctx_id,
        full_page,
        &prep,
    )
    .map_err(|e| {
        AppError::User(format!(
            "screenshot: screenshotActor.capture failed ({e}) — \
             screenshots require headless mode; relaunch with: ff-rdp launch --headless"
        ))
    })?;

    Ok(data)
}

/// Resolve the eval result to an `Option<String>`:
/// - `Grip::Null` → `None` (drawWindow unavailable)
/// - `Grip::Value(String)` → `Some(string)`
/// - `Grip::LongString` → fetch the full string and return `Some`
fn resolve_string(
    ctx: &mut super::connect_tab::ConnectedTab,
    grip: &Grip,
) -> Result<Option<String>, AppError> {
    match grip {
        Grip::Null | Grip::Undefined => Ok(None),
        Grip::Value(serde_json::Value::String(s)) => Ok(Some(s.clone())),
        Grip::LongString {
            actor,
            length,
            initial: _,
        } => {
            let full = LongStringActor::full_string(ctx.transport_mut(), actor.as_ref(), *length)
                .map_err(AppError::from)?;
            Ok(Some(full))
        }
        other => Err(AppError::User(format!(
            "screenshot: unexpected result type: {}",
            other.to_json()
        ))),
    }
}

/// Determine the output file path.
///
/// If the caller provided an explicit path, use it.  Otherwise generate a
/// timestamped filename in the current directory.
fn resolve_output_path(output_path: Option<&str>) -> anyhow::Result<PathBuf> {
    if let Some(p) = output_path {
        return Ok(PathBuf::from(p));
    }

    // Generate a unique filename from the current system time (millisecond
    // resolution to avoid collisions when taking multiple screenshots quickly).
    let ts = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .context("system clock is before UNIX epoch")?
        .as_millis();

    Ok(PathBuf::from(format!("screenshot-{ts}.png")))
}

/// Extract width and height from a PNG file's IHDR chunk.
///
/// PNG structure: 8-byte signature, then chunks. The first chunk is always
/// IHDR which contains `width` (4 bytes, big-endian) at offset 16 and
/// `height` (4 bytes, big-endian) at offset 20.
fn png_dimensions(data: &[u8]) -> Option<(u32, u32)> {
    if data.len() < 24 {
        return None;
    }
    let width = u32::from_be_bytes(data[16..20].try_into().ok()?);
    let height = u32::from_be_bytes(data[20..24].try_into().ok()?);
    Some((width, height))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cli::args::{Cli, Command};
    use clap::Parser as _;

    // ── clap parse tests ────────────────────────────────────────────────────

    #[test]
    fn clap_screenshot_full_page_flag_parsed() {
        let cli = Cli::try_parse_from(["ff-rdp", "screenshot", "--full-page"])
            .expect("should parse --full-page");
        let Command::Screenshot { full_page, .. } = cli.command else {
            panic!("expected Screenshot command");
        };
        assert!(full_page, "--full-page flag must be set");
    }

    #[test]
    fn clap_a11y_limit_and_format_text_parsed() {
        let cli = Cli::try_parse_from(["ff-rdp", "a11y", "--limit", "5", "--format", "text"])
            .expect("should parse a11y --limit 5 --format text");
        assert_eq!(cli.limit, Some(5));
        assert_eq!(cli.format, "text");
        assert!(matches!(cli.command, Command::A11y { .. }));
    }

    #[test]
    fn png_dimensions_minimal_png() {
        // Minimal 1x1 white PNG generated from known bytes.
        let b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC";
        let bytes = base64::engine::general_purpose::STANDARD
            .decode(b64)
            .unwrap();
        let (w, h) = png_dimensions(&bytes).unwrap();
        assert_eq!(w, 1);
        assert_eq!(h, 1);
    }

    #[test]
    fn png_dimensions_too_short() {
        assert!(png_dimensions(&[0u8; 10]).is_none());
    }

    #[test]
    fn resolve_output_path_explicit() {
        let path = resolve_output_path(Some("/tmp/test.png")).unwrap();
        assert_eq!(path, PathBuf::from("/tmp/test.png"));
    }

    #[test]
    fn resolve_output_path_auto_timestamped() {
        let path = resolve_output_path(None).unwrap();
        let name = path.to_string_lossy();
        assert!(
            name.starts_with("screenshot-")
                && path
                    .extension()
                    .is_some_and(|ext| ext.eq_ignore_ascii_case("png")),
            "unexpected auto path: {name}"
        );
    }

    #[test]
    fn strip_data_url_prefix() {
        let url = format!("{PNG_DATA_URL_PREFIX}abc123");
        let b64 = url.strip_prefix(PNG_DATA_URL_PREFIX);
        assert_eq!(b64, Some("abc123"));
    }

    #[test]
    fn strip_data_url_prefix_mismatch() {
        let url = "data:image/jpeg;base64,abc";
        assert!(url.strip_prefix(PNG_DATA_URL_PREFIX).is_none());
    }

    #[test]
    fn build_js_default_uses_inner_height() {
        let js = build_screenshot_js(None);
        assert!(js.contains("window.innerHeight"));
        assert!(!js.contains("scrollHeight"));
    }

    #[test]
    fn build_js_full_page_uses_scroll_height() {
        let js = build_screenshot_js(Some(HeightOverride::FullPage));
        assert!(js.contains("scrollHeight"));
        assert!(js.contains("scrollingElement"));
    }

    #[test]
    fn build_js_explicit_height_inlines_value() {
        let js = build_screenshot_js(Some(HeightOverride::Explicit(4321)));
        assert!(
            js.contains("var h = 4321;"),
            "expected explicit height to be inlined, got: {js}"
        );
    }
}