(************************************************************************)
(*         *      The Rocq Prover / The Rocq Development Team           *)
(*  v      *         Copyright INRIA, CNRS and contributors             *)
(* <O___,, * (see version control and CREDITS file for authors & dates) *)
(*   \VV/  **************************************************************)
(*    //   *    This file is distributed under the terms of the         *)
(*         *     GNU Lesser General Public License Version 2.1          *)
(*         *     (see LICENSE file for the text of the license)         *)
(************************************************************************)

open Util
open Names
open Constr
open Term
open Declarations
open UVars
open Cooking

module NamedDecl = Context.Named.Declaration

let lift_univs info univ_hyps = function
  | Monomorphic ->
    assert (UVars.Instance.is_empty univ_hyps);
    info, univ_hyps, Monomorphic
  | Polymorphic auctx ->
    let info, (qn, un), auctx = lift_poly_univs info auctx in
    let univ_hyps =
      let open UVars.Instance in
      let qs, us = to_array univ_hyps in
      let qs = Array.sub qs 0 (Array.length qs - qn) in
      let us = Array.sub us 0 (Array.length us - un) in
      of_array (qs,us)
    in
    info, univ_hyps, Polymorphic auctx

(********************************)
(* Discharging opaque proof terms *)

let lift_private_univs info = function
  | Opaqueproof.PrivateMonomorphic () as priv ->
    let () = lift_private_mono_univs info () in
    priv
  | Opaqueproof.PrivatePolymorphic ctx ->
    let ctx = lift_private_poly_univs info ctx in
    Opaqueproof.PrivatePolymorphic ctx

let cook_opaque_proofterm info c =
  let fold info (c, priv) =
    let priv = lift_private_univs info priv in
    let c = abstract_as_body (create_cache info) c in
    (c, priv)
  in
  List.fold_right fold info c

(********************************)
(* Discharging constant         *)

let cook_constant env info cb =
  (* Adjust the info so that it is meaningful under the block of quantified universe binders *)
  let info, univ_hyps, univs = lift_univs info cb.const_univ_hyps cb.const_universes in
  let cache = create_cache info in
  let map c = abstract_as_body cache c in
  let body = match cb.const_body with
  | Undef _ as x -> x
  | Def cs -> Def (map cs)
  | OpaqueDef o ->
    OpaqueDef (Opaqueproof.discharge_opaque info o)
  | Primitive _ -> CErrors.anomaly (Pp.str "Primitives cannot be cooked")
  | Symbol _ -> CErrors.anomaly (Pp.str "Symbols cannot be cooked")
  in
  let typ = abstract_as_type cache cb.const_type in
  let names = names_info info in
  let hyps = List.filter (fun d -> not (Id.Set.mem (NamedDecl.get_id d) names)) cb.const_hyps in
  {
    const_hyps = hyps;
    const_univ_hyps = univ_hyps;
    const_body = body;
    const_type = typ;
    const_body_code = ();
    const_universes = univs;
    const_relevance =  lift_relevance info cb.const_relevance;
    const_inline_code = cb.const_inline_code;
    const_typing_flags = cb.const_typing_flags;
    const_def_height = Some (Environ.constant_definitional_height env body);
  }

(********************************)
(* Discharging mutual inductive *)

let cook_rel_context cache ctx =
  (* Dealing with substitutions between contexts is too annoying, so
     we reify [ctx] into a big [forall] term and work on that. *)
  let t = it_mkProd_or_LetIn mkProp ctx in
  let t = abstract_as_type cache t in
  let ctx, t = decompose_prod_decls t in
  assert (Constr.equal t mkProp);
  ctx

let cook_projection cache ~params t =
  let t = mkArrowR mkProp t in (* dummy type standing in for the inductive *)
  let t = it_mkProd_or_LetIn t params in
  let t = abstract_as_type cache t in
  let nrels = Context.Rel.nhyps (rel_context_of_cooking_cache cache) in
  let _, t = decompose_prod_n_decls (Context.Rel.length params + 1 + nrels) t in
  t

let lift_squashed info = let open Declarations in function
  | AlwaysSquashed -> AlwaysSquashed
  | SometimesSquashed s ->
    let s = Sorts.Quality.Set.fold (fun x acc ->
        let x = lift_quality info x in
        Sorts.Quality.Set.add x acc)
        s Sorts.Quality.Set.empty
    in
    SometimesSquashed s

let cook_one_ind info cache ~params mip =
  let mind_user_arity = abstract_as_type cache mip.mind_user_arity in
  let mind_sort = abstract_as_sort cache mip.mind_sort in
  let mind_arity_ctxt = cook_rel_context cache mip.mind_arity_ctxt in
  let mind_user_lc = Array.map (abstract_as_type cache) mip.mind_user_lc in
  let mind_nf_lc = Array.map (fun (ctx,t) ->
      let lc = it_mkProd_or_LetIn t ctx in
      let lc = abstract_as_type cache lc in
      decompose_prod_decls lc)
      mip.mind_nf_lc
  in
  let mind_record = match mip.mind_record with
    | NotRecord -> NotRecord
    | FakeRecord -> FakeRecord
    | PrimRecord ({ relevances; tys; _ } as pinfo) ->
        let tys = Array.map (cook_projection cache ~params) tys in
        let relevances = Array.map (lift_relevance info) relevances in
        PrimRecord { pinfo with relevances ; tys }
  in
  let squashed = Option.map (lift_squashed info) mip.mind_squashed in

  { mind_typename = mip.mind_typename;
    mind_record;
    mind_arity_ctxt;
    mind_user_arity;
    mind_sort;
    mind_consnames = mip.mind_consnames;
    mind_user_lc;
    mind_nrealargs = mip.mind_nrealargs;
    mind_nrealdecls = mip.mind_nrealdecls;
    mind_squashed = squashed;
    mind_nf_lc;
    mind_consnrealargs = mip.mind_consnrealargs;
    mind_consnrealdecls = mip.mind_consnrealdecls;
    mind_automaton = mip.mind_automaton;
    mind_relevance = lift_relevance info mip.mind_relevance;
    mind_relies_on_indices_not_mattering = mip.mind_relies_on_indices_not_mattering;
    mind_nb_constant = mip.mind_nb_constant;
    mind_nb_args = mip.mind_nb_args;
    mind_reloc_tbl = mip.mind_reloc_tbl;
  }

let cook_inductive info mib =
  let info, univ_hyps, mind_universes = lift_univs info mib.mind_univ_hyps mib.mind_universes in
  let cache = create_cache info in
  let nnewparams = Context.Rel.nhyps (rel_context_of_cooking_cache cache) in
  let mind_params_ctxt = cook_rel_context cache mib.mind_params_ctxt in
  let mind_packets = Array.map (cook_one_ind info cache ~params:mib.mind_params_ctxt) mib.mind_packets in
  let names = names_info info in
  let mind_hyps =
    List.filter (fun d -> not (Id.Set.mem (NamedDecl.get_id d) names))
      mib.mind_hyps
  in
  let mind_variance, mind_sec_variance =
    match mib.mind_variance, mib.mind_sec_variance with
    | None, None -> None, None
    | None, Some _ | Some _, None -> assert false
    | Some variance, Some sec_variance ->
      (* no variance for qualities *)
      let ulen  = snd (AbstractContext.size (universe_context_of_cooking_info info)) in
      let sec_variance, newvariance =
        Array.chop (Array.length sec_variance - ulen) sec_variance
      in
      Some (Array.append newvariance variance), Some sec_variance
  in
  let mind_template = match mib.mind_template with
  | None -> None
  | Some {template_param_arguments=levels; template_context; template_concl; template_defaults;} ->
      let sec_levels = List.make (Context.Rel.nhyps (rel_context_of_cooking_cache cache)) None in
      let levels = List.rev_append sec_levels levels in
      Some {template_param_arguments=levels; template_context; template_concl; template_defaults}
  in
  {
    mind_packets;
    mind_finite = mib.mind_finite;
    mind_hyps;
    mind_univ_hyps = univ_hyps;
    mind_nparams = mib.mind_nparams + nnewparams;
    mind_nparams_rec = mib.mind_nparams_rec + nnewparams;
    mind_params_ctxt;
    mind_universes;
    mind_template;
    mind_variance;
    mind_sec_variance;
    mind_private = mib.mind_private;
    mind_typing_flags = mib.mind_typing_flags;
  }

let cook_rel_context info ctx =
  cook_rel_context (create_cache info) ctx