// Copyright (C) 2026 Danilo M. <danix@danix.xyz>
// SPDX-License-Identifier: GPL-2.0-only
//
// Upload an image, GIF or video to the AJAZZ AK820 Pro 128x128 screen over
// Linux hidraw. The protocol below was taken from a USB capture of the
// official Windows software (v1.0.0.5) uploading a 3-frame GIF.

use std::fs::{File, OpenOptions};
use std::io::{Read, Write};
use std::os::unix::io::AsRawFd;
use std::process::{Command, exit};
use std::thread::sleep;
use std::time::Duration;

const HID_ID: &str = "HID_ID=0003:00000C45:00008009";
const CONTROL_IFACE: &str = "/input3";
const DATA_IFACE: &str = "/input2";

const FRAME_BYTES: usize = 128 * 128 * 2; // RGB565 little-endian
const CHUNK: usize = 4096; // sent as 64 full USB packets, no short packet
const HEADER: usize = 256; // frame count + one delay byte per frame, 0xFF padded
const MAX_FRAMES: usize = 255; // frame count is a u8 in the header
// Captured: 50 ms in the app -> 0x19, 250 ms -> 0x7d. Max ~510 ms per frame.
const DELAY_UNIT_MS: u32 = 2;

const CMD_START: u8 = 0x18;
const CMD_IMAGE: u8 = 0x72;
const CMD_SAVE: u8 = 0x02;

unsafe extern "C" {
    fn ioctl(fd: i32, req: u64, ...) -> i32;
    fn poll(fds: *mut PollFd, n: u64, timeout_ms: i32) -> i32;
}

#[repr(C)]
struct PollFd {
    fd: i32,
    events: i16,
    revents: i16,
}

// HIDIOC[SG]FEATURE(len) = _IOC(_IOC_READ|_IOC_WRITE, 'H', 0x06/0x07, len)
fn hidioc(nr: u64, len: usize) -> u64 {
    (3 << 30) | ((len as u64) << 16) | ((b'H' as u64) << 8) | nr
}

fn die(msg: impl std::fmt::Display) -> ! {
    eprintln!("ak820-upload: {msg}");
    exit(1)
}

/// Find the keyboard's hidraw node for one USB interface.
fn open_iface(iface: &str) -> File {
    let dir = std::fs::read_dir("/sys/class/hidraw").unwrap_or_else(|e| die(e));
    for entry in dir.flatten() {
        let uevent = std::fs::read_to_string(entry.path().join("device/uevent")).unwrap_or_default();
        if uevent.contains(HID_ID) && uevent.lines().any(|l| l.starts_with("HID_PHYS=") && l.ends_with(iface)) {
            let dev = format!("/dev/{}", entry.file_name().to_string_lossy());
            return OpenOptions::new().read(true).write(true).open(&dev)
                .unwrap_or_else(|e| die(format!("{dev}: {e}")));
        }
    }
    die("AK820 Pro not found. Connect it with the USB-C cable in wired mode.")
}

/// Feature report buffer: report ID 0, then the 64-byte packet
/// [0x04, cmd, sub, 0, 0, 0, 0, 0, lo, hi, 0...].
fn control(cmd: u8, sub: u8, arg: u16) -> [u8; 65] {
    let mut p = [0u8; 65];
    p[1] = 0x04;
    p[2] = cmd;
    p[3] = sub;
    p[9] = arg as u8;
    p[10] = (arg >> 8) as u8;
    p
}

/// SET_REPORT, then the 64-byte GET_REPORT read-back the app always follows
/// it with. Without the read-back the keyboard ignores the upload.
fn send_control(ctl: &File, pkt: [u8; 65]) {
    let fd = ctl.as_raw_fd();
    if unsafe { ioctl(fd, hidioc(0x06, pkt.len()), pkt.as_ptr()) } < 0 {
        die(format!("control packet {:#04x}: {}", pkt[2], std::io::Error::last_os_error()));
    }
    sleep(Duration::from_millis(40));
    let mut buf = [0u8; 65]; // buf[0] = report ID 0
    if unsafe { ioctl(fd, hidioc(0x07, buf.len()), buf.as_mut_ptr()) } < 0 {
        die(format!("read-back after {:#04x}: {}", pkt[2], std::io::Error::last_os_error()));
    }
    sleep(Duration::from_millis(40));
}

/// Wait for the keyboard's per-chunk ack (an input report on the control
/// endpoint 0x84 of the data interface, `01 5a 02 ...`).
fn wait_ack(dat: &mut File, chunk: usize) {
    let mut pfd = PollFd { fd: dat.as_raw_fd(), events: 1, revents: 0 };
    if unsafe { poll(&mut pfd, 1, 2000) } <= 0 {
        die(format!("no ack from the keyboard for chunk {chunk}"));
    }
    let mut buf = [0u8; 64];
    let _ = dat.read(&mut buf);
}

