///|
/// Trait implemented by every type that can travel through the type
/// registry as an erased value.
///
/// Implementations are typically generated by `desert-codegen`
/// (`type_id = "pkg.TypeName"`); hand-written ones just need a unique,
/// stable type name and a delegate to the type's codec.
pub(open) trait RegisteredValue {
  /// The unique, stable name this type is registered under. Wire data
  /// references registry ids, but the serialize side resolves the id from
  /// this name, so names must match between producers and consumers.
  fn registered_type_name(Self) -> String
  /// Serialize the erased value with the type's own codec (the registry
  /// adds the id prefix).
  fn serialize_registered(Self, SerializationContext) -> Unit raise DesertFailure
}

///|
/// One registered type: its id, name and typed deserialization constructor.
pub struct RegisteredEntry {
  /// Registry id written on the wire (1-based, registration order).
  id : Int32
  /// The registered type name.
  type_name : String
  /// Reads the value's payload (after the id prefix) and boxes it.
  deserialize : (DeserializationContext) -> &RegisteredValue raise DesertFailure
}

///|
/// Builder API assigning consecutive ids starting from 1, mirroring the
/// Scala `DefaultTypeRegistry`:
///
/// ```mbt nocheck
/// let registry = TypeRegistryBuilder::new()
///   .register::(fn(ctx) { MyTypeA::deserialize(ctx) })
///   .register_placeholder()
///   .register::(fn(ctx) { MyTypeB::deserialize(ctx) })
///   .freeze()
/// ```
///
/// `register_placeholder` reserves an id slot without a type, keeping ids
/// stable across releases when a type is retired. Lookups are exact
/// type-name matches (MoonBit has no runtime class hierarchy test).
pub struct TypeRegistryBuilder {
  priv entries : Array[RegisteredEntry?]
  priv names : Array[String]
}

///|
pub fn TypeRegistryBuilder::new() -> TypeRegistryBuilder {
  { entries: [], names: [], }
}

///|
/// Register the next id for `deserialize`, a typed constructor of a
/// `RegisteredValue` implementation. Fails if the type name is already
/// registered.
pub fn TypeRegistryBuilder::register(
  self : TypeRegistryBuilder,
  type_name : String,
  deserialize : (DeserializationContext) -> &RegisteredValue raise DesertFailure,
) -> TypeRegistryBuilder raise DesertFailure {
  if self.names.contains(type_name) {
    raise SerializationFailure("Type \{type_name} is already registered")
  }
  self.names.push(type_name)
  let id = self.entries.length() + 1
  let entry : RegisteredEntry = { id, type_name, deserialize, }
  self.entries.push(Some(entry))
  self
}

///|
/// Reserve the next id without registering a type.
pub fn TypeRegistryBuilder::register_placeholder(
  self : TypeRegistryBuilder,
) -> TypeRegistryBuilder {
  self.entries.push(None)
  self
}

///|
pub fn TypeRegistryBuilder::freeze(self : TypeRegistryBuilder) -> TypeRegistry {
  let by_name : @hashmap.HashMap[String, Int32] = @hashmap.HashMap([])
  let by_id : @hashmap.HashMap[Int32, RegisteredEntry] = @hashmap.HashMap([])
  for entry in self.entries {
    match entry {
      Some(e) => {
        by_name.set(e.type_name, e.id)
        by_id.set(e.id, e)
      }
      None => ()
    }
  }
  { by_name, by_id, }
}

///|
/// An immutable type registry: maps registry ids to typed constructors and
/// type names to ids.
pub struct TypeRegistry {
  priv by_name : @hashmap.HashMap[String, Int32]
  priv by_id : @hashmap.HashMap[Int32, RegisteredEntry]
}

///|
/// A registry where every lookup fails (equivalent of `TypeRegistry.empty`).
pub fn TypeRegistry::empty() -> TypeRegistry {
  { by_name: @hashmap.HashMap([]), by_id: @hashmap.HashMap([]), }
}

///|
/// Look up the id registered for a type name.
pub fn TypeRegistry::id_of(self : TypeRegistry, type_name : String) -> Int32? {
  self.by_name.get(type_name)
}

///|
/// Look up the entry for a registry id.
pub fn TypeRegistry::entry_of(
  self : TypeRegistry,
  id : Int32,
) -> RegisteredEntry? {
  self.by_id.get(id)
}

///|
/// A deserialized unknown value: the resolved registry id and entry name,
/// plus the boxed value. Extract the concrete value by matching on
/// `type_name` (MoonBit has no runtime downcast; model open sets of known
/// types as an enum of yours and register the enum).
pub struct RegisteredBox {
  id : Int32
  type_name : String
  value : &RegisteredValue
}

///|
/// Serialize a value through the registry: writes `var_u32(type_id)`
/// followed by the registered codec's payload, inline — no markers, no
/// length prefix (wire-compatible with both original implementations).
///
/// Fails with `TypeNotRegistered` when the value's type name has no id.
pub fn SerializationContext::write_unknown(
  self : SerializationContext,
  registry : TypeRegistry,
  value : &RegisteredValue,
) -> Unit raise DesertFailure {
  let type_name = value.registered_type_name()
  match registry.id_of(type_name) {
    Some(id) => {
      self.write_var_u32(id.reinterpret_as_uint())
      value.serialize_registered(self)
    }
    None => raise TypeNotRegistered(type_name)
  }
}

///|
/// Read a registry-dispatched value: reads the `var_u32` id, looks the
/// entry up and calls its typed constructor. Unknown ids fail with
/// `InvalidTypeId` — there is no way to skip an unknown payload, since the
/// format carries no length information for it.
pub fn DeserializationContext::read_unknown(
  self : DeserializationContext,
  registry : TypeRegistry,
) -> RegisteredBox raise DesertFailure {
  let id = self.read_var_u32().reinterpret_as_int()
  match registry.entry_of(id) {
    Some(entry) => {
      let value = (entry.deserialize)(self)
      { id, type_name: entry.type_name, value, }
    }
    None => raise InvalidTypeId(id)
  }
}

///|
/// Top-level convenience: serialize an erased value to bytes.
pub fn[T : RegisteredValue] serialize_unknown_to_bytes(
  value : T,
  registry : TypeRegistry,
  options~ : Options = Options::default(),
) -> Bytes raise DesertFailure {
  let out = ByteOutput::new()
  let ctx = SerializationContext::new(out, SerializerState::new(), options)
  ctx.write_unknown(registry, value)
  out.to_bytes()
}

///|
/// Top-level convenience: read a registry-dispatched value.
pub fn deserialize_unknown_from_bytes(
  data : Bytes,
  registry : TypeRegistry,
  options~ : Options = Options::default(),
) -> RegisteredBox raise DesertFailure {
  let ctx = DeserializationContext::new(
    BinaryInput::new(data),
    SerializerState::new(),
    options,
  )
  ctx.read_unknown(registry)
}