///|
pub(all) struct GenerationLimits {
  max_nodes : Int32
  max_depth : Int32
  max_attempts : Int32
} derive(Eq, @debug.Debug)

///|
pub fn GenerationLimits::default() -> GenerationLimits {
  { max_nodes: 128, max_depth: 6, max_attempts: 256, }
}

///|
pub(all) suberror GenerationError {
  InvalidLimits(String)
  AttemptsExhausted(Int32)
} derive(Eq, @debug.Debug)

///|
pub fn GenerationLimits::validate(
  self : GenerationLimits,
) -> Unit raise GenerationError {
  guard self.max_nodes >= 1 && self.max_nodes <= 2048 else {
    raise InvalidLimits("max_nodes must be between 1 and 2048")
  }
  guard self.max_depth >= 0 && self.max_depth <= 16 else {
    raise InvalidLimits("max_depth must be between 0 and 16")
  }
  guard self.max_attempts >= 1 && self.max_attempts <= 8192 else {
    raise InvalidLimits("max_attempts must be between 1 and 8192")
  }
}

///|
pub(all) struct GeneratedProgram {
  program : CompleteProgram
  construction_attempts : Int32
  generator : Generator
}

///|
priv struct Construction {
  rng : @random.Rand
  limits : GenerationLimits
  generator : Generator
  enums : Array[EnumDefinition]
  // Existing holes reserve a node, so closing them never needs more space.
  mut nodes : Int32
  mut attempts : Int32
}

///|
fn Construction::new(
  seed : Int32,
  limits : GenerationLimits,
  generator~ : Generator = Generator::default(),
) -> Construction {
  let bytes = Bytes::makei(32, i => {
    if i < 4 {
      (seed >> (i * 8)).to_byte()
    } else {
      0
    }
  })
  {
    rng: @random.Rand::chacha8(seed=bytes),
    limits,
    generator,
    enums: [],
    nodes: 1,
    attempts: 0,
  }
}

///|
fn Construction::attempt(self : Construction) -> Unit raise GenerationError {
  guard self.attempts < self.limits.max_attempts else {
    raise AttemptsExhausted(self.attempts)
  }
  self.attempts += 1
}

///|
fn Construction::ty(self : Construction) -> Ty {
  if self.rng.boolean() {
    Int
  } else {
    Bool
  }
}

///|
priv enum Rule {
  Terminal
  Binding
  Branch
  Invoke(Int32)
  Plus
  Minus
  Compare
  Negate
  Composite(Expr, EnumDefinition?)
}

///|
fn Construction::tuple_rule(
  self : Construction,
  ty : Ty,
  scope : Array[Binder],
) -> (Expr, EnumDefinition?) {
  let types = Array::makei(2 + self.rng.int(limit=2), _ => self.ty())
  let index = self.rng.int(limit=types.length())
  types[index] = ty
  let binder = {
    id: scope.fold(init=-1, (id, binder) => id.max(binder.id)) + 1,
    ty: Ty::Tuple(types),
  }
  (
    Let(
      binder,
      Tuple(types.map(ty => Hole(ty))),
      Project(Ref(binder.id), index),
    ),
    None,
  )
}

///|
fn Construction::enum_rule(
  self : Construction,
  ty : Ty,
  scope : Array[Binder],
) -> (Expr, EnumDefinition?) {
  let (id, definition, added) = if self.enums.is_empty() {
    let definition : EnumDefinition = {
      variants: Array::makei(2 + self.rng.int(limit=2), _ => {
        Array::makei(self.rng.int(limit=3), _ => self.ty())
      }),
    }
    (0, definition, Some(definition))
  } else {
    let id = self.rng.int(limit=self.enums.length())
    (id, self.enums[id], None)
  }
  let tag = self.rng.int(limit=definition.variants.length())
  let other = if definition.variants.length() > 1 {
    (tag + 1 + self.rng.int(limit=definition.variants.length() - 1)) %
    definition.variants.length()
  } else {
    tag
  }
  let binder = {
    id: scope.fold(init=-1, (id, binder) => id.max(binder.id)) + 1,
    ty: Ty::Enum(id),
  }
  let arms = definition.variants.map(fields => {
    parameters: Array::makei(fields.length(), index => Binder::{
      id: binder.id + index + 1,
      ty: fields[index],
    }),
    body: Hole(ty),
  })
  (
    Let(
      binder,
      If(
        Hole(Bool),
        Variant(id, tag, definition.variants[tag].map(ty => Hole(ty))),
        Variant(id, other, definition.variants[other].map(ty => Hole(ty))),
      ),
      Match(id, Ref(binder.id), arms),
    ),
    added,
  )
}