/// Header + frames, padded with 0xFF to whole chunks.
fn payload(frames: &[u8], delay_ms: u32) -> Vec<u8> {
    let n = frames.len() / FRAME_BYTES;
    let mut out = vec![0xffu8; HEADER];
    out[0] = n as u8;
    out[1..=n].fill((delay_ms / DELAY_UNIT_MS).clamp(1, 255) as u8);
    out.extend_from_slice(frames);
    out.resize(out.len().div_ceil(CHUNK) * CHUNK, 0xff);
    out
}

/// Run ffmpeg on `input` with the keyboard's crop/scale/fps/frame cap, then
/// `extra` (filter suffix) and `out_args` (output options and target).
fn ffmpeg(input: &str, fps: u32, extra: &str, out_args: &[&str]) -> Vec<u8> {
    let vf = format!("fps={fps},crop='min(iw,ih)':'min(iw,ih)',scale=128:128:flags=lanczos{extra}");
    let out = Command::new("ffmpeg")
        .args(["-v", "error", "-y", "-i", input, "-filter_complex", &vf, "-frames:v"])
        .arg(MAX_FRAMES.to_string())
        .args(out_args)
        .output()
        .unwrap_or_else(|e| die(format!("ffmpeg: {e}")));
    if !out.status.success() {
        die(format!("ffmpeg failed on {input}: {}", String::from_utf8_lossy(&out.stderr).trim()));
    }
    out.stdout
}

/// Decode any ffmpeg-readable input to raw 128x128 RGB565LE frames.
fn decode(input: &str, fps: u32) -> Vec<u8> {
    let raw = ffmpeg(input, fps, "", &["-f", "rawvideo", "-pix_fmt", "rgb565le", "-"]);
    if raw.len() < FRAME_BYTES {
        die(format!("{input}: no frames decoded"));
    }
    raw
}

fn main() {
    let mut args: Vec<String> = std::env::args().skip(1).collect();
    let gif = args.iter().position(|a| a == "--gif").map(|i| {
        args.remove(i);
        if i < args.len() { args.remove(i) } else { die("--gif needs an output file") }
    });
    let (input, fps) = match args.as_slice() {
        [i] => (i.as_str(), 15),
        [i, f] => (i.as_str(), f.parse().ok().filter(|f| (2..=50).contains(f))
            .unwrap_or_else(|| die("fps must be 2..50 (max frame time is ~510 ms)"))),
        _ => die("usage: ak820-upload [--gif out.gif] <image|gif|video> [fps, default 15]"),
    };

    // --gif: write what would be uploaded as a GIF instead of uploading.
    if let Some(out) = gif {
        ffmpeg(input, fps, ",split[a][b];[a]palettegen=stats_mode=diff[p];[b][p]paletteuse=dither=bayer:bayer_scale=4",
            &["-loop", "0", &out]);
        eprintln!("wrote {out}");
        return;
    }

    let frames = decode(input, fps);
    let n = frames.len() / FRAME_BYTES;
    let data = payload(&frames, 1000 / fps);
    let chunks = data.len() / CHUNK;

    let ctl = open_iface(CONTROL_IFACE);
    let mut dat = open_iface(DATA_IFACE);
    eprintln!("uploading {n} frame(s), {chunks} chunks");

    // START -> IMAGE sub 2 (chunk count) -> chunks -> SAVE.
    send_control(&ctl, control(CMD_START, 0, 0));
    send_control(&ctl, control(CMD_IMAGE, 2, chunks as u16));

    let mut report = vec![0u8; 1 + CHUNK]; // leading 0 = unnumbered output report
    for (i, chunk) in data.chunks(CHUNK).enumerate() {
        report[1..].copy_from_slice(chunk);
        dat.write_all(&report).unwrap_or_else(|e| die(format!("chunk {i}: {e}")));
        wait_ack(&mut dat, i);
        eprint!("\r{}/{chunks}", i + 1);
    }
    eprintln!();

    send_control(&ctl, control(CMD_SAVE, 0, 0));
    eprintln!("done");
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn layout() {
        // Bytes as captured from the official app.
        let c = control(CMD_IMAGE, 2, 0x19);
        assert_eq!(&c[..11], &[0, 0x04, 0x72, 2, 0, 0, 0, 0, 0, 0x19, 0]);

        // 3 frames at 50 ms: 25 chunks, header 03 19 19 19 ff.
        let three = payload(&vec![0u8; 3 * FRAME_BYTES], 50);
        assert_eq!(three.len(), 25 * CHUNK);
        assert_eq!(&three[..5], &[3, 0x19, 0x19, 0x19, 0xff]);
        assert_eq!(three[HEADER], 0);
        assert_eq!(three[HEADER + 3 * FRAME_BYTES], 0xff);

        let one = payload(&vec![0u8; FRAME_BYTES], 66);
        assert_eq!(&one[..3], &[1, 33, 0xff]);
        assert_eq!(one.len(), 9 * CHUNK);
    }
}