///|
priv struct Checker {
  mut env : TypeEnv
  inferred : Map[Int, CelType]
  diagnostics : Array[Diagnostic]
}

///|
fn Checker::new(env : TypeEnv) -> Checker {
  { env, inferred: {}, diagnostics: [] }
}

///|
fn Checker::record(self : Checker, expr : Expr, cel_type : CelType) -> CelType {
  self.inferred[expr.id] = cel_type
  cel_type
}

///|
fn Checker::error(
  self : Checker,
  code : String,
  message : String,
  span : SourceSpan,
) -> Unit {
  self.diagnostics.push(Diagnostic::error(code, message, span))
}

///|
fn Checker::require(
  self : Checker,
  actual : CelType,
  expected : CelType,
  span : SourceSpan,
  context : String,
) -> Bool {
  if expected.accepts(actual) {
    true
  } else {
    self.error(
      "CEL-CHECK-TYPE",
      "\{context} requires \{expected.display()}, found \{actual.display()}",
      span,
    )
    false
  }
}

///|
fn literal_type(literal : Literal) -> CelType {
  match literal {
    Literal::Null => CelType::Null
    Literal::Bool(_) => CelType::Bool
    Literal::Int(_) => CelType::Int
    Literal::UInt(_) => CelType::UInt
    Literal::Double(_) => CelType::Double
    Literal::String(_) => CelType::String
    Literal::Bytes(_) => CelType::Bytes
  }
}

///|
fn Checker::infer_identifier(
  self : Checker,
  expr : Expr,
  name : String,
) -> CelType {
  match self.env.lookup(name) {
    Some(cel_type) => self.record(expr, cel_type)
    None => {
      self.error("CEL-CHECK-NAME", "undeclared variable '\{name}'", expr.span)
      self.record(expr, CelType::Dyn)
    }
  }
}

///|
fn Checker::infer_list(
  self : Checker,
  expr : Expr,
  items : Array[Expr],
) -> CelType {
  let element_type = items.fold(init=CelType::Dyn, (current, item) => {
    current.common(self.infer(item))
  })
  self.record(expr, CelType::List(element_type))
}

///|
fn valid_map_key(cel_type : CelType) -> Bool {
  cel_type
  is (CelType::Dyn
  | CelType::Bool
  | CelType::Int
  | CelType::UInt
  | CelType::String)
}

///|
fn Checker::infer_map(
  self : Checker,
  expr : Expr,
  entries : Array[MapExprEntry],
) -> CelType {
  let mut key_type = CelType::Dyn
  let mut value_type = CelType::Dyn
  for entry in entries {
    let found_key = self.infer(entry.key)
    let found_value = self.infer(entry.value)
    if !valid_map_key(found_key) {
      self.error(
        "CEL-CHECK-MAP-KEY",
        "map keys must be bool, int, uint, or string; found \{found_key.display()}",
        entry.key.span,
      )
    }
    key_type = key_type.common(found_key)
    value_type = value_type.common(found_value)
  }
  self.record(expr, CelType::Map(key_type, value_type))
}

///|
fn Checker::infer_unary(
  self : Checker,
  expr : Expr,
  operator : UnaryOp,
  operand : Expr,
) -> CelType {
  let operand_type = self.infer(operand)
  let result = match operator {
    UnaryOp::Not => {
      ignore(
        self.require(
          operand_type,
          CelType::Bool,
          operand.span,
          "logical negation",
        ),
      )
      CelType::Bool
    }
    UnaryOp::Negate => {
      if !operand_type.is_numeric() && operand_type != CelType::Dyn {
        self.error(
          "CEL-CHECK-NEGATE",
          "numeric negation requires int, uint, or double; found \{operand_type.display()}",
          operand.span,
        )
      }
      if operand_type == CelType::UInt {
        self.error(
          "CEL-CHECK-NEGATE-UINT",
          "uint values cannot be negated",
          operand.span,
        )
      }
      operand_type
    }
  }
  self.record(expr, result)
}

///|
fn Checker::numeric_pair(
  self : Checker,
  left : CelType,
  right : CelType,
  span : SourceSpan,
  context : String,
) -> CelType {
  if left == CelType::Dyn || right == CelType::Dyn {
    return CelType::Dyn
  }
  if !left.is_numeric() || !right.is_numeric() {
    self.error(
      "CEL-CHECK-NUMERIC",
      "\{context} requires numeric operands; found \{left.display()} and \{right.display()}",
      span,
    )
    return CelType::Dyn
  }
  if left != right {
    self.error(
      "CEL-CHECK-NUMERIC-MIX",
      "CEL numeric operators require matching types; found \{left.display()} and \{right.display()}",
      span,
    )
    CelType::Dyn
  } else {
    left
  }
}