///|
fn Construction::terminal(
  self : Construction,
  ty : Ty,
  scope : Array[Binder],
  template : Bool,
) -> Expr {
  let visible = scope.filter(binder => binder.ty == ty)
  if !visible.is_empty() && self.rng.boolean() {
    return Ref(visible[self.rng.int(limit=visible.length())].id)
  }
  if template {
    return Hole(ty)
  }
  match ty {
    Bool => BoolLiteral(self.rng.boolean())
    Int => {
      let boundaries = [
        -2147483648, -2147483647, -1, 0, 1, 2147483646, 2147483647,
      ]
      IntLiteral(
        if self.rng.boolean() {
          boundaries[self.rng.int(limit=boundaries.length())]
        } else {
          self.rng.int(limit=17) - 8
        },
      )
    }
    _ => abort("this generator only requests scalar terminal types")
  }
}

///|
/// Both passes use one typed rule set and charge the same construction budget.
fn Construction::expand(
  self : Construction,
  ty : Ty,
  scope : Array[Binder],
  functions : Array[Function],
  room : Int32,
  template : Bool,
) -> Expr raise GenerationError {
  self.attempt()
  let rules = [Terminal]
  let remaining = self.limits.max_nodes - self.nodes
  if room > 0 {
    if remaining >= 2 {
      rules.push(Binding)
    }
    if remaining >= 3 {
      rules.push(Branch)
    }
    for id, function in functions {
      if function.returns == ty && remaining >= function.parameters.length() {
        rules.push(Invoke(id))
      }
    }
    if !template {
      match ty {
        Int =>
          if remaining >= 2 {
            rules.push(Plus)
            rules.push(Minus)
          }
        Bool => {
          if remaining >= 2 {
            rules.push(Compare)
          }
          if remaining >= 1 {
            rules.push(Negate)
          }
        }
        _ => ()
      }
    }
  }
  if self.generator == CompositeV2 && room >= 2 {
    for
      (expression, added) in [
        self.tuple_rule(ty, scope),
        self.enum_rule(ty, scope),
      ] {
      let extra = match added {
        Some(definition) => definition.size()
        None => 0
      }
      let cost = expression.size() - 1 + extra
      if expression.depth() <= room && cost <= remaining {
        rules.push(Composite(expression, added))
      }
    }
  }
  let expression = match rules[self.rng.int(limit=rules.length())] {
    Terminal => return self.terminal(ty, scope, template)
    Binding => {
      let binder = {
        id: scope.fold(init=-1, (id, binder) => id.max(binder.id)) + 1,
        ty: self.ty(),
      }
      Let(binder, Hole(binder.ty), Hole(ty))
    }
    Branch => If(Hole(Bool), Hole(ty), Hole(ty))
    Invoke(id) => Call(id, functions[id].parameters.map(p => Hole(p.ty)))
    Plus => Add(Hole(Int), Hole(Int))
    Minus => Subtract(Hole(Int), Hole(Int))
    Compare => Less(Hole(Int), Hole(Int))
    Negate => Not(Hole(Bool))
    Composite(expression, added) => {
      if added is Some(definition) {
        self.attempt()
        for _ in definition.variants {
          self.attempt()
        }
        self.enums.push(definition)
        self.nodes += definition.size()
      }
      expression
    }
  }
  self.nodes += expression.size() - 1
  self.fill(expression, scope, functions, room, template)
}

