///|
/// The first line of each generated Python file.
let generated_header : String = "# Code generated by python-jsonrpc-codegen. DO NOT EDIT."

///|
/// The model configuration of a generated structure, before the `extra` policy.
let model_options : String = "frozen=True, validate_by_name=True, validate_by_alias=True"

///|
/// The state of one generator run.
priv struct Gen {
  model : @smithy.Model
  settings : @jsonrpc.PeerSettings
  names : @codegen.ServiceNames
  package_name : String
  delegator : @codegen.WriterDelegator
  // The prepared structure contracts of the peer's services, by shape.
  values : Map[@smithy.ShapeId, @validation.StructureContract]
  value_config : @validation.Config
  // The shapes that I wrote, the error shapes and the other structures, in the
  // order that I wrote them.
  written : Map[@smithy.ShapeId, Unit]
  errors : Array[@smithy.ShapeId]
  structures : Array[@smithy.ShapeId]
  // The generated class names, to find two shapes with one name.
  type_names : Map[String, @smithy.ShapeId]
  // The docstring of each module, by the namespace of its writer.
  docs : Map[String, String]
}

///|
/// The namespace of a module, relative to the generated package, such as
/// `.models`. A test module has an absolute name, such as `tests.test_generated`.
fn module_namespace(module_ : String) -> @codegen.Namespace {
  @codegen.Namespace::parse(module_, separator=".").unwrap()
}

///|
/// Record that the file imports `name` from `module_`, and give the name.
fn use_name(
  writer : @codegen.CodeWriter,
  module_ : String,
  name : String,
) -> String {
  guard @codegen.SymbolName::parse(name) is Some(symbol_name) else {
    return name
  }
  writer.add_import(
    @codegen.Symbol::new(symbol_name, ns=module_namespace(module_)),
  )
  name
}

///|
/// Record the import of a type annotation, and give its text.
fn use_ty(writer : @codegen.CodeWriter, ty : Ty) -> String {
  if ty.import_of() is Some((module_, name)) {
    let _ = use_name(writer, module_, name)
  }
  ty.text()
}

///|
/// The standard library modules that generated code imports.
let stdlib_modules : Array[String] = [
  "collections.abc", "decimal", "enum", "typing",
]

///|
/// The isort section of a module: 0 `__future__`, 1 the standard library, 2
/// third-party packages, 3 the generated package, 4 relative imports.
fn import_section(module_ : String, package_name : String) -> Int32 {
  if module_ == "__future__" {
    0
  } else if module_.has_prefix(".") {
    4
  } else if stdlib_modules.contains(module_) {
    1
  } else if module_ == package_name || module_.has_prefix(package_name + ".") {
    3
  } else {
    2
  }
}

///|
/// The kind of an imported name in the isort order of ruff: a constant, a
/// class, then any other name.
fn member_kind(name : String) -> Int32 {
  let chars = name.to_array()
  let cased = chars.iter().any(c => c.is_ascii_alphabetic())
  let upper = chars.iter().all(c => !c.is_ascii_lowercase())
  if chars.length() > 1 && cased && upper {
    0
  } else if chars.length() > 0 && chars[0].is_ascii_uppercase() {
    1
  } else {
    2
  }
}

///|
fn compare_members(a : String, b : String) -> Int32 {
  match member_kind(a).compare(member_kind(b)) {
    0 =>
      match a.to_lower().lexical_compare(b.to_lower()) {
        0 => a.lexical_compare(b)
        order => order
      }
    order => order
  }
}

///|
/// The import block of a file, in the order of ruff's isort rules: sections in
/// order, a blank line between sections, modules in lexical order without
/// case, and names by kind, then without case. A line longer than 100
/// characters has one name on each line.
fn render_imports(
  entries : Array[@codegen.ImportEntry],
  package_name : String,
  future : Bool,
) -> String {
  let by_module : Map[String, Array[String]] = Map([])
  if future {
    by_module["__future__"] = ["annotations"]
  }
  for entry in entries {
    let module_ = entry.ns.text
    let name = entry.name.to_string()
    match by_module.get(module_) {
      Some(names) => if !names.contains(name) { names.push(name) }
      None => by_module[module_] = [name]
    }
  }
  let modules = by_module.keys().collect()
  modules.sort_by((a, b) => {
    match
      import_section(a, package_name).compare(import_section(b, package_name)) {
      0 =>
        match a.to_lower().lexical_compare(b.to_lower()) {
          0 => a.lexical_compare(b)
          order => order
        }
      order => order
    }
  })
  let out = StringBuilder()
  let mut section = -1
  for module_ in modules {
    let current = import_section(module_, package_name)
    if section >= 0 && current != section {
      out.write_string("\n")
    }
    section = current
    let names = by_module[module_]
    names.sort_by(compare_members)
    let line = "from \{module_} import \{names.join(", ")}"
    if line.length() <= 100 {
      out.write_string(line + "\n")
    } else {
      out.write_string("from \{module_} import (\n")
      for name in names {
        out.write_string("    \{name},\n")
      }
      out.write_string(")\n")
    }
  }
  out.to_string()
}

