26use std::collections::hash_map::DefaultHasher;
27use std::hash::{Hash, Hasher};
28use std::path::PathBuf;
29use std::time::{Duration, Instant};
30use std::{env, fs, process};

32use console::{Console, Saves};
33use replay::format::{FrameEntry, apply_wram_diff, diff_wram};
34use replay::keyframes::{self, KeyframeIndex, Tier};
35use replay::log::LogWriter;
36use zbanks::{Bot, States, apply_pad};

44const FPS: f64 = 60.0988;

46struct Dir(PathBuf);
47
48impl States for Dir {
49    fn load(&mut self, name: &str) -> Option<Vec<u8>> {
50        fs::read(self.0.join(format!("{name}.state"))).ok()
51    }
52    fn save(&mut self, name: &str, snapshot: &[u8]) -> bool {
53        fs::write(self.0.join(format!("{name}.state")), snapshot).is_ok()
54    }
55}

62fn hash_state(wram: &[u8], frame_rgba: &[u8]) -> u64 {
63    let mut h = DefaultHasher::new();
64    wram.hash(&mut h);
65    frame_rgba.hash(&mut h);
66    h.finish()
67}

74fn timed_replay(rom: Vec<u8>, restore_from: &[u8], entries: &[FrameEntry]) -> Duration {
75    let mut console = Console::boot(rom, Saves::default()).expect("booting for replay");
76    console.restore(restore_from).expect("restoring for replay");
77    let started = Instant::now();
78    for entry in entries {
79        apply_wram_diff(console.wram_mut(), &entry.wram_diff);
80        apply_pad(&mut console, entry.console_pad);
81        console.run_frame().expect("replay frame");
82        console.samples.0.clear();
83    }
84    started.elapsed()
85}