///|
fn Construction::fill(
  self : Construction,
  expression : Expr,
  scope : Array[Binder],
  functions : Array[Function],
  room : Int32,
  template : Bool,
) -> Expr raise GenerationError {
  match expression {
    Hole(ty) => self.expand(ty, scope, functions, room, template)
    IntLiteral(_) | BoolLiteral(_) | Ref(_) => expression
    Let(binder, value, body) => {
      let value = self.fill(value, scope, functions, room - 1, template)
      let nested = scope.copy()
      nested.push(binder)
      Let(binder, value, self.fill(body, nested, functions, room - 1, template))
    }
    Add(a, b) =>
      Add(
        self.fill(a, scope, functions, room - 1, template),
        self.fill(b, scope, functions, room - 1, template),
      )
    Subtract(a, b) =>
      Subtract(
        self.fill(a, scope, functions, room - 1, template),
        self.fill(b, scope, functions, room - 1, template),
      )
    Less(a, b) =>
      Less(
        self.fill(a, scope, functions, room - 1, template),
        self.fill(b, scope, functions, room - 1, template),
      )
    Not(e) => Not(self.fill(e, scope, functions, room - 1, template))
    If(condition, yes, no) =>
      If(
        self.fill(condition, scope, functions, room - 1, template),
        self.fill(yes, scope, functions, room - 1, template),
        self.fill(no, scope, functions, room - 1, template),
      )
    Call(id, arguments) =>
      Call(
        id,
        arguments.map(e => self.fill(e, scope, functions, room - 1, template)),
      )
    Tuple(fields) =>
      Tuple(fields.map(e => self.fill(e, scope, functions, room - 1, template)))
    Project(value, index) =>
      Project(self.fill(value, scope, functions, room - 1, template), index)
    Variant(id, tag, arguments) =>
      Variant(
        id,
        tag,
        arguments.map(e => self.fill(e, scope, functions, room - 1, template)),
      )
    Match(id, value, arms) =>
      Match(
        id,
        self.fill(value, scope, functions, room - 1, template),
        arms.map(arm => {
          let nested = scope.copy()
          nested.append(arm.parameters)
          {
            ..arm,
            body: self.fill(arm.body, nested, functions, room - 1, template),
          }
        }),
      )
  }
}

///|
fn Construction::compose_template(
  self : Construction,
) -> Template raise GenerationError {
  let functions : Array[Function] = []
  let count = if self.limits.max_depth == 0 { 0 } else { self.rng.int(limit=4) }
  for _ in 0.. Binder::{
      id,
      ty: self.ty(),
    })
    let cost = 2 + parameters.length()
    if self.nodes + cost > self.limits.max_nodes {
      break
    }
    self.attempt()
    self.nodes += cost
    let returns = self.ty()
    functions.push({ parameters, returns, body: Hole(returns), })
  }
  // Only preceding signatures are visible to a function; cycles cannot be constructed.
  let room = self.limits.max_depth.min(3)
  for id, function in functions {
    functions[id] = {
      ..function,
      body: self.fill(
        function.body,
        function.parameters,
        functions[:id].to_owned(),
        room,
        true,
      ),
    }
  }
  let result = self.expand(self.ty(), [], functions, room, true)
  { enums: self.enums.copy(), functions, result, }
}

///|
fn Construction::complete(
  self : Construction,
  template : Template,
) -> CompleteProgram raise {
  self.enums.clear()
  self.enums.append(template.enums)
  let functions : Array[Function] = []
  for id, function in template.functions {
    functions.push({
      ..function,
      body: self.fill(
        function.body,
        function.parameters,
        template.functions[:id].to_owned(),
        self.limits.max_depth,
        false,
      ),
    })
  }
  let result = self.fill(
    template.result,
    [],
    functions,
    self.limits.max_depth,
    false,
  )
  Template::{ enums: self.enums.copy(), functions, result, }.finish()
}

///|
/// Deterministic, versioned template rules with one random stream and shared budgets.
pub fn synthesize(
  seed : Int32,
  limits~ : GenerationLimits = GenerationLimits::default(),
  generator~ : Generator = Generator::default(),
) -> GeneratedProgram raise {
  limits.validate()
  let state = Construction::new(seed, limits, generator~)
  let template = state.compose_template()
  let program = state.complete(template)
  { program, construction_attempts: state.attempts, generator, }
}

///|
pub extend GenerationLimits with Eq::{equal, not_equal}

///|
pub extend GenerationLimits with @debug.Debug::{to_repr}

///|
pub extend GenerationError with Eq::{equal, not_equal}

///|
pub extend GenerationError with @debug.Debug::{to_repr}