///|
/// The text with at most two blank lines in a row, no blank line at the start
/// and one newline at the end.
fn tidy(text : String) -> String {
  let out = StringBuilder()
  let mut blank = 0
  let mut started = false
  for line in text.split("\n") {
    if line.trim_end(chars=" ").is_empty() {
      blank += 1
      continue
    }
    if started {
      out.write_string("\n")
      for _ in 0..<(if blank > 2 { 2 } else { blank }) {
        out.write_string("\n")
      }
    }
    started = true
    blank = 0
    out.write_string(line.to_owned())
  }
  out.write_string("\n")
  out.to_string()
}

///|
fn Gen::shape(
  self : Gen,
  id : @smithy.ShapeId,
) -> @smithy.Shape raise @codegen.CodegenError {
  match self.model.shapes.get(id) {
    Some(shape) => shape
    None =>
      raise @codegen.InvalidSettings(
        "I could not find shape `\{id}` in the model.",
      )
  }
}

///|
fn Gen::is_error(self : Gen, id : @smithy.ShapeId) -> Bool {
  match self.model.shapes.get(id) {
    Some(shape) => shape.trait_value("smithy.api#error") is Some(_)
    None => false
  }
}

///|
/// The class name of a generated shape. An error shape names its exception.
fn Gen::class_of(self : Gen, id : @smithy.ShapeId) -> String {
  class_name(self.names.shape_name(id))
}

///|
/// The class name of the data model of an error shape.
fn Gen::data_class_of(self : Gen, id : @smithy.ShapeId) -> String {
  class_name(self.names.shape_name(id) + "Data")
}

///|
/// The type of a structure, union or enum value: its class and module. A
/// value of an error shape is its data model.
fn Gen::named(self : Gen, id : @smithy.ShapeId) -> Ty {
  if self.is_error(id) {
    Named(self.data_class_of(id), errors_module)
  } else {
    Named(self.class_of(id), models_module)
  }
}

///|
/// The timestamp format of a member and its target.
fn Gen::timestamp_format(
  self : Gen,
  member_shape : @smithy.Shape?,
  target : @smithy.Shape?,
) -> String {
  let from_trait = (shape : @smithy.Shape?) => {
    match shape {
      Some(shape) =>
        match shape.trait_value("smithy.api#timestampFormat") {
          Some(String(format)) => Some(format)
          _ => None
        }
      None => None
    }
  }
  match (from_trait(member_shape), from_trait(target)) {
    (Some(format), _) | (None, Some(format)) => format
    (None, None) =>
      match self.value_config.timestamp_format {
        DateTime => "date-time"
        EpochSeconds => "epoch-seconds"
      }
  }
}

///|
/// The type of the value of a member. A list, a map and a sparse wrapper
/// become layers; the walk keeps its own stack of layers.
fn Gen::member_ty(
  self : Gen,
  member_id : @smithy.ShapeId,
) -> Ty raise @codegen.CodegenError {
  let layers : Array[Int32] = [] // 0 list, 1 map, 2 optional item
  let mut member_shape = self.shape(member_id)
  let leaf = for ;; {
    guard member_shape.body is Member(target~) else {
      raise @codegen.InvalidSettings(
        "I expected `\{member_shape.id}` to be a member shape.",
      )
    }
    if target.namespace() == "smithy.api" {
      let target_shape = self.model.shapes.get(target)
      let format = self.timestamp_format(Some(member_shape), target_shape)
      match prelude_ty(target.name(), format) {
        Some(ty) => break ty
        None =>
          raise @codegen.InvalidSettings(
            "I cannot map prelude shape `\{target}` to a Python type.",
          )
      }
    }
    let target_shape = self.shape(target)
    match target_shape.body {
      Simple(simple) =>
        break simple_ty(
          simple,
          self.timestamp_format(Some(member_shape), Some(target_shape)),
        )
      List(item) => {
        layers.push(0)
        if target_shape.trait_value("smithy.api#sparse") is Some(_) {
          layers.push(2)
        }
        member_shape = self.shape(item)
      }
      Map(key~, value~) => {
        let key_shape = self.shape(key)
        if key_shape.body is Member(target=key_target) &&
          !self.is_string(key_target) {
          raise @codegen.InvalidSettings(
            "I expected the key of map `\{target}` to be a string.",
          )
        }
        layers.push(1)
        if target_shape.trait_value("smithy.api#sparse") is Some(_) {
          layers.push(2)
        }
        member_shape = self.shape(value)
      }
      Structure(_) | Union(_) | Enum(_) | IntEnum(_) => break self.named(target)
      _ =>
        raise @codegen.InvalidSettings(
          "I cannot use shape `\{target}` (\{target_shape.body.keyword()}) as a value in Python.",
        )
    }
  }
  let mut ty = leaf
  for index = layers.length() - 1; index >= 0; index = index - 1 {
    ty = match layers[index] {
      0 => ListOf(ty)
      1 => MapOf(ty)
      _ => Opt(ty)
    }
  }
  ty
}