///|
fn Checker::infer_addition(
  self : Checker,
  left : CelType,
  right : CelType,
  span : SourceSpan,
) -> CelType {
  match (left, right) {
    (CelType::String, CelType::String) => CelType::String
    (CelType::Bytes, CelType::Bytes) => CelType::Bytes
    (CelType::List(left_element), CelType::List(right_element)) =>
      CelType::List(left_element.common(right_element))
    (CelType::Dyn, _) | (_, CelType::Dyn) => CelType::Dyn
    _ => self.numeric_pair(left, right, span, "addition")
  }
}

///|
fn Checker::infer_relation(
  self : Checker,
  left : CelType,
  right : CelType,
  span : SourceSpan,
) -> CelType {
  if left == CelType::Dyn || right == CelType::Dyn {
    return CelType::Bool
  }
  if left == right &&
    (left.is_numeric() || left == CelType::String || left == CelType::Bytes) {
    CelType::Bool
  } else {
    self.error(
      "CEL-CHECK-COMPARE",
      "relational comparison is not defined for \{left.display()} and \{right.display()}",
      span,
    )
    CelType::Bool
  }
}

///|
fn Checker::infer_membership(
  self : Checker,
  left : CelType,
  right : CelType,
  span : SourceSpan,
) -> CelType {
  match right {
    CelType::Dyn => CelType::Bool
    CelType::List(element) => {
      if !element.accepts(left) {
        self.error(
          "CEL-CHECK-IN",
          "membership value \{left.display()} does not match list element \{element.display()}",
          span,
        )
      }
      CelType::Bool
    }
    CelType::Map(key, _) => {
      if !key.accepts(left) {
        self.error(
          "CEL-CHECK-IN",
          "membership key \{left.display()} does not match map key \{key.display()}",
          span,
        )
      }
      CelType::Bool
    }
    _ => {
      self.error(
        "CEL-CHECK-IN",
        "right operand of 'in' must be list or map, found \{right.display()}",
        span,
      )
      CelType::Bool
    }
  }
}

///|
fn Checker::infer_binary(
  self : Checker,
  expr : Expr,
  operator : BinaryOp,
  left_expr : Expr,
  right_expr : Expr,
) -> CelType {
  let left = self.infer(left_expr)
  let right = self.infer(right_expr)
  let result = match operator {
    BinaryOp::Or | BinaryOp::And => {
      ignore(
        self.require(left, CelType::Bool, left_expr.span, "logical operator"),
      )
      ignore(
        self.require(right, CelType::Bool, right_expr.span, "logical operator"),
      )
      CelType::Bool
    }
    BinaryOp::Equal | BinaryOp::NotEqual => {
      if !left.accepts(right) {
        self.error(
          "CEL-CHECK-EQUALITY",
          "equality operands must have compatible types; found \{left.display()} and \{right.display()}",
          expr.span,
        )
      }
      CelType::Bool
    }
    BinaryOp::Less
    | BinaryOp::LessEqual
    | BinaryOp::Greater
    | BinaryOp::GreaterEqual => self.infer_relation(left, right, expr.span)
    BinaryOp::In => self.infer_membership(left, right, expr.span)
    BinaryOp::Add => self.infer_addition(left, right, expr.span)
    BinaryOp::Subtract | BinaryOp::Multiply | BinaryOp::Divide =>
      self.numeric_pair(left, right, expr.span, operator.symbol())
    BinaryOp::Modulo => {
      let result = self.numeric_pair(left, right, expr.span, operator.symbol())
      if result == CelType::Double {
        self.error(
          "CEL-CHECK-MODULO",
          "modulo is only defined for int and uint",
          expr.span,
        )
      }
      result
    }
  }
  self.record(expr, result)
}

///|
fn Checker::infer_conditional(
  self : Checker,
  expr : Expr,
  condition : Expr,
  when_true : Expr,
  when_false : Expr,
) -> CelType {
  let condition_type = self.infer(condition)
  ignore(
    self.require(
      condition_type,
      CelType::Bool,
      condition.span,
      "conditional condition",
    ),
  )
  let true_type = self.infer(when_true)
  let false_type = self.infer(when_false)
  if !true_type.accepts(false_type) {
    self.error(
      "CEL-CHECK-CONDITIONAL",
      "conditional branches have incompatible types \{true_type.display()} and \{false_type.display()}",
      expr.span,
    )
  }
  self.record(expr, true_type.common(false_type))
}

