graded

import glance.{type Span, type Statement}
import gleam/option.{type Option}
import gleam/set.{type Set}
import gleam/string

// Whether a name's first character is an uppercase letter. This single
// lexical rule distinguishes effect labels from effect variables and record
// constructors from ordinary functions — the comparison against both cases
// excludes non-cased graphemes (digits, symbols) that compare equal to
// themselves under both `uppercase` and `lowercase`.
pub fn is_upper_initial(name: String) -> Bool {
  case string.first(name) {
    Ok(first) ->
      first == string.uppercase(first) && first != string.lowercase(first)
    Error(Nil) -> False
  }
}

// A fully qualified function name: module path + function name.
pub type QualifiedName {
  QualifiedName(module: String, function: String)
}

// Distinguishes auto-inferred annotations from enforced invariants.
pub type AnnotationKind {
  // Auto-generated by `graded infer`. Not enforced.
  Effects
  // Hand-written invariant. Violations break the build.
  Check
}

// An effect set: either a concrete set of named effects, the universal
// wildcard [_], or a polymorphic set with effect variables.
//
// Effect variables (lowercase identifiers like `e`, `e1`) represent
// "whatever effects the corresponding callback has". They are bound at
// call sites by argument-to-parameter matching and substituted away.
pub type EffectSet {
  // [_] — the universal set, top element of the effect lattice.
  // Declaring [_] means "any effects are permitted here".
  Wildcard
  // A concrete set of named effects, e.g. [Http, Stdout] or [].
  Specific(set: Set(String))
  // A polymorphic set mixing concrete labels with effect variables,
  // e.g. [e] or [Stdout, e]. Variables are resolved by substitution
  // at call sites.
  Polymorphic(labels: Set(String), variables: Set(String))
}

// True iff `actual` is a subset of `declared` in the effect lattice.
// Wildcard as declared always passes; Wildcard as actual against a
// finite declared set always fails. Polymorphic sets are conservatively
// handled: as `actual`, any unresolved variables fail the subset check
// (substitution should happen before comparison); as `declared`,
// variables are treated as open slots that accept anything.
pub fn is_subset(actual: EffectSet, declared: EffectSet) -> Bool {
  case declared, actual {
    Wildcard, _ -> True
    _, Wildcard -> False
    Polymorphic(d_labels, _), Specific(a) -> set.is_subset(a, of: d_labels)
    Polymorphic(d_labels, d_vars), Polymorphic(a_labels, a_vars) ->
      set.is_subset(a_labels, of: d_labels) && set.is_subset(a_vars, of: d_vars)
    Specific(d), Specific(a) -> set.is_subset(a, of: d)
    // Polymorphic with unresolved variables cannot be verified against
    // a finite concrete set. Return False conservatively.
    Specific(_), Polymorphic(_, _) -> False
  }
}

// Union of two effect sets. Wildcard absorbs everything. Polymorphic
// sets merge labels and variables component-wise.
pub fn union(a: EffectSet, b: EffectSet) -> EffectSet {
  case a, b {
    Wildcard, _ -> Wildcard
    _, Wildcard -> Wildcard
    Specific(x), Specific(y) -> Specific(set.union(x, y))
    Polymorphic(l1, v1), Polymorphic(l2, v2) ->
      Polymorphic(set.union(l1, l2), set.union(v1, v2))
    Specific(x), Polymorphic(l, v) -> Polymorphic(set.union(x, l), v)
    Polymorphic(l, v), Specific(x) -> Polymorphic(set.union(l, x), v)
  }
}

// The empty (pure) effect set.
pub fn empty() -> EffectSet {
  Specific(set.new())
}

// Construct a Specific effect set from a list of label strings.
pub fn from_labels(labels: List(String)) -> EffectSet {
  Specific(set.from_list(labels))
}

// True iff this effect set contains unresolved effect variables.
pub fn has_variables(effect_set: EffectSet) -> Bool {
  case effect_set {
    Polymorphic(_, variables) -> !set.is_empty(variables)
    _ -> False
  }
}

// A richer effect representation: a small lambda-calculus-with-union over
// effects. `EffectSet` is its ground normal form — anything an `EffectSet`
// can express, a variable-free, application-free `EffectTerm` can too.
//
// Effect variables come in two kinds, kept implicit in structure:
//   - `Eff`       — a flat effect (a bare `TVar`, e.g. `e`)
//   - `Eff -> Eff` — an effect *operator* (a `TAbs`, used under `TApp`,
//                    e.g. a higher-order parameter `action`)
//
// This is what lets graded express *second-order* effect polymorphism: an
// effect variable that is itself parameterized by a callback. See
// docs/second-order-effects.md.
pub type EffectTerm {
  // Ground labels. `TLabels(∅)` is pure. Kind `Eff`.
  TLabels(labels: Set(String))
  // The wildcard `[_]`, absorbing under union. Kind `Eff`.
  TTop
  // A free effect variable. Kind `Eff` when bare.
  TVar(name: String)
  // Operator application: `operator` applied to `argument`, e.g.
  // `action(Stdout)`. Stuck (left symbolic) when `operator` is an unresolved
  // variable. Kind `Eff`.
  TApp(operator: EffectTerm, argument: EffectTerm)
  // An effect operator `λparam. body` — a higher-order parameter's latent
  // effect as a function of its callback. Kind `Eff -> Eff`.
  TAbs(param: String, body: EffectTerm)
  // Composition of effects (set union). Kind `Eff`.
  TUnion(terms: List(EffectTerm))
}