///|
fn Gen::is_string(self : Gen, id : @smithy.ShapeId) -> Bool {
  if id == @smithy.ShapeId::parse("smithy.api#String").unwrap() {
    return true
  }
  match self.model.shapes.get(id) {
    Some({ body: Simple(String) | Enum(_), .. }) => true
    _ => false
  }
}

///|
/// A field of a generated model.
priv struct Field {
  member_name : String
  field_name : String
  wire_name : String
  ty : Ty
  required : Bool
}

///|
/// The fields of a structure, in model order, with the wire names and the
/// required flags of the prepared contract.
fn Gen::fields(
  self : Gen,
  shape : @smithy.Shape,
) -> Array[Field] raise @codegen.CodegenError {
  guard shape.body is Structure(members) else { return [] }
  let contract = self.values.get(shape.id)
  let out : Array[Field] = []
  let seen : Map[String, String] = Map([])
  for member_id in members {
    let member_name = member_id.member_name().unwrap_or("")
    let member_contract = contract.bind(c => {
      c.members.iter().find_first(m => m.shape == member_id)
    })
    let wire_name = match member_contract {
      Some(m) => m.wire_name
      None => member_name
    }
    let required = match member_contract {
      Some(m) => m.required
      None =>
        self.shape(member_id).trait_value("smithy.api#required") is Some(_)
    }
    let field_name = snake_name(member_name)
    if seen.get(field_name) is Some(other) {
      raise @codegen.InvalidSettings(
        "I found members `\{other}` and `\{member_name}` of `\{shape.id}` with the same Python field name `\{field_name}`.",
      )
    }
    seen[field_name] = member_name
    let ty = self.member_ty(member_id)
    out.push({ member_name, field_name, wire_name, ty, required, })
  }
  out
}

///|
fn Gen::extra_policy(self : Gen, id : @smithy.ShapeId) -> String {
  match self.values.get(id) {
    Some(contract) if contract.unknown_members is Reject => "forbid"
    _ => "ignore"
  }
}

///|
fn Gen::claim_name(
  self : Gen,
  name : String,
  owner : @smithy.ShapeId,
) -> Unit raise @codegen.CodegenError {
  match self.type_names.get(name) {
    Some(other) if other != owner =>
      raise @codegen.InvalidSettings(
        "I found shapes `\{other}` and `\{owner}` with the same Python class name `\{name}`.",
      )
    _ => self.type_names[name] = owner
  }
}

///|
fn Gen::package_path(self : Gen, file : String) -> @codegen.OutputPath {
  @codegen.OutputPath::parse(self.package_name + "/" + file).unwrap()
}

///|
/// Run `body` with the writer of a module of the package. `module_` is the
/// relative module name, such as `.models`.
fn Gen::with_module(
  self : Gen,
  module_ : String,
  doc : String,
  body : (@codegen.CodeWriter) -> Unit raise @codegen.CodegenError,
) -> Unit raise @codegen.CodegenError {
  if !self.docs.contains(module_) {
    self.docs[module_] = doc
  }
  let file = module_.unsafe_substring(start=1, end=module_.length()) + ".py"
  self.delegator.use_file_writer(
    self.package_path(file),
    module_namespace(module_),
    body,
  )
}

///|
fn Gen::models_doc(self : Gen) -> String {
  "The data types of the services of peer `\{self.settings.peer.name.to_string()}`."
}

///|
fn Gen::errors_doc(self : Gen) -> String {
  "The modeled errors of the services of peer `\{self.settings.peer.name.to_string()}`."
}

