16use std::collections::BTreeSet;
17use std::fs;
18use std::io::Write;
19use std::path::{Path, PathBuf};
20use std::process;

22use console::{Console, Saves};
23use zbanks::{Bot, States, apply_pad, recovery, stall};
24
25struct Dir(PathBuf);
26
27impl States for Dir {
28    fn load(&mut self, name: &str) -> Option<Vec<u8>> {
29        fs::read(self.0.join(format!("{name}.state"))).ok()
30    }
31    fn save(&mut self, name: &str, snapshot: &[u8]) -> bool {
32        fs::write(self.0.join(format!("{name}.state")), snapshot).is_ok()
33    }
34}
35
36fn png(path: &Path, console: &Console) {
37    let (w, h) = console.frame.size;
38    let file = fs::File::create(path).expect("creating png");
39    let mut encoder = png::Encoder::new(std::io::BufWriter::new(file), w, h);
40    encoder.set_color(png::ColorType::Rgba);
41    encoder.set_depth(png::BitDepth::Eight);
42    let mut writer = encoder.write_header().expect("png header");
43    writer.write_image_data(console.frame.rgba()).expect("png data");
44}
45
46fn main() {
47    let args: Vec<String> = std::env::args().skip(1).collect();
48    let mut positional = Vec::new();
49    let mut states_dir = None;
50    let mut start = None;
51    let mut every = 600u32;
52    let mut make_home: Option<String> = None;
53    let mut import_map: Option<PathBuf> = None;
54    // Jev at the goal choice: off unless asked for (`--jev` or ZB_JEV=1).
55    let mut jev_on = std::env::var("ZB_JEV").is_ok_and(|v| v == "1");
56    let mut jev_dollars: Option<f64> = None;
57    // `--no-retry`: goals the bot gives up on stay given up, as upstream's.
58    let mut retry = true;
59    // `--save-at N`: keep the machine after frame N as `<out-dir>/at-N.state`,
60    // to start a later run there (as its states directory's `home`).
61    let mut save_at: Vec<u32> = Vec::new();
62    let mut i = 0;
63    while i < args.len() {
64        match args[i].as_str() {
65            "--states" => { states_dir = Some(PathBuf::from(&args[i + 1])); i += 1; }
66            "--start" => { start = Some(args[i + 1].clone()); i += 1; }
67            "--every" => { every = args[i + 1].parse().expect("--every N"); i += 1; }
68            "--import-map" => { import_map = Some(PathBuf::from(&args[i + 1])); i += 1; }
69            "--make-home" => { make_home = Some(args[i + 1].clone()); i += 1; }
70            "--jev" => jev_on = true,
71            "--no-retry" => retry = false,
72            "--save-at" => { save_at.push(args[i + 1].parse().expect("--save-at FRAME")); i += 1; }
73            "--jev-dollars" => { jev_dollars = Some(args[i + 1].parse().expect("--jev-dollars USD")); i += 1; }
74            a => positional.push(a.to_string()),
75        }
76        i += 1;
77    }
78    let [rom_path, frames, out] = positional.as_slice() else {
79        eprintln!("usage: zbanks <rom.sfc> <frames> <out-dir> [--states DIR] [--start NAME] [--every N] [--import-map FILE] [--make-home NAME] [--jev] [--jev-dollars USD] [--no-retry] [--save-at FRAME]...");
80        process::exit(2);
81    };
82    let frames: u32 = frames.parse().expect("frames is a number");
83    let out = PathBuf::from(out);
84    fs::create_dir_all(&out).expect("creating out dir");
85    let mut rom = fs::read(rom_path).expect("reading rom");
86    // The two ROM differences the bot depends on (zbanks::rando).
87    zbanks::rando::patch_rom(&mut rom).expect("patching the ROM");
88    let mut states = Dir(states_dir.unwrap_or_else(|| PathBuf::from(format!("{}.states", rom_path.trim_end_matches(".sfc")))));
89
90    let mut console = Console::boot(rom.clone(), Saves::default()).expect("booting");
91    if let Some(name) = &start {
92        let snapshot = states.load(name).unwrap_or_else(|| panic!("no state {name}"));
93        console.restore(&snapshot).expect("restoring start state");
94    }
95    if let Some(name) = &make_home {
96        make_home_state(&mut console, &rom, &mut states, name, &out);
97        return;
98    }
99    // `--import-map FILE`: the map the bot's first tick imports, under the
100    // name upstream's code asks for (alttp.c:34).
101    if let Some(map) = &import_map {
102        fs::copy(map, out.join("map.19.txt")).expect("copying the map to import");
103    }
104    let bot_dir = fs::canonicalize(&out).expect("out dir");
105    let mut bot = Bot::start(&mut console, &rom, &mut states, &bot_dir);
106    bot.retry_given_up = retry;
107    eprintln!("bot started; state calls: {:?}; retrying given-up goals: {retry}", bot.state_log());
108    let jev = jev_on.then(|| {
109        let client = jev_http::Jev::from_env("zbanks", decisions::model())
110            .expect("--jev needs TYPESAFE_API_KEY (run under op-env-run)")
111            .expect("opening the Jev client");
112        let config = decisions::goal_choice::Config { run_dollars: jev_dollars, ..decisions::goal_choice::Config::default() };
113        let mut chooser = decisions::goal_choice::Chooser::new(client, config).expect("jev chooser");
114        chooser.log_to(&out.join("jev.jsonl")).expect("jev log");
115        let shared = decisions::goal_choice::Shared::new(chooser);
116        bot.set_goal_chooser(Some(Box::new(shared.clone())), config.margin);
117        eprintln!("jev: on, margin {}, run limit {:?}", config.margin, jev_dollars);
118        shared
119    });
120
121    let mut progress = fs::File::create(out.join("progress.tsv")).expect("progress.tsv");
122    writeln!(progress, "frame\tmodule\tsub\tindoors\troom\tarea\tlink_x\tlink_y\thp\tmax_hp\tsword\tbot_link\ttraps\tstray\tinfo\ttask\tmanual\tpendants\tcrystals\tbow\tbook\tboots\tgloves\tgiven_up\tput_back").unwrap();
123    let mut places = BTreeSet::new();
124    // ZB_TRACE=first-last prints the pad and a few bytes for those frames.
125    let trace = std::env::var("ZB_TRACE").ok().and_then(|t| {
126        let (a, b) = t.split_once('-')?;
127        Some(a.parse::<u32>().ok()?..=b.parse::<u32>().ok()?)
128    }).unwrap_or(1..=0);
129    // ZB_HUMAN=first-last:PAD holds PAD (Snes9x layout, hex) as the human's
130    // pad on those frames. With X (0x0040) in it the bot steps aside
131    // (alttp.c:57), so this tests what Link can do from a state.
132    let human = std::env::var("ZB_HUMAN").ok().and_then(|h| {
133        let (frames, pad) = h.split_once(':')?;
134        let (a, b) = frames.split_once('-')?;
135        Some((a.parse::<u32>().ok()?..=b.parse::<u32>().ok()?, u16::from_str_radix(pad.trim_start_matches("0x"), 16).ok()?))
136    });
137    let started = std::time::Instant::now();
138    // The same stall detector and bounded recovery the window uses
139    // (packages/zbanks/src/stall.rs, src/recovery.rs) - a goal that keeps
140    // failing for a reason no possession change will ever fix (the door
141    // `D 0x80` stall, research/zbanks-alttp.md) does not get to sit in
142    // manual mode for the rest of a headless run either. This replaces the
143    // old "stop a minute after manual mode" behaviour: now the run keeps
144    // going, retrying, until recovery itself gives up (or the frame budget
145    // runs out).
146    let mut detector = stall::Detector::new();
147    let mut stalled: Option<stall::Stall> = None;
148    let mut recovery = recovery::Recovery::new();
149    let mut gave_up_since: Option<u32> = None;
150    for n in 0..frames {
151        let held = human.as_ref().filter(|(frames, _)| frames.contains(&n)).map_or(0, |(_, pad)| *pad);
152        if let Some(jev) = &jev {
153            jev.0.borrow_mut().set_snapshot(decisions::goal_choice::Snapshot::read(console.wram()));
154        }
155        let pad = bot.tick(&mut console, &mut states, held);
156        if trace.contains(&n) {
157            let r = console.wram();
158            eprintln!("trace {n}: pad {pad:#06x} $1A {:#04x} module {:#04x} sub {:#04x} subsub {:#04x} link_state {:#04x} x {:#06x} y {:#06x} layer $EE {}", r[0x1A], r[0x10], r[0x11], r[0xB0], r[0x5D], u16::from_le_bytes([r[0x22], r[0x23]]), u16::from_le_bytes([r[0x20], r[0x21]]), r[0xEE]);
159        }
160        apply_pad(&mut console, pad);
161        console.run_frame().expect("running frame");
162        if save_at.contains(&(n + 1)) {
163            fs::write(out.join(format!("at-{}.state", n + 1)), console.snapshot().expect("snapshot")).unwrap();
164        }
165        let ram = console.wram();
166        let indoors = ram[0x1B];
167        let place = if indoors != 0 { 0x1000 | u16::from_le_bytes([ram[0xA0], ram[0xA1]]) } else { u16::from(ram[0x8A]) };
168        places.insert(place);
169
170        let report0 = bot.report(0);
171        let state = alttp::State::read(ram);
172        let signal = stall::Signal {
173            frame: report0.frame,
174            manual_mode: report0.manual_mode,
175            no_task: report0.tasks.is_empty(),
176            eligible: state.control && state.mode.is_play(),
177            position: (state.place, state.link_x, state.link_y),
178        };
179        let was = stalled;
180        stalled = detector.update(signal);
181        if stalled != was {
182            match stalled {
183                Some(s) => eprintln!("STALLED from frame {} ({}), now {n}", s.since, s.reason.words()),
184                None => eprintln!("stall cleared at frame {n}"),
185            }
186        }
187        if let Some(event) = recovery.consider(&mut bot, &console, stalled.is_some(), report0.frame, signal.position) {
188            if event.gave_up {
189                eprintln!(
190                    "RECOVERY gave up after {} recoveries with no progress (frame {}, restored {} goal(s) on the last try)",
191                    event.attempt, event.frame, event.restored
192                );
193            } else {
194                eprintln!(
195                    "RECOVERY attempt {} of {}: restored {} goal(s) at frame {}",
196                    event.attempt, recovery::CAP, event.restored, event.frame
197                );
198            }
199        }
200        // `recovery.status().gave_up` is a lifetime latch (recovery.rs: "for
201        // the life of this Recovery... giving up is for good, not just for
202        // this backoff") - but that only means Recovery itself will not
203        // retry any MORE given-up goals. It does not mean the bot is doomed:
204        // `zb_goal_add`'s own "a goal was added while in manual mode;
205        // resuming the plan" clears `ap_manual_mode` on ANY fresh goal,
206        // recovery-driven or not, so an unrelated goal (a new EXPLORE tile
207        // reached, an item picked up) can clear a stall Recovery had nothing
208        // to do with. Measured 2026-09-22 (`run/zbanks-c/task1-fix.log`):
209        // recovery gave up at frame 186826 with the bot reporting stalled,
210        // and the very next frame's own log line is "stall cleared at frame
211        // 186827" - through this exact mechanism, unrelated to Recovery.
212        // Stopping unconditionally 3600 frames after Recovery's cap once
213        // fired (the old rule below) ended that run at 190426 while the bot
214        // had been happily exploring, un-stalled, for over 3,500 frames -
215        // discarding real remaining budget on a goal that had already moved
216        // on. So the grace-period timer must track "stalled AND recovery
217        // exhausted" together, and reset the moment either stops being true.
218        if recovery.status().gave_up && stalled.is_some() {
219            if gave_up_since.is_none() {
220                gave_up_since = Some(n);
221            }
222        } else {
223            gave_up_since = None;
224        }
225        if gave_up_since.is_some_and(|g| n >= g + 3600) {
226            eprintln!("stopping at frame {n}: recovery gave up and the bot has stayed stalled for a minute");
227            break;
228        }
229
230        if (n + 1) % every == 0 || n + 1 == frames {
231            let report = bot.report(8);
232            bot.flush_log();
233            let line = format!(
234                "{}\t{:#04x}\t{:#04x}\t{}\t{:#06x}\t{:#04x}\t{:#06x}\t{:#06x}\t{}\t{}\t{}\t{:#06x},{:#06x}\t{}\t{}\t{}\t{}\t{}\t{:#04x}\t{:#04x}\t{}\t{}\t{}\t{}\t{}\t{}",
235                n + 1, ram[0x10], ram[0x11], indoors,
236                u16::from_le_bytes([ram[0xA0], ram[0xA1]]), ram[0x8A],
237                u16::from_le_bytes([ram[0x22], ram[0x23]]), u16::from_le_bytes([ram[0x20], ram[0x21]]),
238                ram[0xF36D], ram[0xF36C], ram[0xF359],
239                report.link.0, report.link.1, report.traps, report.stray_reads,
240                report.info.replace(['\t', '\n'], " "),
241                report.tasks.first().map(String::as_str).unwrap_or("").replace('\t', " "),
242                report.manual_mode,
243                // $7EF374 pendants, $7EF37A crystals, $7EF340 bow, $7EF34E
244                // book, $7EF355 boots, $7EF354 gloves (research/alttp-ram-map.md §3).
245                ram[0xF374], ram[0xF37A], ram[0xF340], ram[0xF34E], ram[0xF355], ram[0xF354],
246                report.given_up, report.retried,
247            );
248            writeln!(progress, "{line}").unwrap();
249            progress.flush().unwrap();
250            png(&out.join(format!("frame-{:06}.png", n + 1)), &console);
251            eprintln!("{line}");
252        }
253    }
254    let report = bot.report(40);
255    eprintln!(
256        "done: {frames} frames in {:.1}s; {} places seen; traps {} (last frame {}, ip anchor+{:#x}); stray reads {}",
257        started.elapsed().as_secs_f64(), places.len(), report.traps, report.last_trap_frame, report.last_trap_ip, report.stray_reads
258    );
259    eprintln!("places: {:x?}", places);
260    eprintln!(
261        "given up: {} waiting; {} goals put back over {} changes of Link's possessions",
262        report.given_up, report.retried, report.retries
263    );
264    let rec = recovery.status();
265    eprintln!(
266        "recovery: {} attempt(s) since the last progress (cap {}), gave up: {}{}",
267        rec.last.map_or(0, |e| e.attempt),
268        recovery::CAP,
269        rec.gave_up,
270        rec.last.map_or(String::new(), |e| format!(" (last at frame {}, restored {})", e.frame, e.restored)),
271    );
272    if let Some(jev) = &jev {
273        let chooser = jev.0.borrow();
274        let t = chooser.engine.totals.clone();
275        let overrides = chooser.overrides;
276        let minutes = f64::from(report.frame + 1) / 3600.0; // frames actually played
277        eprintln!(
278            "jev: {} choices, {} questions, {} reused, {} one-choice (not asked), {} throttled, {} defaults, {} overrides of upstream's pick; {} input tokens ({:.0} per question); ${:.6} over {:.2} game-minutes: {:.2} questions/min, ${:.6}/game-minute",
279            t.choices, t.questions, t.reused, t.no_question, t.throttled, t.defaults, overrides, t.input_tokens,
280            t.input_tokens as f64 / f64::from(t.questions.max(1)), t.dollars, minutes,
281            f64::from(t.questions) / minutes, t.dollars / minutes
282        );
283    }
284    eprintln!("goals ({} total):", report.goal_count);
285    for g in &report.goals {
286        eprintln!("  {} {} @ {} attempts={} score={}", g.kind, g.node, g.screen, g.attempts, zbanks::score_words(g.last_score).map(String::from).unwrap_or(g.last_score.to_string()));
287    }
288    eprintln!("tasks: {:#?}", report.tasks);
289    fs::write(out.join("final.state"), console.snapshot().expect("snapshot")).unwrap();
290    // What the bot learned, in its own format: copy it to `map.19.txt` in a
291    // later run's out-dir and that run's first tick imports it
292    // (alttp.c:31-38), as upstream carried its map from run to run.
293    bot.export_map("map_export.txt");
294    eprintln!("manual mode: {}; map exported to {}", report.manual_mode, out.join("map_export.txt").display());
295}

