From 218700616c72a07fdaa3859552e8000797c96d1e Mon Sep 17 00:00:00 2001
From: Chris Mounce <christophermounce@gmail.com>
Date: Sun, 6 Jul 2025 12:44:50 -0700
Subject: [PATCH] Clean up CLI, add arguments and some validation

---
 Cargo.lock  |  7 ++++
 Cargo.toml  |  1 +
 src/main.rs | 96 ++++++++++++++++++++++++++++++++++++++++++++---------
 3 files changed, 89 insertions(+), 15 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index c5af765..a89fab3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -16,6 +16,7 @@ dependencies = [
  "codepage-437",
  "compact_str",
  "insta",
+ "lexopt",
  "nom",
  "peg_macro",
  "pest",
@@ -133,6 +134,12 @@ version = "1.0.15"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c"
 
+[[package]]
+name = "lexopt"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9fa0e2a1fcbe2f6be6c42e342259976206b383122fc152e872795338b5a3f3a7"
+
 [[package]]
 name = "libc"
 version = "0.2.171"
diff --git a/Cargo.toml b/Cargo.toml
index 4e0a4bc..b20c324 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -14,6 +14,7 @@ nom = "7.1.3"
 pest = "2.7.15"
 compact_str = "0.9.0"
 rustc-hash = "2.1.1"
+lexopt = "0.3.1"
 
 [dev-dependencies]
 insta = "1.39.0"
diff --git a/src/main.rs b/src/main.rs
index 006331d..45b9b31 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,26 +5,68 @@ mod peg;
 mod preprocess;
 mod world;
 
-use anyhow::anyhow;
+use anyhow::{Result, anyhow};
 use error::Context as ErrContext;
 use labels::process_labels;
+use lexopt::prelude::*;
 use preprocess::eval::Context;
-use std::{env, error::Error, fs, path::PathBuf, process::exit};
+use std::{
+    env, fs,
+    path::{Path, PathBuf},
+    process::exit,
+};
 use world::World;
 
-fn main() -> Result<(), Box<dyn Error>> {
-    let args: Vec<String> = env::args().collect();
-    if args.len() != 2 {
-        eprintln!("Usage: {} WORLD_FILE", args[0]);
+fn main() -> Result<()> {
+    let mut input_file = None;
+    let mut output_file = None;
+    let mut parser = lexopt::Parser::from_env();
+    let mut has_args = false;
+
+    while let Some(arg) = parser.next()? {
+        has_args = true;
+        match arg {
+            Short('o') | Long("output") => {
+                output_file = Some(parser.value()?.string()?);
+            }
+            Value(val) => {
+                if input_file.is_none() {
+                    input_file = Some(val.string()?);
+                } else {
+                    return Err(anyhow!("Unknown arg: {}", val.to_string_lossy()));
+                }
+            }
+            _ => return Err(arg.unexpected().into()),
+        }
+    }
+
+    if !has_args {
+        eprintln!("Usage: {} INPUT -o OUTPUT", env::args().next().unwrap());
         exit(1);
     }
 
-    let world_filename = &args[1];
-    let bytes = fs::read(world_filename)?;
+    let input_filename = input_file.ok_or_else(|| anyhow!("No input file specified"))?;
+    let output_filename = output_file.ok_or_else(|| anyhow!("No output file specified"))?;
+
+    // Prevent overwriting input file
+    let input_path = Path::new(&input_filename).canonicalize()?;
+    let output_path = Path::new(&output_filename);
+    let output_dir = output_path
+        .parent()
+        .filter(|x| !x.as_os_str().is_empty())
+        .unwrap_or_else(|| Path::new("."))
+        .canonicalize()?;
+    let output_path = output_dir.join(output_path.file_name().unwrap());
+    if input_path == output_path {
+        eprintln!("Error: Output file cannot be the same as input file");
+        exit(1);
+    }
+
+    let bytes = fs::read(&input_path)?;
     let mut world = World::from_bytes(&bytes)?;
 
     // Prepare to evaluate macros from the world file's directory
-    let world_pathbuf = PathBuf::from(&world_filename);
+    let world_pathbuf = PathBuf::from(&input_path);
     let world_dir = world_pathbuf
         .parent()
         .ok_or(anyhow!("Couldn't get world's directory"))?;
@@ -39,7 +81,7 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     // Resolve labels to proper ZZT-OOP
     let base_ctx = ErrContext::new();
-    let ctx = base_ctx.with_file_path(&world_filename);
+    let ctx = base_ctx.with_file_path(&input_filename);
     for (i, board) in world.boards.iter_mut().enumerate() {
         let ctx = ctx.with_board(i);
         *board = process_labels(&board, &ctx);
@@ -47,15 +89,39 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     // Print diagnostics
     let messages = base_ctx.into_messages();
-    for (i, message) in messages.iter().enumerate() {
-        if i > 0 {
-            println!();
+    for message in messages.iter() {
+        println!("{}\n", message.rich_format(&world));
+    }
+    if !messages.is_empty() {
+        let mut warnings = 0;
+        let mut errors = 0;
+        for message in messages {
+            match message.level {
+                error::Level::Error => errors += 1,
+                error::Level::Warning => warnings += 1,
+            }
+        }
+        let plural = |x| if x == 1 { "" } else { "s" };
+        if errors == 0 {
+            println!(
+                "Compilation succeeded with {warnings} warning{}.",
+                plural(warnings)
+            )
+        } else {
+            println!(
+                "Compilation failed with {warnings} warning{} and {errors} error{}.",
+                plural(warnings),
+                plural(errors)
+            );
+            exit(1);
         }
-        println!("{}", message.rich_format(&world));
     }
 
     // Try to write a modified world file
-    fs::write("tmp.zzt", world.to_bytes()?)?;
+    let world_bytes = world
+        .to_bytes()
+        .map_err(|e| anyhow!("Couldn't serialize world, {}", e))?;
+    fs::write(&output_path, world_bytes)?;
 
     Ok(())
 }