///|
/// Write the `model_config` line of a model.
fn write_config(
  writer : @codegen.CodeWriter,
  extra : String,
) -> Unit raise @codegen.CodegenError {
  let config = use_name(writer, "pydantic", "ConfigDict")
  writer.write("model_config = $L(", args=[Text(config)])
  writer.indent()
  writer.write("extra=$L, $L", args=[
    Text(py_string(extra)),
    Text(model_options),
  ])
  writer.dedent()
  writer.write(")")
}

///|
/// Write the field declarations of a model.
fn write_fields(
  writer : @codegen.CodeWriter,
  fields : Array[Field],
) -> Unit raise @codegen.CodegenError {
  for field in fields {
    let ty = if field.required { field.ty } else { Opt(field.ty) }
    let text = use_ty(writer, ty)
    let aliased = field.wire_name != field.field_name
    let default = if field.required {
      if aliased {
        let field_fn = use_name(writer, "pydantic", "Field")
        " = \{field_fn}(validation_alias=\{py_string(field.wire_name)}, serialization_alias=\{py_string(field.wire_name)})"
      } else {
        ""
      }
    } else if aliased {
      let field_fn = use_name(writer, "pydantic", "Field")
      " = \{field_fn}(\n        default=None,\n        validation_alias=\{py_string(field.wire_name)},\n        serialization_alias=\{py_string(field.wire_name)},\n    )"
    } else {
      " = None"
    }
    writer.write("$L: $L$L", args=[
      Text(field.field_name),
      Text(text),
      Text(default),
    ])
  }
}

///|
/// Write the Pydantic model of a structure.
fn Gen::write_model(
  self : Gen,
  writer : @codegen.CodeWriter,
  name : String,
  shape : @smithy.Shape,
  doc : String,
) -> Unit raise @codegen.CodegenError {
  let fields = self.fields(shape)
  writer.reserve_name(@codegen.SymbolName::parse(name).unwrap())
  let base = use_name(writer, "pydantic", "BaseModel")
  writer.blank_line()
  writer.blank_line()
  writer.write("class $L($L):", args=[Text(name), Text(base)])
  writer.indent()
  writer.write("\"\"\"$L\"\"\"", args=[Text(doc_text(doc))])
  writer.blank_line()
  write_config(writer, self.extra_policy(shape.id))
  if !fields.is_empty() {
    writer.blank_line()
    write_fields(writer, fields)
  }
  writer.dedent()
}

///|
/// Write a structure that is not an error.
fn Gen::write_structure(
  self : Gen,
  shape : @smithy.Shape,
) -> Unit raise @codegen.CodegenError {
  let name = self.class_of(shape.id)
  self.claim_name(name, shape.id)
  self.with_module(models_module, self.models_doc(), w => {
    self.write_model(w, name, shape, "The structure `\{shape.id}`.")
  })
}

///|
/// The member tag of a union member: its member name, or its `jsonName` when
/// the value configuration uses JSON names.
fn Gen::union_tag(
  self : Gen,
  member_id : @smithy.ShapeId,
) -> String raise @codegen.CodegenError {
  let member_name = member_id.member_name().unwrap_or("")
  if self.value_config.member_names is JsonName &&
    self.shape(member_id).trait_value("smithy.api#jsonName")
    is Some(String(name)) {
    return name
  }
  member_name
}