///|
fn Checker::infer_select(
  self : Checker,
  expr : Expr,
  target : Expr,
  field : String,
) -> CelType {
  let target_type = self.infer(target)
  let result = match target_type {
    CelType::Dyn => CelType::Dyn
    CelType::Map(CelType::String | CelType::Dyn, value) => value
    _ => {
      self.error(
        "CEL-CHECK-SELECT",
        "field '\{field}' cannot be selected from \{target_type.display()}",
        expr.span,
      )
      CelType::Dyn
    }
  }
  self.record(expr, result)
}

///|
fn Checker::infer_index(
  self : Checker,
  expr : Expr,
  target : Expr,
  index : Expr,
) -> CelType {
  let target_type = self.infer(target)
  let index_type = self.infer(index)
  let result = match target_type {
    CelType::Dyn => CelType::Dyn
    CelType::List(element) => {
      ignore(self.require(index_type, CelType::Int, index.span, "list index"))
      element
    }
    CelType::Map(key, value) => {
      ignore(self.require(index_type, key, index.span, "map index"))
      value
    }
    _ => {
      self.error(
        "CEL-CHECK-INDEX",
        "indexing is not defined for \{target_type.display()}",
        target.span,
      )
      CelType::Dyn
    }
  }
  self.record(expr, result)
}

///|
fn Checker::infer_call(
  self : Checker,
  expr : Expr,
  name : String,
  args : Array[Expr],
  target? : Expr,
) -> CelType {
  if target is None && name == "has" {
    return self.infer_has_macro(expr, args)
  }
  match target {
    Some(receiver) =>
      if is_comprehension_macro(name) {
        return self.infer_comprehension_macro(expr, receiver, name, args)
      }
    None => ()
  }
  let receiver_type = target.map(value => self.infer(value))
  let argument_types = args.map(arg => self.infer(arg))
  let result = self.resolve_standard_call(
    name,
    receiver_type,
    argument_types,
    expr.span,
  )
  self.record(expr, result)
}

///|
fn Checker::infer(self : Checker, expr : Expr) -> CelType {
  match expr.kind {
    ExprKind::Literal(value) => self.record(expr, literal_type(value))
    ExprKind::Identifier(name) => self.infer_identifier(expr, name)
    ExprKind::List(items) => self.infer_list(expr, items)
    ExprKind::Map(entries) => self.infer_map(expr, entries)
    ExprKind::Unary(operator, operand) =>
      self.infer_unary(expr, operator, operand)
    ExprKind::Binary(operator, left, right) =>
      self.infer_binary(expr, operator, left, right)
    ExprKind::Conditional(condition, when_true, when_false) =>
      self.infer_conditional(expr, condition, when_true, when_false)
    ExprKind::Select(target, field) => self.infer_select(expr, target, field)
    ExprKind::Index(target, index) => self.infer_index(expr, target, index)
    ExprKind::Call(name, args) => self.infer_call(expr, name, args)
    ExprKind::ReceiverCall(target, name, args) =>
      self.infer_call(expr, name, args, target~)
  }
}

///|
/// Type-check a parsed CEL expression in the supplied variable environment.
pub fn check(expr : ParsedExpr, env : TypeEnv) -> CheckResult {
  let checker = Checker::new(env)
  let root_type = checker.infer(expr.root())
  let checked = if checker.diagnostics.any(diagnostic => {
      diagnostic.severity == Severity::Error
    }) {
    None
  } else {
    Some(CheckedExpr::{ parsed: expr, root_type, inferred: checker.inferred })
  }
  { checked, diagnostics: checker.diagnostics }
}

///|
/// Parse and type-check a CEL expression, producing a reusable program.
pub fn compile(
  source : StringView,
  env : TypeEnv,
) -> Program raise CompileError {
  let parsed = parse(source) catch {
    ParseError(code~, message~, span~) =>
      raise CompileError(diagnostics=[Diagnostic::error(code, message, span)])
  }
  let result = check(parsed, env)
  match result.checked {
    Some(checked) => Program::{ checked, }
    None => raise CompileError(diagnostics=result.diagnostics)
  }
}