(** Constraint system for relational membership criteria.

    A relation's intension includes constraints that refine the space of
    admissible tuples. Since constraints are themselves relations, checking
    a constraint = checking membership in that constraint-relation.

    Negation requires an explicit universe relation name, forcing
    closed-world reasoning in first-order logic. *)

(* TODO: lt/lte/gt/gte/eq/neq/between at the bottom reference target names
   like "less_than" that don't match the prelude relation names
   (e.g. "natural_natural_less_than"). These constructors are currently dead
   code and will fail at runtime if used. *)

type relation_name = string
type attr_name = string

type binding_expr =
  | Var of attr_name
  | Const of Conventions.AbstractValue.t

module BindingMap = Map.Make (String)

type binding = binding_expr BindingMap.t

type t =
  | MemberOf of { target : relation_name; binding : binding }
  | Not of { body : t; universe : relation_name }
  | And of t list
  | Or of t list
  | Exists of { variable : attr_name; quantifier : relation_name; body : t }
  | Forall of { variable : attr_name; quantifier : relation_name; body : t }

type diagnostic =
  | MembershipFailed of {
      target : string;
      bound_values : (attr_name * Conventions.AbstractValue.t) list;
    }
  | RelationNotFound of string
  | UnboundedQuantifier of {
      variable : attr_name;
      quantifier : relation_name;
    }
  | ConstraintFailures of (string * diagnostic) list

let rec vars_in : t -> attr_name list = function
  | MemberOf { binding; _ } ->
    BindingMap.fold
      (fun _ expr acc ->
        match expr with
        | Var v -> if List.mem v acc then acc else v :: acc
        | Const _ -> acc)
      binding []
  | Not { body; _ } -> vars_in body
  | And cs | Or cs -> List.concat_map vars_in cs |> List.sort_uniq String.compare
  | Exists { variable; body; _ } | Forall { variable; body; _ } ->
    let inner = vars_in body in
    if List.mem variable inner then inner
    else variable :: inner

let rec rename_vars (renames : (attr_name * attr_name) list) : t -> t = function
  | MemberOf { target; binding } ->
    let binding' =
      BindingMap.fold
        (fun k expr acc ->
          let k' =
            match List.assoc_opt k renames with Some n -> n | None -> k
          in
          let expr' =
            match expr with
            | Var v -> (
              match List.assoc_opt v renames with
              | Some n -> Var n
              | None -> expr)
            | Const _ -> expr
          in
          BindingMap.add k' expr' acc)
        binding BindingMap.empty
    in
    MemberOf { target; binding = binding' }
  | Not { body; universe } ->
    Not { body = rename_vars renames body; universe }
  | And cs -> And (List.map (rename_vars renames) cs)
  | Or cs -> Or (List.map (rename_vars renames) cs)
  | Exists { variable; quantifier; body } ->
    let variable' =
      match List.assoc_opt variable renames with
      | Some n -> n
      | None -> variable
    in
    Exists { variable = variable'; quantifier; body = rename_vars renames body }
  | Forall { variable; quantifier; body } ->
    let variable' =
      match List.assoc_opt variable renames with
      | Some n -> n
      | None -> variable
    in
    Forall { variable = variable'; quantifier; body = rename_vars renames body }

(** Drop constraints referencing attributes not in the given set.
    Returns [None] if the constraint cannot be kept. *)
let rec filter_by_attrs (attrs : attr_name list) (c : t) : t option =
  let all_present c =
    List.for_all (fun v -> List.mem v attrs) (vars_in c)
  in
  match c with
  | MemberOf _ -> if all_present c then Some c else None
  | Not { body; universe } -> (
    match filter_by_attrs attrs body with
    | Some body' -> Some (Not { body = body'; universe })
    | None -> None)
  | And cs ->
    let kept = List.filter_map (filter_by_attrs attrs) cs in
    if kept = [] then None else Some (And kept)
  | Or cs ->
    (* All branches must survive for Or to remain sound *)
    let kept = List.filter_map (filter_by_attrs attrs) cs in
    if List.length kept = List.length cs then Some (Or kept) else None
  | Exists { variable; quantifier; body } -> (
    match filter_by_attrs (variable :: attrs) body with
    | Some body' -> Some (Exists { variable; quantifier; body = body' })
    | None -> None)
  | Forall { variable; quantifier; body } -> (
    match filter_by_attrs (variable :: attrs) body with
    | Some body' -> Some (Forall { variable; quantifier; body = body' })
    | None -> None)