87fn main() {
88    let args: Vec<String> = env::args().skip(1).collect();
89    let mut rom_path = None;
90    let mut warmup = 3000u32;
91    let mut frames = 10_000u32;
92    let mut i = 0;
93    while i < args.len() {
94        match args[i].as_str() {
95            "--warmup" => {
96                warmup = args[i + 1].parse().expect("--warmup N");
97                i += 1;
98            }
99            "--frames" => {
100                frames = args[i + 1].parse().expect("--frames N");
101                i += 1;
102            }
103            a => rom_path = Some(a.to_string()),
104        }
105        i += 1;
106    }
107    let Some(rom_path) = rom_path else {
108        eprintln!("usage: replay-probe <rom.sfc> [--warmup N] [--frames N]");
109        process::exit(2);
110    };
111    // The seek cases sit in the third coarse span, so every tier's deepest
112    // delta chain exists by then.
113    let coarse_min = keyframes::COARSE_FRAMES as u32 * 3;
114    if frames < coarse_min {
115        eprintln!("--frames must be at least {coarse_min} (3 x COARSE_FRAMES) so every keyframe tier's worst case exists");
116        process::exit(2);
117    }
118
119    let mut rom = fs::read(&rom_path).expect("reading rom");
120    zbanks::rando::patch_rom(&mut rom).expect("patching the rom for the bot");
121    let states_dir = PathBuf::from(format!("{}.states", rom_path.trim_end_matches(".sfc")));
122    let mut states = Dir(states_dir);
123
124    let mut console = Console::boot(rom.clone(), Saves::default()).expect("booting");
125    let home = states.load("home").unwrap_or_else(|| panic!("no home.state; make one with apps/zbanks --make-home first"));
126    console.restore(&home).expect("restoring home");
127
128    let out = PathBuf::from("run/replay-probe");
129    let _ = fs::remove_dir_all(&out);
130    fs::create_dir_all(&out).expect("creating out dir");
131    let mut bot = Bot::start(&mut console, &rom, &mut states, &out);
132    eprintln!("bot started from home.state; state calls: {:?}", bot.state_log());
133
134    // Warm up with the real bot so the recorded window is actual gameplay,
135    // not the first frames out of a savestate.
136    for _ in 0..warmup {
137        let pad = bot.tick(&mut console, &mut states, 0);
138        apply_pad(&mut console, pad);
139        console.run_frame().expect("warmup frame");
140    }
141    eprintln!("warmed up {warmup} frames");
142
143    // Frame 0 of the recording: every replay below restores this first.
144    let header_snapshot = console.snapshot().expect("header snapshot");
145
146    // The reference run: the real bot, ticking for real, for `frames`
147    // frames. Every frame's console pad and WRAM diff go in `entries` (what
148    // a real recording would write to its log); every frame's hash goes in
149    // `reference_hashes`, the thing the two headless replays below must
150    // reproduce exactly. A keyframe is taken at every `DENSE_FRAMES` -
151    // `replay::recorder::Recorder`'s own write cadence - so we get real
152    // dense-tier snapshots to measure, at the frame numbers a real branch
153    // directory would have them.
154    let mut entries = Vec::with_capacity(frames as usize);
155    let mut reference_hashes = Vec::with_capacity(frames as usize);
156    let mut dense_snapshots: Vec<(u64, Vec<u8>)> = Vec::new();
157    let started = Instant::now();
158    for n in 0..frames {
159        let before = *console.wram();
160        let pad = bot.tick(&mut console, &mut states, 0);
161        let after = *console.wram();
162        let wram_diff = diff_wram(&before, &after);
163        apply_pad(&mut console, pad);
164        console.run_frame().expect("reference frame");
165        reference_hashes.push(hash_state(console.wram(), console.frame.rgba()));
166        entries.push(FrameEntry { console_pad: pad, human_pad: 0, wram_diff, decision: None });
167        let frame = u64::from(n + 1);
168        if frame % keyframes::DENSE_FRAMES == 0 {
169            dense_snapshots.push((frame, console.snapshot_fixed().expect("keyframe snapshot")));
170        }
171    }
172    let record_elapsed = started.elapsed();
173    eprintln!(
174        "recorded {frames} frames with the real bot in {:.1}s ({:.0} frames/s)",
175        record_elapsed.as_secs_f64(),
176        f64::from(frames) / record_elapsed.as_secs_f64()
177    );
178
179    // ---- log bytes/hour (unchanged shape from before tiering) ----
180    let log_path = out.join("measure-log.bin");
181    let mut writer = LogWriter::create_or_append(&log_path).expect("log writer");
182    for e in &entries {
183        writer.append(e).expect("append");
184    }
185    writer.flush().expect("flush");
186    let log_bytes = fs::metadata(&log_path).expect("log metadata").len();
187    let frames_per_hour = FPS * 3600.0;
188    let log_bytes_per_hour = log_bytes as f64 / f64::from(frames) * frames_per_hour;
189    eprintln!(
190        "log: {log_bytes} bytes for {frames} frames -> {:.0} bytes/hour ({:.2} MB/hour)",
191        log_bytes_per_hour,
192        log_bytes_per_hour / 1_000_000.0
193    );
194
195    // ---- keyframe sizes per tier, as the recorder writes them ----
196    // `Console::snapshot_fixed` (what `replay::recorder` stores), through the
197    // real KeyframeIndex, so each tier deltas against the reference it
198    // really gets: dense against the keyframe 2 s back, medium against 10 s
199    // back, coarse against nothing.
200    let kf_index = KeyframeIndex::open(&out.join("measure-keyframes")).expect("keyframe index");
201    // Timed too: the live recorder does this on the window's UI thread
202    // (decode the reference's chain, XOR, zstd, write), once per keyframe.
203    let mut slowest_store = Duration::ZERO;
204    let store_started = Instant::now();
205    for (frame, snap) in &dense_snapshots {
206        let t = Instant::now();
207        kf_index.store(*frame, snap).expect("storing keyframe");
208        slowest_store = slowest_store.max(t.elapsed());
209    }
210    eprintln!(
211        "keyframe store (on the UI thread when live): avg {:.1} ms, slowest {:.1} ms",
212        store_started.elapsed().as_secs_f64() * 1000.0 / dense_snapshots.len().max(1) as f64,
213        slowest_store.as_secs_f64() * 1000.0
214    );
215    let raw_len = dense_snapshots.first().map_or(0, |(_, s)| s.len());
216    let distinct_lengths = {
217        let mut l: Vec<usize> = dense_snapshots.iter().map(|(_, s)| s.len()).collect();
218        l.sort_unstable();
219        l.dedup();
220        l.len()
221    };
222    let tier_sizes = |tier: Tier| -> Vec<u64> {
223        dense_snapshots
224            .iter()
225            .filter(|(f, _)| keyframes::tier_of(*f) == tier)
226            .map(|(f, _)| fs::metadata(kf_index.file_path(*f)).expect("keyframe file").len())
227            .collect()
228    };
229    let avg = |v: &[u64]| v.iter().sum::<u64>() as f64 / v.len().max(1) as f64;
230    let (dense, medium, coarse) = (tier_sizes(Tier::Dense), tier_sizes(Tier::Medium), tier_sizes(Tier::Coarse));
231    eprintln!("snapshot_fixed: {raw_len} bytes raw, {distinct_lengths} distinct length(s) across {} snapshots", dense_snapshots.len());
232    for (name, v) in [("dense (2 s delta)", &dense), ("medium (10 s delta)", &medium), ("coarse (base)", &coarse)] {
233        eprintln!(
234            "  {name}: {} measured, avg {:.1} KB, max {:.1} KB",
235            v.len(),
236            avg(v) / 1000.0,
237            v.iter().copied().max().unwrap_or(0) as f64 / 1000.0
238        );
239    }
240
241    // Disk for an hour of play once the recorder runs tiered: the dense
242    // window (~5 min) and the medium window (~1 h) are a standing pool that
243    // prune() keeps at a fixed size; coarse keyframes and the log grow.
244    // Counts per window follow the tier arithmetic: of every 5 dense slots
245    // 4 are dense-only, of every 3 medium slots 2 are medium-only.
246    let dense_pool = (keyframes::DENSE_RETAIN_FRAMES / keyframes::DENSE_FRAMES) as f64 * 0.8 * avg(&dense);
247    let medium_pool = (keyframes::MEDIUM_RETAIN_FRAMES / keyframes::MEDIUM_FRAMES) as f64 * (2.0 / 3.0) * avg(&medium);
248    let coarse_per_hour = frames_per_hour / keyframes::COARSE_FRAMES as f64 * avg(&coarse);
249    eprintln!(
250        "standing pool: dense {:.1} MB + medium {:.1} MB = {:.1} MB; growth per hour: coarse {:.1} MB + log {:.1} MB",
251        dense_pool / 1e6,
252        medium_pool / 1e6,
253        (dense_pool + medium_pool) / 1e6,
254        coarse_per_hour / 1e6,
255        log_bytes_per_hour / 1e6
256    );
257    eprintln!(
258        "first hour of play, total on disk: {:.1} MB",
259        (dense_pool + medium_pool + coarse_per_hour + log_bytes_per_hour) / 1e6
260    );
261
262    // ---- real headless replay throughput ----
263    let throughput_slice = &entries[..(keyframes::COARSE_FRAMES as usize).min(entries.len())];
264    let throughput_elapsed = timed_replay(rom.clone(), &header_snapshot, throughput_slice);
265    let replay_fps = throughput_slice.len() as f64 / throughput_elapsed.as_secs_f64();
266    eprintln!(
267        "headless replay throughput: {} frames in {:.3}s -> {:.0} frames/s ({:.1}x real time)",
268        throughput_slice.len(),
269        throughput_elapsed.as_secs_f64(),
270        replay_fps,
271        replay_fps / FPS
272    );
273
274    // ---- worst-case seek per tier, the real path ----
275    // What `replay::seek::seek_to` does: boot, load the keyframe off disk
276    // (decoding its whole delta chain), `restore_fixed`, replay the log to
277    // one frame short of the next keyframe of that tier. Keyframes chosen
278    // with the deepest chain each tier can have, in the third coarse span.
279    let base = keyframes::COARSE_FRAMES * 2;
280    let seek_case = |label: &str, kf: u64, gap: u64| {
281        let started = Instant::now();
282        let mut c = Console::boot(rom.clone(), Saves::default()).expect("booting for seek");
283        let bytes = kf_index.load(kf).expect("loading keyframe");
284        let loaded = started.elapsed();
285        c.restore_fixed(&bytes).expect("restore_fixed");
286        for entry in &entries[kf as usize..(kf + gap - 1) as usize] {
287            apply_wram_diff(c.wram_mut(), &entry.wram_diff);
288            apply_pad(&mut c, entry.console_pad);
289            c.run_frame().expect("seek frame");
290            c.samples.0.clear();
291        }
292        let elapsed = started.elapsed();
293        eprintln!(
294            "seek worst case, {label}: keyframe {kf}, {} frames replayed, {:.3}s (boot + keyframe load {:.3}s)",
295            gap - 1,
296            elapsed.as_secs_f64(),
297            loaded.as_secs_f64()
298        );
299        elapsed
300    };
301    let dense_kf = base + keyframes::MEDIUM_FRAMES + keyframes::MEDIUM_FRAMES - keyframes::DENSE_FRAMES;
302    let medium_kf = base + keyframes::COARSE_FRAMES - keyframes::MEDIUM_FRAMES;
303    let dense_worst = seek_case("last 5 min (dense, 2 s)", dense_kf, keyframes::DENSE_FRAMES);
304    let medium_worst = seek_case("last hour (medium, 10 s)", medium_kf, keyframes::MEDIUM_FRAMES);
305    let coarse_worst = seek_case("older (coarse, 30 s)", base, keyframes::COARSE_FRAMES);
306    eprintln!(
307        "worst-case seek: last 5 min {:.3}s, last hour {:.3}s, older {:.3}s",
308        dense_worst.as_secs_f64(),
309        medium_worst.as_secs_f64(),
310        coarse_worst.as_secs_f64()
311    );
312
313    // ---- the determinism probe itself ----
314    let mut mismatch: Option<(u32, &'static str)> = None;
315    let mut replay_hashes_1 = Vec::with_capacity(frames as usize);
316    for pass in 1..=2u32 {
317        let mut replay_console = Console::boot(rom.clone(), Saves::default()).expect("booting for replay");
318        replay_console.restore(&header_snapshot).expect("restoring header for replay");
319        let seek_started = Instant::now();
320        for (n, entry) in entries.iter().enumerate() {
321            apply_wram_diff(replay_console.wram_mut(), &entry.wram_diff);
322            apply_pad(&mut replay_console, entry.console_pad);
323            replay_console.run_frame().expect("replay frame");
324            replay_console.samples.0.clear();
325            let hash = hash_state(replay_console.wram(), replay_console.frame.rgba());
326            if pass == 1 {
327                replay_hashes_1.push(hash);
328            }
329            let baseline = if pass == 1 { reference_hashes[n] } else { replay_hashes_1[n] };
330            if hash != baseline && mismatch.is_none() {
331                mismatch = Some((n as u32, if pass == 1 { "replay 1 vs the original recording" } else { "replay 2 vs replay 1" }));
332            }
333        }
334        let elapsed = seek_started.elapsed();
335        eprintln!(
336            "determinism replay pass {pass}: {frames} frames in {:.3}s ({:.0} frames/s, {:.1}x real time)",
337            elapsed.as_secs_f64(),
338            f64::from(frames) / elapsed.as_secs_f64(),
339            f64::from(frames) / FPS / elapsed.as_secs_f64()
340        );
341    }
342
343    match mismatch {
344        None => {
345            println!(
346                "DETERMINISTIC: {frames} frames, two independent replay passes, both byte-identical \
347                 to the original recording (work RAM + frame hash compared every frame)."
348            );
349        }
350        Some((frame, which)) => {
351            println!("NOT DETERMINISTIC: first mismatch at local frame {frame} ({which}).");
352            process::exit(1);
353        }
354    }
355}