///|
/// desert-codegen: generate `BinaryCodec` implementations for types
/// annotated with `#derive.binary_codec`.
///
/// Usage: `moon run codegen --  [...]`
///
/// For every `.mbt` file in the directory (excluding generated outputs and
/// test files), the tool parses the source, finds annotated struct/enum
/// declarations, and (re)writes `/desert_derive.mbt` with the codec
/// implementations. The target package's `moon.pkg` gains a
/// `vigoo/desert` import if missing.

///|
const GENERATED_FILE : String = "desert_derive.mbt"

///|
const GENERATED_BANNER : String = "// Generated by desert-codegen (vigoo/desert) — DO NOT EDIT.\n"

///|
fn main {
  let all_args = @env.args()
  let args : Array[String] = []
  for a in all_args {
    args.push(a)
  }
  // Drop the program path (argv[0]) when present.
  if (args.length() > 0 && !args[0].is_empty() && args[0].has_suffix(".wasm")) ||
    (args.length() > 0 && args[0].has_prefix("/")) {
    let _ = args.remove(0)
  }
  if args.length() == 0 {
    println("usage: desert-codegen  [...]")
    abort("no directories given")
  }
  for dir in args {
    try {
      generate_for_directory(dir)
      println("generated \{dir}/\{GENERATED_FILE}")
    } catch {
      CodegenError(msg) => {
        let m = "ERROR: " + dir + ": " + msg
        println(m)
      }
    }
  }
}

///|
fn generate_for_directory(dir : String) -> Unit raise CodegenError {
  let (structs, enums) = scan_directory(dir)
  let source = render_all(structs, enums)
  let path = dir + "/" + GENERATED_FILE
  if file_exists(path) {
    let existing = read_file(path)
    if existing == source {
      return
    }
  }
  write_file(path, source)
  reconcile_imports(dir)
}

///|
fn render_all(
  structs : Array[StructIR],
  enums : Array[EnumIR],
) -> String raise CodegenError {
  let body = StringBuilder()
  body.write_string(GENERATED_BANNER)
  body.write_string("\n")
  for st in structs {
    body.write_string(emit_struct(st))
  }
  for en in enums {
    body.write_string(emit_enum(en))
  }
  format_source(body.to_string())
}

///|
/// Scan a directory for .mbt sources with `#derive.binary_codec` types.
fn scan_directory(
  dir : String,
) -> (Array[StructIR], Array[EnumIR]) raise CodegenError {
  let structs : Array[StructIR] = []
  let enums : Array[EnumIR] = []
  for filename in list_mbt_files(dir) {
    if filename == GENERATED_FILE {
      continue
    }
    let content = read_file(dir + "/" + filename)
    let (parsed, reports) = @parser.parse_string(content)
    if reports.length() > 0 {
      raise CodegenError(
        "parse error in \{filename}: \{Show::to_string(reports[0])}",
      )
    }
    let (s, e) = scan_impls(parsed)
    for x in s {
      structs.push(x)
    }
    for x in e {
      enums.push(x)
    }
  }
  (structs, enums)
}

///|
/// Round-trip generated source through the parser and the official
/// formatter so the output is `moon fmt`-canonical, and validate that it
/// parses cleanly.
fn format_source(source : String) -> String raise CodegenError {
  let (parsed, reports) = @parser.parse_string(source)
  if reports.length() > 0 {
    raise CodegenError(
      "generated code failed to parse: \{Show::to_string(reports[0])}",
    )
  }
  let formatted = @formatter.impls_to_string(parsed)
  GENERATED_BANNER + "\n" + formatted
}

///|
fn list_mbt_files(dir : String) -> Array[String] raise CodegenError {
  let files : Array[String] = []
  let names : ArrayView[String] = match read_dir_result(dir) {
    Ok(view) => view
    Err(err) => raise CodegenError("cannot read directory \{dir}: \{err}")
  }
  for name in names {
    if name.has_suffix(".mbt") && !name.has_prefix(".") {
      files.push(name)
    }
  }
  files
}

///|
fn read_dir_result(dir : String) -> Result[ArrayView[String], String] {
  let res : Result[ArrayView[String], String] = Ok(@fs.read_dir(dir)) catch {
    e => Err(Repr(e).to_string())
  }
  res
}

///|
fn read_file(path : String) -> String raise CodegenError {
  let res : Result[String, String] = Ok(@fs.read_file_to_string(path)) catch {
    e => Err(Repr(e).to_string())
  }
  match res {
    Ok(content) => content
    Err(err) => raise CodegenError("cannot read \{path}: \{err}")
  }
}

///|
fn file_exists(path : String) -> Bool {
  @fs.path_exists(path)
}

///|
fn write_file(path : String, content : String) -> Unit raise CodegenError {
  let res : Result[Unit, String] = Ok(@fs.write_string_to_file(path, content)) catch {
    e => Err(Repr(e).to_string())
  }
  match res {
    Ok(_) => ()
    Err(err) => raise CodegenError("cannot write \{path}: \{err}")
  }
}

///|
/// Ensure the package's `moon.pkg` imports `vigoo/desert`.
fn reconcile_imports(dir : String) -> Unit raise CodegenError {
  let pkg_path = dir + "/moon.pkg"
  if !file_exists(pkg_path) {
    return
  }
  let content = read_file(pkg_path)
  if content.contains("\"vigoo/desert\"") {
    return
  }
  // Insert into the first import block, or create one.
  let marker = "import {\n"
  let idx = match content.find(marker) {
    Some(i) => i
    None => -1
  }
  let updated = if idx >= 0 {
    let insert_at = idx + marker.length()
    content[0:insert_at].to_owned() +
    "  \"vigoo/desert\",\n" +
    content[insert_at:].to_owned()
  } else {
    "import {\n  \"vigoo/desert\",\n}\n\n" + content
  }
  write_file(pkg_path, updated)
}