// An effect bound on a function-typed parameter. The `effects` is an
// `EffectTerm`: a flat `Eff` term for a first-order callback (`f: [e]`), or
// an operator `TAbs` for a higher-order one (`action: fn(cb) -> [cb]`).
//
// `name` is the bare parameter name (`f`) for a parameter bound, or a
// `param.field` path (`handler.on_click`) for a *field bound* — a hand-written
// declaration of a record field's effect at the function boundary, the
// boundary-scoped counterpart to a `type` line. A field bound's path carries a
// dot; a parameter name never can, so the two forms don't collide.
pub type ParamBound {
  ParamBound(name: String, effects: EffectTerm)
}

// An effect annotation from a .graded sidecar file. `effects` is an
// `EffectTerm` — for the common first-order case a variable-free or flat-
// variable term equivalent to an `EffectSet`, but it may carry operator
// applications (`[action(Stdout)]`) for second-order signatures.
pub type EffectAnnotation {
  EffectAnnotation(
    kind: AnnotationKind,
    function: String,
    params: List(ParamBound),
    effects: EffectTerm,
  )
}

// The operator a function *returns*, for a function whose result is itself a
// function (`fn pick() -> fn(fn() -> _) -> _`). Serialized into the spec file so
// the signature crosses module and package boundaries — a downstream
// `let h = pick(); with(h)` resolves `h` to this operator. `function` is
// module-qualified in the spec.
pub type ReturnsAnnotation {
  ReturnsAnnotation(function: String, operator: EffectTerm)
}

// A single line in an .graded file, preserving structure for round-trip rewrites.
pub type GradedLine {
  AnnotationLine(annotation: EffectAnnotation)
  TypeFieldLine(type_field: TypeFieldAnnotation)
  ExternalLine(external: ExternalAnnotation)
  ReturnsLine(returns: ReturnsAnnotation)
  CommentLine(text: String)
  BlankLine
}

// A complete parsed .graded file, preserving line-level structure.
pub type GradedFile {
  GradedFile(lines: List(GradedLine))
}

// A resolved call site found in a function body.
pub type ResolvedCall {
  ResolvedCall(name: QualifiedName, span: Span)
}

// What kind of value is at a call-site argument position? Used for
// binding effect variables during call-site substitution.
pub type ArgumentValue {
  // A qualified function reference, e.g. `io.println` or `types.OutOfRange`.
  // Effects are looked up in the knowledge base.
  FunctionRef(name: QualifiedName)
  // A bare identifier — a local function, a parameter, or an unbound
  // local variable. Resolved against the caller's param bounds or
  // local function map.
  LocalRef(name: String)
  // A record constructor (uppercase-initial qualified or bare name).
  // Pure by Gleam's semantics.
  ConstructorRef
  // An inline closure `fn(params) { body }`. `params` are its parameter names
  // (`_` for discarded), `body` its statements. When passed to an operator
  // parameter, the checker analyses the body — treating the first parameter as
  // the callback — and lifts it to an effect operator so the application
  // beta-reduces (rather than collapsing to `[Unknown]`).
  Closure(params: List(String), body: List(Statement))
  // A value selected among several function-like options by a `case`/`if`
  // (`case c { True -> f  False -> g }`). When passed to an operator parameter,
  // each option is lifted and the results are joined (`(f ⊔ g)(cb) = f(cb) ⊔
  // g(cb)`), so the effect over-approximates every branch. Any non-function
  // branch makes the whole expression `OtherExpression` instead.
  Choice(options: List(ArgumentValue))
  // A value produced by *calling* a function that returns a function
  // (`let h = pick_handler()`). `callee` names the producer; a `""` module is
  // the same-module sentinel (resolved on-demand via the function map),
  // otherwise it's resolved from the producer's inferred returned-operator in
  // the knowledge base. `args` are the producer call's arguments, bound to the
  // producer's parameters when its returned operator is *polymorphic* in them
  // (a decorator, `fn traced(action) { fn(cb) { action(cb) } }`). Lets
  // `with_logger(pick_handler())` resolve instead of `[Unknown]`.
  ReturnedOperator(callee: QualifiedName, args: List(CallArgument))
  // Anything else (a computed expression, literal, etc.). Effects come from
  // the enclosing walk; at the argument level we have no concrete function to
  // propagate.
  OtherExpression
}