let merge (cs1 : (string * t) list) (cs2 : (string * t) list)
    : (string * t) list =
  let combined = cs1 @ cs2 in
  let tbl = Hashtbl.create 16 in
  List.iter
    (fun (name, c) ->
      match Hashtbl.find_opt tbl name with
      | None -> Hashtbl.replace tbl name [ c ]
      | Some cs -> Hashtbl.replace tbl name (c :: cs))
    combined;
  Hashtbl.fold
    (fun name cs acc ->
      let merged = match cs with [ c ] -> c | _ -> And cs in
      (name, merged) :: acc)
    tbl []

let member_of ~target ~binding = MemberOf { target; binding }
let not_ ~universe body = Not { body; universe }
let and_ cs = match cs with [ c ] -> c | _ -> And cs
let or_  cs = match cs with [ c ] -> c | _ -> Or cs
let exists ~variable ~quantifier body = Exists { variable; quantifier; body }
let forall ~variable ~quantifier body = Forall { variable; quantifier; body }

type eval_context = {
  check_membership :
    relation_name ->
    (attr_name * Conventions.AbstractValue.t) list ->
    bool;
  iterate_finite :
    relation_name ->
    (attr_name * Conventions.AbstractValue.t) list list option;
}

let bind (b : binding) (tuple : Tuple.materialized)
    : (attr_name * Conventions.AbstractValue.t) list =
  BindingMap.fold
    (fun target_attr expr acc ->
      let value =
        match expr with
        | Const v -> Some v
        | Var src_attr -> (
          match Tuple.AttributeMap.find_opt src_attr tuple.attributes with
          | Some attr -> Some attr.Attribute.value
          | None -> None)
      in
      match value with
      | Some v -> (target_attr, v) :: acc
      | None -> acc)
    b []

let rec evaluate (ctx : eval_context) (tuple : Tuple.materialized) (c : t)
    : (bool, diagnostic) result =
  match c with
  | MemberOf { target; binding } ->
    let bound_values = bind binding tuple in
    if ctx.check_membership target bound_values then Ok true
    else Error (MembershipFailed { target; bound_values })
  | Not { body; universe = _ } ->
    (* The universe is a declarative annotation for closed-world reasoning.
       At runtime we negate the body; the universe enforces logical soundness
       at the schema level, not as a runtime tuple-existence check. *)
    (match evaluate ctx tuple body with
     | Ok true -> Ok false
     | Ok false -> Ok true
     | Error (MembershipFailed _) -> Ok true
     | Error d -> Error d)
  | And cs -> evaluate_and ctx tuple cs
  | Or cs -> evaluate_or ctx tuple cs
  | Exists { variable; quantifier; body } -> (
    match ctx.iterate_finite quantifier with
    | None -> Error (UnboundedQuantifier { variable; quantifier })
    | Some rows ->
      let rec check = function
        | [] -> Ok false
        | row :: rest ->
          let extended_tuple = extend_tuple tuple variable row in
          (match evaluate ctx extended_tuple body with
           | Ok true -> Ok true
           | Ok false -> check rest
           | Error _ -> check rest)
      in
      check rows)
  | Forall { variable; quantifier; body } -> (
    match ctx.iterate_finite quantifier with
    | None -> Error (UnboundedQuantifier { variable; quantifier })
    | Some rows ->
      let rec check = function
        | [] -> Ok true
        | row :: rest ->
          let extended_tuple = extend_tuple tuple variable row in
          (match evaluate ctx extended_tuple body with
           | Ok true -> check rest
           | Ok false -> Ok false
           | Error (MembershipFailed _) -> Ok false
           | Error d -> Error d)
      in
      check rows)