305fn walk(console: &mut Console, frames: u32, done: impl Fn(&[u8; console::WRAM_LEN]) -> bool) -> Option<u32> {
306    let mut naming_frames = 0u32;
307    for n in 0..frames {
308        let ram = console.wram();
309        if done(ram) {
310            return Some(n);
311        }
312        let (module, sub, link_state) = (ram[0x10], ram[0x11], ram[0x5D]);
313        if n % 1000 == 999 {
314            eprintln!(
315                "walk {n}: module {module:#04x} sub {sub:#04x} room {:#06x} indoors {} link state {link_state:#04x}",
316                u16::from_le_bytes([ram[0xA0], ram[0xA1]]), ram[0x1B]
317            );
318        }
319        let naming = module == 0x04;
320        if naming {
321            naming_frames += 1;
322        }
323        let type_letter = naming && (30..34).contains(&naming_frames);
324        let menu = (module < 0x06 && (!naming || naming_frames > 60)) || (module == 0x0E && sub == 0x01);
325        // Spawn select (0x1B), shown for a file past the rain: A takes the
326        // point under the cursor, Link's house.
327        let text = (module == 0x0E && sub == 0x02) || module == 0x1B;
328        let in_bed = module == 0x07 && sub == 0x00 && link_state == 0x16;
329        console.set_button(console::SnesButton::Start, menu && n % 60 < 4);
330        console.set_button(console::SnesButton::A, (text && n % 20 < 4) || type_letter);
331        console.set_button(console::SnesButton::Right, in_bed && n % 30 < 4);
332        console.run_frame().expect("frame");
333    }
334    None
335}

