type outcome = Solved of Pattern.Subst.t list | Stuck

let forced_label name =
  match Label.of_string name with Ok l -> l | Error _ -> assert false

let value_l = forced_label "value"
let left_l = forced_label "left"
let right_l = forced_label "right"
let result_l = forced_label "result"
let lower_l = forced_label "lower"
let upper_l = forced_label "upper"

let names =
  [
    "True";
    "Integer";
    "Natural";
    "Rational";
    "String";
    "Equal";
    "Plus";
    "Times";
    "LessThan";
    "LessThanEqual";
    "GreaterThan";
    "GreaterEqual";
    "Between";
  ]

let is_builtin name = List.mem name names

let fail = Solved []
let yes subst = Solved [ subst ]

(* Numeric view: integers and rationals compute together in Q, and
   results come back through Term.rational, which restores the
   minimal shape. An integer-valued result is an Integer. *)
let to_q = function
  | Pattern.Integer z -> Some (Q.of_bigint z)
  | Pattern.Rational q -> Some q
  | _ -> None

let of_q q =
  match Term.rational ~num:(Q.num q) ~den:(Q.den q) with
  | Ok term -> Pattern.of_term term
  | Error `Zero_denominator -> assert false

let arg subst bindings label =
  Option.map (Pattern.Subst.resolve subst) (Label.Map.find_opt label bindings)

(* Domain membership. Unbound arguments are Stuck, never
   enumerated: membership is not generation. The inclusions are the
   literal ones of decision 12. *)
let domain_check subst bindings check =
  match arg subst bindings value_l with
  | None -> Stuck
  | Some pattern ->
      if Pattern.is_ground pattern then if check pattern then yes subst else fail
      else Stuck

let comparison subst bindings compare =
  match (arg subst bindings left_l, arg subst bindings right_l) with
  | Some left, Some right -> (
      match (to_q left, to_q right) with
      | Some l, Some r -> if compare (Q.compare l r) 0 then yes subst else fail
      | _ ->
          if Pattern.is_ground left && Pattern.is_ground right then fail
          else Stuck)
  | _ -> Stuck

(* Plus and Times solve backwards: any two ground numeric arguments
   determine the third by unification. Times through zero cannot
   determine the missing factor and stays Stuck rather than guess. *)
let arithmetic subst bindings ~forward ~from_result_left ~from_result_right =
  let left = arg subst bindings left_l
  and right = arg subst bindings right_l
  and result = Label.Map.find_opt result_l bindings in
  let unify_result value =
    match result with
    | Some pattern -> (
        match Pattern.Subst.unify subst pattern value with
        | Some subst -> yes subst
        | None -> fail)
    | None -> yes subst
  in
  let solve_side binding solved =
    match binding with
    | Some pattern -> (
        match solved with
        | Some value -> (
            match Pattern.Subst.unify subst pattern value with
            | Some subst -> yes subst
            | None -> fail)
        | None -> Stuck)
    | None -> Stuck
  in
  match (Option.bind left to_q, Option.bind right to_q) with
  | Some l, Some r -> unify_result (of_q (forward l r))
  | Some l, None -> (
      match Option.bind result (fun p -> to_q (Pattern.Subst.resolve subst p)) with
      | Some res ->
          solve_side right (Option.map of_q (from_result_left res l))
      | None -> Stuck)
  | None, Some r -> (
      match Option.bind result (fun p -> to_q (Pattern.Subst.resolve subst p)) with
      | Some res ->
          solve_side left (Option.map of_q (from_result_right res r))
      | None -> Stuck)
  | None, None -> Stuck

let between subst bindings =
  let value = Label.Map.find_opt value_l bindings
  and lower = arg subst bindings lower_l
  and upper = arg subst bindings upper_l in
  match (lower, upper) with
  | Some (Pattern.Integer lo), Some (Pattern.Integer hi) -> (
      match Option.map (Pattern.Subst.resolve subst) value with
      | Some (Pattern.Integer v) ->
          if Z.geq v lo && Z.leq v hi then yes subst else fail
      | Some (Pattern.Var _) | None ->
          (* Bounded generation: the one place enumeration is
             admitted, because the range is finite by construction. *)
          let rec enumerate v acc =
            if Z.gt v hi then List.rev acc
            else
              let acc =
                match value with
                | Some pattern -> (
                    match
                      Pattern.Subst.unify subst pattern (Pattern.Integer v)
                    with
                    | Some subst -> subst :: acc
                    | None -> acc)
                | None -> subst :: acc
              in
              enumerate (Z.succ v) acc
          in
          Solved (enumerate lo [])
      | Some _ -> fail)
  | Some p, Some q when Pattern.is_ground p && Pattern.is_ground q -> fail
  | _ -> Stuck

let run subst name bindings =
  match name with
  | "True" -> Some (yes subst)
  | "Integer" ->
      Some
        (domain_check subst bindings (function
          | Pattern.Integer _ -> true
          | _ -> false))
  | "Natural" ->
      Some
        (domain_check subst bindings (function
          | Pattern.Integer z -> Z.geq z Z.zero
          | _ -> false))
  | "Rational" ->
      Some
        (domain_check subst bindings (function
          | Pattern.Integer _ | Pattern.Rational _ -> true
          | _ -> false))
  | "String" ->
      Some
        (domain_check subst bindings (function
          | Pattern.String _ -> true
          | _ -> false))
  | "Equal" -> (
      match
        (Label.Map.find_opt left_l bindings, Label.Map.find_opt right_l bindings)
      with
      | Some left, Some right ->
          Some
            (match Pattern.Subst.unify subst left right with
            | Some subst -> yes subst
            | None -> fail)
      | _ -> Some Stuck)
  | "Plus" ->
      Some
        (arithmetic subst bindings ~forward:Q.add
           ~from_result_left:(fun res l -> Some (Q.sub res l))
           ~from_result_right:(fun res r -> Some (Q.sub res r)))
  | "Times" ->
      Some
        (arithmetic subst bindings ~forward:Q.mul
           ~from_result_left:(fun res l ->
             if Q.equal l Q.zero then None else Some (Q.div res l))
           ~from_result_right:(fun res r ->
             if Q.equal r Q.zero then None else Some (Q.div res r)))
  | "LessThan" -> Some (comparison subst bindings ( < ))
  | "LessThanEqual" -> Some (comparison subst bindings ( <= ))
  | "GreaterThan" -> Some (comparison subst bindings ( > ))
  | "GreaterEqual" -> Some (comparison subst bindings ( >= ))
  | "Between" -> Some (between subst bindings)
  | _ -> None