use std::env;
use std::fs::{self, File};
use std::io::BufWriter;
use std::path::Path;
use std::time::{Duration, Instant};

use ericrfb::handshake::Config;
use ericrfb::session::{ActiveSession, Event};

fn save_png(fb: &ericrfb::framebuffer::Framebuffer, path: &Path) {
    let rgba = fb.to_rgba();
    let file = File::create(path).expect("cannot create PNG file");
    let bw = BufWriter::new(file);
    let mut encoder = png::Encoder::new(bw, fb.width as u32, fb.height as u32);
    encoder.set_color(png::ColorType::Rgba);
    encoder.set_depth(png::BitDepth::Eight);
    let mut writer = encoder.write_header().expect("png header failed");
    writer.write_image_data(&rgba).expect("png write failed");
}

fn main() {
    let args: Vec<String> = env::args().collect();

    let host = args
        .iter()
        .position(|a| a == "--host")
        .and_then(|i| args.get(i + 1))
        .expect("usage: --host <ip> --applet-id <token> [--port <port>] [--duration <secs>]");

    let applet_id = args
        .iter()
        .position(|a| a == "--applet-id")
        .and_then(|i| args.get(i + 1))
        .expect("usage: --host <ip> --applet-id <token> [--port <port>] [--duration <secs>]");

    let port: u16 = args
        .iter()
        .position(|a| a == "--port")
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(443);

    let duration_secs: u64 = args
        .iter()
        .position(|a| a == "--duration")
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(30);

    let out_dir = Path::new("out");
    fs::create_dir_all(out_dir).expect("cannot create out/");

    let cfg = Config::new(host, port, applet_id);
    println!("Connecting to {}:{}...", cfg.host, cfg.port);

    // Request Hextile (5), CopyRect (1), Raw (0)
    let mut session = ActiveSession::connect(&cfg, &[5, 1, 0]).expect("connect failed");
    println!(
        "Connected: {}x{}, recording for {duration_secs}s...",
        session.framebuffer.width, session.framebuffer.height
    );

    let start = Instant::now();
    let duration = Duration::from_secs(duration_secs);
    let mut last_save = Instant::now() - Duration::from_secs(2); // force first save
    let mut frame_count = 0u32;

    loop {
        if start.elapsed() >= duration {
            break;
        }

        match session.process_one() {
            Ok(Some(Event::FramebufferDirty)) => {
                // Save at most 1 PNG per second
                if last_save.elapsed() >= Duration::from_secs(1) {
                    let path = out_dir.join(format!("frame_{frame_count:04}.png"));
                    save_png(&session.framebuffer, &path);
                    println!(
                        "[{:.1}s] saved {}",
                        start.elapsed().as_secs_f64(),
                        path.display()
                    );
                    frame_count += 1;
                    last_save = Instant::now();
                }
                // Request next update
                if let Err(e) = session.request_update() {
                    eprintln!("Error requesting update: {e}");
                    break;
                }
            }
            Ok(Some(Event::Resize { width, height })) => {
                println!("Resized to {width}x{height}");
            }
            Ok(Some(Event::Debug(s))) => {
                eprintln!("[debug] {s}");
            }
            Ok(Some(_)) | Ok(None) => {}
            Err(e) => {
                eprintln!("Error: {e}");
                break;
            }
        }
    }

    println!("Done. Saved {frame_count} frames to {}/", out_dir.display());
}