///|
suberror CodegenError {
  InvalidInput(String)
}

///|
const GENERATED_HEADER : String = "// Generated by Selene codegen. DO NOT EDIT.\n"

///|
fn generate(source : String, input : String) -> String raise CodegenError {
  let (items, reports) = @parser.parse_string(source, name=input)
  guard reports.is_empty() else {
    raise InvalidInput("cannot parse MoonBit source: \{input}")
  }
  let out = StringBuilder()
  let type_attributes = []
  for item in items {
    guard item is TopTypeDef(decl) else { continue }
    let name = decl.tycon
    let mut marked = false
    for attr in decl.attrs {
      let marker = match attr.parsed {
        Some(Ident(id)) if id.qual == Some("selene") => id.name
        Some(Apply(id, _)) if id.qual == Some("selene") => id.name
        _ => {
          guard !attr.raw.has_prefix("#selene.") else {
            raise InvalidInput(
              "\{input}: invalid attribute \{attr.raw} on \{name}",
            )
          }
          continue
        }
      }
      type_attributes.push(attr.loc.start.cnum)
      guard !marked else {
        raise InvalidInput("\{input}: multiple Selene declarations on \{name}")
      }
      marked = true
      guard decl.params.is_empty() else {
        raise InvalidInput("\{input}: generic Selene declaration \{name}")
      }
      guard decl.components is Record(..) ||
        decl.components is TupleStruct(_) ||
        decl.components is Variant(_) else {
        raise InvalidInput("\{input}: \{name} must be a struct or enum")
      }
      let visibility = if decl.type_vis is Pub(..) { "pub " } else { "" }
      match (marker, attr.parsed) {
        ("component", Some(Ident(_))) => {
          let key = "selene_component_key_\{name}"
          out.write_string(
            "\n///|\nlet \{key} : @ecs.ComponentKey[\{name}] = ComponentKey()\n",
          )
          out.write_string(
            "\n///|\n\{visibility}impl @ecs.Component for \{name} with component_key() {\n  \{key}\n}\n",
          )
          if !visibility.is_empty() {
            hide_method(out, name, "@ecs.Component", "component_key")
          }
        }
        ("resource", Some(Ident(_))) => {
          let key = "selene_resource_key_\{name}"
          out.write_string(
            "\n///|\nlet \{key} : @ecs.ResourceKey[\{name}] = ResourceKey()\n",
          )
          out.write_string(
            "\n///|\n\{visibility}impl @ecs.Resource for \{name} with resource_key() {\n  \{key}\n}\n",
          )
          if !visibility.is_empty() {
            hide_method(out, name, "@ecs.Resource", "resource_key")
          }
        }
        _ =>
          raise InvalidInput(
            "\{input}: invalid \{attr.raw} on \{name}; expected #selene.component or #selene.resource",
          )
      }
    }
  }
  for
    (token, start, _) in @lexer.tokens_from_string(source, comment=false).tokens {
    if token is ATTRIBUTE(("selene", Some(marker), _)) &&
      !type_attributes.contains(start.cnum) {
      raise InvalidInput(
        "\{input}:\{start.lnum}: #selene.\{marker} must be on a top-level struct or enum",
      )
    }
  }
  let result = StringBuilder()
  result.write_string(GENERATED_HEADER)
  let body = @fmt.format(out.to_string()).trim()
  if !body.is_empty() {
    result.write_char('\n')
    for line in body.split("\n") {
      result.write_string(line.trim_end().to_owned())
      result.write_char('\n')
    }
  }
  result.to_string()
}

///|
fn hide_method(
  out : StringBuilder,
  name : String,
  trait_name : String,
  method_name : String,
) -> Unit {
  out.write_string(
    "\n///|\n#deprecated(\"Use the typed Selene API.\", skip_current_package=true)\n#doc(hidden)\n",
  )
  out.write_string("pub extend \{name} with \{trait_name}::{\{method_name}}\n")
}

///|
async fn check_source(input : String) -> Unit {
  if input.has_suffix(".g.mbt") {
    let source = input[:input.length() - ".g.mbt".length()].to_owned() + ".mbt"
    if !@fs.exists(source) &&
      @fs.read_file(input).text().has_prefix(GENERATED_HEADER) {
      raise InvalidInput(
        "\{input}: generated companion has no source file \{source}",
      )
    }
    return
  }
  guard input.has_suffix(".mbt") else {
    raise InvalidInput("\{input}: expected a MoonBit source file")
  }
  let output = input[:input.length() - ".mbt".length()].to_owned() + ".g.mbt"
  let source = @fs.read_file(input).text()
  let existing = if @fs.exists(output) {
    Some(@fs.read_file(output).text())
  } else {
    None
  }
  if !source.contains("#selene.") &&
    !(existing is Some(content) && content.has_prefix(GENERATED_HEADER)) {
    return
  }
  let generated = generate(source, input)
  match existing {
    None if generated != GENERATED_HEADER =>
      raise InvalidInput("\{input}: missing generated companion \{output}")
    Some(content) if content != generated &&
      (generated != GENERATED_HEADER || content.has_prefix(GENERATED_HEADER)) =>
      raise InvalidInput("\{output}: generated companion is out of date")
    _ => ()
  }
}

///|
async fn run_cli(args : Array[String]) -> Unit {
  if args.get(1) == Some("--check") {
    guard args.length() > 2 else {
      raise InvalidInput("usage: codegen --check ...")
    }
    for i in 2.. ")
  }
  let input = args[1]
  let output = args[2]
  guard input.has_suffix(".mbt") &&
    !input.has_suffix(".g.mbt") &&
    output.has_suffix(".g.mbt") else {
    raise InvalidInput("expected an .mbt source and a separate .g.mbt output")
  }
  let generated = generate(@fs.read_file(input).text(), input)
  if @fs.exists(output) && @fs.read_file(output).text() == generated {
    return
  }
  @fs.write_file(output, generated, create_mode=CreateOrTruncate)
}

///|
async fn main {
  run_cli(@env.args()) catch {
    InvalidInput(message) => {
      @stdio.stderr.write("\{message}\n")
      @env.exit(1)
    }
    err => raise err
  }
}