///|
/// Write a union: one model for each member, and a type alias that selects
/// the member model by its tag.
fn Gen::write_union(
  self : Gen,
  shape : @smithy.Shape,
) -> Unit raise @codegen.CodegenError {
  guard shape.body is Union(members) else { return }
  let name = self.class_of(shape.id)
  self.claim_name(name, shape.id)
  let variants : Array[(String, Field)] = []
  for member_id in members {
    let member_name = member_id.member_name().unwrap_or("")
    let variant = class_name(name + pascal_name(member_name))
    self.claim_name(variant, shape.id)
    let field : Field = {
      member_name,
      field_name: snake_name(member_name),
      wire_name: self.union_tag(member_id),
      ty: self.member_ty(member_id),
      required: true,
    }
    variants.push((variant, field))
  }
  self.with_module(models_module, self.models_doc(), w => {
    for entry in variants {
      let (variant, field) = entry
      w.reserve_name(@codegen.SymbolName::parse(variant).unwrap())
      let base = use_name(w, "pydantic", "BaseModel")
      let class_var = use_name(w, "typing", "ClassVar")
      w.blank_line()
      w.blank_line()
      w.write("class $L($L):", args=[Text(variant), Text(base)])
      w.indent()
      w.write("\"\"\"$L\"\"\"", args=[
        Text(
          doc_text("The member `\{field.member_name}` of union `\{shape.id}`."),
        ),
      ])
      w.blank_line()
      write_config(w, "forbid")
      w.write("union_member: $L[str] = $L", args=[
        Text(class_var),
        Text(py_string(field.wire_name)),
      ])
      w.blank_line()
      write_fields(w, [field])
      w.dedent()
    }
    w.reserve_name(@codegen.SymbolName::parse(name).unwrap())
    let type_alias = use_name(w, "typing", "TypeAlias")
    w.blank_line()
    w.blank_line()
    if variants.length() == 1 {
      w.write("$L: $L = $L", args=[
        Text(name),
        Text(type_alias),
        Text(variants[0].0),
      ])
    } else {
      let annotated = use_name(w, "typing", "Annotated")
      let tag = use_name(w, "pydantic", "Tag")
      let discriminator = use_name(w, "pydantic", "Discriminator")
      let union_tag = use_name(w, runtime_module, "union_tag")
      w.write("$L: $L = $L[", args=[
        Text(name),
        Text(type_alias),
        Text(annotated),
      ])
      w.indent()
      for index, entry in variants {
        let (variant, field) = entry
        let prefix = if index == 0 { "" } else { "| " }
        let suffix = if index == variants.length() - 1 { "," } else { "" }
        w.write("$L$L[$L, $L($L)]$L", args=[
          Text(prefix),
          Text(annotated),
          Text(variant),
          Text(tag),
          Text(py_string(field.wire_name)),
          Text(suffix),
        ])
      }
      w.write("$L($L),", args=[Text(discriminator), Text(union_tag)])
      w.dedent()
      w.write("]")
    }
    w.write("\"\"\"$L\"\"\"", args=[
      Text(
        doc_text(
          "The union `\{shape.id}`: one model for each member, selected by its tag.",
        ),
      ),
    ])
  })
}

///|
/// Write an enum or an int enum.
fn Gen::write_enum(
  self : Gen,
  shape : @smithy.Shape,
) -> Unit raise @codegen.CodegenError {
  let (members, int_enum) = match shape.body {
    Enum(members) => (members, false)
    IntEnum(members) => (members, true)
    _ => return
  }
  let name = self.class_of(shape.id)
  self.claim_name(name, shape.id)
  let variants : Array[(String, String)] = []
  let seen : Map[String, Unit] = Map([])
  for member_id in members {
    let member_name = member_id.member_name().unwrap_or("")
    let variant = enum_member_name(member_name)
    if seen.contains(variant) {
      raise @codegen.InvalidSettings(
        "I found two members of enum `\{shape.id}` with the Python member name `\{variant}`.",
      )
    }
    seen[variant] = ()
    let value = match
      self.shape(member_id).trait_value("smithy.api#enumValue") {
      Some(String(text)) => py_string(text)
      Some(Number(number)) => number.to_string()
      _ => py_string(member_name)
    }
    variants.push((variant, value))
  }
  self.with_module(models_module, self.models_doc(), w => {
    w.reserve_name(@codegen.SymbolName::parse(name).unwrap())
    let base = use_name(w, "enum", if int_enum { "IntEnum" } else { "StrEnum" })
    w.blank_line()
    w.blank_line()
    w.write("class $L($L):", args=[Text(name), Text(base)])
    w.indent()
    w.write("\"\"\"$L\"\"\"", args=[Text(doc_text("The enum `\{shape.id}`."))])
    w.blank_line()
    for variant in variants {
      w.write("$L = $L", args=[Text(variant.0), Text(variant.1)])
    }
    w.dedent()
  })
}

///|
fn Gen::write_shape(
  self : Gen,
  shape : @smithy.Shape,
) -> Unit raise @codegen.CodegenError {
  if self.written.contains(shape.id) || shape.id.namespace() == "smithy.api" {
    return
  }
  match shape.body {
    Structure(_) => {
      self.written[shape.id] = ()
      if shape.trait_value("smithy.api#error") is Some(_) {
        self.errors.push(shape.id)
      } else {
        self.structures.push(shape.id)
        self.write_structure(shape)
      }
    }
    Union(_) => {
      self.written[shape.id] = ()
      self.write_union(shape)
    }
    Enum(_) | IntEnum(_) => {
      self.written[shape.id] = ()
      self.write_enum(shape)
    }
    _ => ()
  }
}