337fn release_all(console: &mut Console) {
338    apply_pad(console, 0);
339}

354fn make_home_state(console: &mut Console, rom: &[u8], states: &mut Dir, name: &str, out: &Path) {
355    // 1. A vanilla file, created and started; wait a second so the core has
356    //    persisted SRAM (it checks every 30 frames, snes-core api.rs:387-392).
357    walk(console, 20_000, |r| r[0x10] == 0x07).expect("never started a new file");
358    release_all(console);
359    for _ in 0..90 {
360        console.run_frame().expect("frame");
361    }
362    let mut saves = console.saves.clone();
363    let sram = saves.0.get_mut("sav").expect("the cartridge persisted no SRAM");
364    // 2. Open mode.
365    zbanks::rando::preset_file(sram, 0).expect("presetting file 1");
366    fs::write(out.join(format!("{name}.sav")), &saves.0["sav"]).expect("writing sav");
367    let mut fresh = Console::boot(rom.to_vec(), saves).expect("booting with the preset save");
368    // 3. Into the game, until Link stands in his house and can move.
369    let controllable = |r: &[u8; console::WRAM_LEN]| {
370        r[0x10] == 0x07 && r[0x11] == 0x00 && r[0x1B] != 0 && u16::from_le_bytes([r[0xA0], r[0xA1]]) == 0x0104 && r[0x5D] == 0x00
371    };
372    walk(&mut fresh, 20_000, controllable).expect("never reached Link in his house");
373    release_all(&mut fresh);
374    for _ in 0..60 {
375        fresh.run_frame().expect("frame");
376        assert!(controllable(fresh.wram()), "Link stopped being controllable");
377    }
378    let r = fresh.wram();
379    eprintln!(
380        "{name}: game state {:#04x}, progress flags {:#04x}, spawn {:#04x}, sword {}",
381        r[0xF3C5], r[0xF3C6], r[0xF3C8], r[0xF359]
382    );
383    assert!(states.save(name, &fresh.snapshot().expect("snapshot")), "saving {name}");
384    png(&out.join(format!("{name}.png")), &fresh);
385    eprintln!("{name}: saved");
386}