// One argument at a call site. `position` respects pipes (the piped
// expression is implicitly position 0 and explicit arguments shift up).
// `label` is `Some(name)` for labeled arguments, `None` otherwise.
pub type CallArgument {
  CallArgument(position: Int, label: Option(String), value: ArgumentValue)
}

// A local (unresolved) call — needs transitive analysis.
pub type LocalCall {
  LocalCall(function: String, span: Span)
}

// A field access call: object.label(args) where object is a local variable.
// `span` is the whole call's span (for diagnostics); `receiver_span` is the
// receiver variable's own span, used to look up its inferred type.
pub type FieldCall {
  FieldCall(object: String, label: String, span: Span, receiver_span: Span)
}

// A *returned operator applied directly*: `let h = pick_handler(); h(cb)`.
// `callee` names the producer (with the `""` same-module sentinel, as in
// `ReturnedOperator`) and `producer_args` are the producer call's arguments;
// together they resolve the operator the producer returns. The direct call's
// own arguments (`cb`) are recorded in `call_args` under `span.start`, so the
// resolved operator is applied to them. Lets a let-bound returned operator
// resolve when *applied directly*, not only when passed to an operator
// parameter.
pub type DirectOperatorCall {
  DirectOperatorCall(
    callee: QualifiedName,
    producer_args: List(CallArgument),
    span: Span,
  )
}

// An inline function-like value used as a *pipe target* and thereby applied to
// the piped value: `x |> fn(f) { f() }` or `x |> case c { _ -> a  _ -> b }`.
// `value` is the lifted operator source (a `Closure` or `Choice`); the piped
// value is recorded in `call_args` under `span.start` as argument 0. Without
// this the closure/branch body's use of the piped value is dropped — an
// *understatement*, so resolving it is a soundness fix, not just precision.
pub type DirectPipeOp {
  DirectPipeOp(value: ArgumentValue, span: Span)
}

// Effect annotation for a type's field (e.g., `type Handler.on_click : [Dom]`).
//
// `module` is `Some(...)` when the annotation comes from a spec file (one
// file per package, qualified names like `myapp.Handler.on_click`) and
// `None` when it comes from a per-module cache file (bare names scoped to
// the file's module by location).
pub type TypeFieldAnnotation {
  TypeFieldAnnotation(
    module: Option(String),
    type_name: String,
    field: String,
    effects: EffectTerm,
  )
}

// A type field's resolved effect in the knowledge base. `effects` is the field
// call's effect set. When the field was inferred from a constructor site that
// wired an effect-polymorphic function, `bounds` and `source` carry that
// function's parameter bounds and qualified name, so a field call can bind the
// effect variables to its arguments (the same substitution resolved calls do).
// Both are empty/`None` for hand-written annotations and concrete field values.
pub type TypeFieldEffect {
  TypeFieldEffect(
    effects: EffectTerm,
    bounds: List(ParamBound),
    source: Option(QualifiedName),
  )
}

// Whether an external targets a whole module or a specific function.
pub type ExternalTarget {
  // `external effects gleam/list : []` — the entire module is pure.
  ModuleExternal
  // `external effects gleam/httpc.send : [Http]` — a specific function.
  FunctionExternal(name: String)
}

// Effect declaration for an external function (e.g., `external effects gleam/httpc.send : [Http]`).
pub type ExternalAnnotation {
  ExternalAnnotation(module: String, target: ExternalTarget, effects: EffectSet)
}

// A single effect violation: an annotated function called something
// that exceeds its declared effect budget.
pub type Violation {
  Violation(
    function: String,
    call: QualifiedName,
    span: Span,
    declared: EffectSet,
    actual: EffectSet,
  )
}

// A warning surfaced during checking.
pub type Warning {
  // A function reference passed as a value whose effects won't be tracked
  // through the callee.
  UntrackedEffectWarning(
    function: String,
    reference: QualifiedName,
    span: Span,
    effects: EffectSet,
  )
  // A field bound (`check f(recv.field: [..])`) whose `recv.field` path matched
  // no field call in the function's body. Such a bound resolves nothing and is
  // silently dead, so it's flagged. `receiver_is_param` distinguishes the cause:
  // when the receiver is a parameter it can't be traced to a construction site,
  // so a missing field call is a genuine typo; when it isn't, the field call may
  // exist but have resolved through value provenance, shadowing the bound.
  UnmatchedFieldBoundWarning(
    function: String,
    field_path: String,
    receiver_is_param: Bool,
  )
  // A plain parameter bound (`check f(g: [..])`) whose name matches no declared
  // parameter of the function — a typo. Matched on parameter *existence*, not
  // call presence: a callback forwarded but never called directly is still a
  // real parameter and isn't flagged, since its bound stays load-bearing during
  // substitution.
  UnmatchedParamBoundWarning(function: String, param: String)
}

// Result of checking one file.
pub type CheckResult {
  CheckResult(
    file: String,
    violations: List(Violation),
    warnings: List(Warning),
  )
}