and evaluate_and ctx tuple = function
  | [] -> Ok true
  | c :: rest -> (
    match evaluate ctx tuple c with
    | Ok true -> evaluate_and ctx tuple rest
    | Ok false -> Ok false
    | Error d -> Error d)

and evaluate_or ctx tuple = function
  | [] -> Ok false
  | c :: rest -> (
    match evaluate ctx tuple c with
    | Ok true -> Ok true
    | Ok false -> evaluate_or ctx tuple rest
    | Error _ -> evaluate_or ctx tuple rest)

and extend_tuple (tuple : Tuple.materialized)
    (variable : attr_name)
    (row : (attr_name * Conventions.AbstractValue.t) list)
    : Tuple.materialized =
  (* TODO: Provenance for tuples at the relation level.
     Attributes from quantifier rows are namespaced by the variable name
     (e.g. "d.dept_id") to avoid silent overwriting of same-named attributes
     from the base tuple or other quantifiers. Without namespacing, the
     following FK constraint silently passes for an invalid insertion:
       - Employee { emp_id=1, dept_id=99 } is being inserted.
       - Constraint: Exists d in Department, MemberOf Department (dept_id = Var "dept_id").
       - Department contains only { dept_id=1, dept_name="Engineering" }.
       - extend_tuple overwrites dept_id=99 with dept_id=1 from the Department row.
       - Var "dept_id" resolves to 1, check_membership succeeds trivially.
       - The insertion is accepted even though dept_id=99 does not exist.
     Namespacing is a syntactic encoding of provenance: the variable prefix
     records which quantifier scope each value came from. A more principled
     solution would attach provenance metadata directly to each attribute in
     the relation, making the origin of every value explicit and queryable
     without relying on name conventions. This would also eliminate the
     ambiguity that arises when two quantifiers range over relations sharing
     attribute names — currently resolved by namespacing, but ideally enforced
     structurally by the type system or relation schema. *)
  let extra_attrs =
    List.fold_left
      (fun acc (k, v) ->
        Tuple.AttributeMap.add (variable ^ "." ^ k) { Attribute.value = v } acc)
      tuple.attributes row
  in
  { tuple with attributes = extra_attrs }

(* TODO: Sometimes it's more efficient and sufficient to simply halt at the first constraint violated. However, it's also important to be able to adjudicate which constraints failed on the tuple materialization. For that reason, we should likely return a lazy stream that computes the constraints on demand. *)
let evaluate_named (ctx : eval_context) (tuple : Tuple.materialized)
    (named : (string * t) list) : (bool, diagnostic) result =
  let failures =
    List.filter_map
      (fun (name, c) ->
        match evaluate ctx tuple c with
        | Ok true -> None
        | Ok false ->
          Some (name, MembershipFailed { target = name; bound_values = [] })
        | Error d -> Some (name, d))
      named
  in
  match failures with
  | [] -> Ok true
  | fs -> Error (ConstraintFailures fs)

let make_comparison_binding ~left ~right =
  BindingMap.empty
  |> BindingMap.add "left" left
  |> BindingMap.add "right" right

let lt ~left ~right =
  MemberOf { target = "less_than"; binding = make_comparison_binding ~left ~right }

let lte ~left ~right =
  MemberOf { target = "less_than_or_equal"; binding = make_comparison_binding ~left ~right }

let gt ~left ~right =
  MemberOf { target = "greater_than"; binding = make_comparison_binding ~left ~right }

let gte ~left ~right =
  MemberOf { target = "greater_than_or_equal"; binding = make_comparison_binding ~left ~right }

let eq ~left ~right =
  MemberOf { target = "equal"; binding = make_comparison_binding ~left ~right }

let neq ~left ~right =
  MemberOf { target = "not_equal"; binding = make_comparison_binding ~left ~right }

let between ~value ~low ~high =
  And
    [ MemberOf { target = "greater_than_or_equal";
                 binding = make_comparison_binding ~left:value ~right:low }
    ; MemberOf { target = "less_than_or_equal";
                 binding = make_comparison_binding ~left:value ~right:high }
    ]