open Types

type specimen = {
  ty : typ;
  term : Source.term;
}

type failure = {
  original : specimen;
  minimal : specimen;
  comparison : Relation.comparison;
}

type result = {
  total : int;
  related : int;
  failures : failure list;
}

let rng = Random.State.make [| 0x51; 0x17; 0x2b |]

let fresh =
  let c = ref 0 in
  fun prefix ->
    incr c;
    prefix ^ string_of_int !c

let pick xs = List.nth xs (Random.State.int rng (List.length xs))

let rec gen_typ depth =
  if depth = 0 then pick [TInt; TBool]
  else
    pick
      [
        TInt;
        TBool;
        TPair (gen_typ (depth - 1), gen_typ (depth - 1));
        TSum (gen_typ (depth - 1), gen_typ (depth - 1));
      ]

let rec gen_closed depth ty =
  if depth = 0 then
    match ty with
    | TInt -> Source.Int (Random.State.int rng 5)
    | TBool -> Source.Bool (Random.State.bool rng)
    | TPair (a, b) -> Source.Pair (gen_closed 0 a, gen_closed 0 b)
    | TSum (a, b) ->
        if Random.State.bool rng then Source.Inl (TSum (a, b), gen_closed 0 a)
        else Source.Inr (TSum (a, b), gen_closed 0 b)
    | _ -> Source.Int 0
  else
    match ty with
    | TInt ->
        pick
          [
            Source.Int (Random.State.int rng 7);
            Source.If (Source.Bool (Random.State.bool rng), Source.Int 1, Source.Int 0);
          ]
    | TBool ->
        pick
          [
            Source.Bool (Random.State.bool rng);
            Source.Eq (GBool, Source.Bool (Random.State.bool rng), Source.Bool (Random.State.bool rng));
          ]
    | TPair (a, b) -> Source.Pair (gen_closed (depth - 1) a, gen_closed (depth - 1) b)
    | TSum (a, b) ->
        if Random.State.bool rng then Source.Inl (TSum (a, b), gen_closed (depth - 1) a)
        else Source.Inr (TSum (a, b), gen_closed (depth - 1) b)
    | _ -> gen_closed 0 ty

let sample_terms ~count ~max_depth () =
  let adversarial =
    [
      {
        ty = TBool;
        term =
          Source.Let
            ( "poly_const_false",
              Source.TyLam ("a", Source.Lam ("x", TVar "a", TArrow (TVar "a", TBool), Source.Lam ("y", TVar "a", TBool, Source.Bool false))),
              Source.App
                ( Source.App (Source.TyApp (Source.Var "poly_const_false", TInt), Source.Int 1),
                  Source.Int 1 ) );
      };
      {
        ty = TInt;
        term =
          Source.Let
            ( "x",
              Source.Roll
                ( TMu ("a", TInt),
                  Source.Fix ("loop", TInt, Source.Var "loop") ),
              Source.Int 0 );
      };
      {
        ty = TInt;
        term = Source.Let ("_", Source.Tick ("alloc", Source.Int 0), Source.Int 1);
      };
    ]
  in
  let rec fill acc i attempts =
    if i = count then List.rev acc
    else if attempts > count * 32 then List.rev acc
    else if i < min count 20 then
      let specimen = List.nth adversarial (i mod List.length adversarial) in
      if Typecheck.is_well_typed specimen.ty specimen.term then fill (specimen :: acc) (i + 1) attempts
      else fill acc i (attempts + 1)
    else
      let ty = gen_typ max_depth in
      let specimen = { ty; term = gen_closed max_depth ty } in
      if Typecheck.is_well_typed specimen.ty specimen.term then fill (specimen :: acc) (i + 1) attempts
      else fill acc i (attempts + 1)
  in
  fill [] 0 0

let rec shrink_term ty term =
  let self = shrink_term in
  let typ_or fallback term =
    match Typecheck.type_of term with
    | Ok ty -> ty
    | Error _ -> fallback
  in
  let rebuild f sub_ty sub =
    List.map (fun sub' -> { ty; term = f sub'.term }) (self sub_ty sub)
  in
  match ty, term with
  | TInt, Source.Int n when n <> 0 -> [{ ty; term = Source.Int 0 }]
  | TBool, Source.Bool true -> [{ ty; term = Source.Bool false }]
  | _, Source.Tick (_, t) -> { ty; term = t } :: rebuild (fun t' -> Source.Tick ("shrunk", t')) ty t
  | _, Source.Let (_, a, b) ->
      { ty; term = b } ::
      rebuild (fun a' -> Source.Let ("_", a', b)) (typ_or ty a) a @
      rebuild (fun b' -> Source.Let ("_", a, b')) ty b
  | _, Source.If (c, t, e) ->
      { ty; term = t } :: { ty; term = e } ::
      rebuild (fun c' -> Source.If (c', t, e)) TBool c @
      rebuild (fun t' -> Source.If (c, t', e)) ty t @
      rebuild (fun e' -> Source.If (c, t, e')) ty e
  | TPair (a_ty, b_ty), Source.Pair (a, b) ->
      rebuild (fun a' -> Source.Pair (a', b)) a_ty a @
      rebuild (fun b' -> Source.Pair (a, b')) b_ty b
  | TSum (a_ty, _), Source.Inl (sum_ty, t) ->
      rebuild (fun t' -> Source.Inl (sum_ty, t')) a_ty t
  | TSum (_, b_ty), Source.Inr (sum_ty, t) ->
      rebuild (fun t' -> Source.Inr (sum_ty, t')) b_ty t
  | _, Source.App (f, a) ->
      let f_ty = typ_or (TArrow (ty, ty)) f in
      let a_ty = typ_or ty a in
      rebuild (fun f' -> Source.App (f', a)) f_ty f @
      rebuild (fun a' -> Source.App (f, a')) a_ty a
  | _, Source.TyApp (t, app_ty) ->
      let t_ty = typ_or (TForall ("a", ty)) t in
      rebuild (fun t' -> Source.TyApp (t', app_ty)) t_ty t
  | _, Source.Eq (g, a, b) ->
      let ground_ty = match g with GInt -> TInt | GBool -> TBool in
      rebuild (fun a' -> Source.Eq (g, a', b)) ground_ty a @
      rebuild (fun b' -> Source.Eq (g, a, b')) ground_ty b
  | _, Source.LetPair (x, y, scrut, body) ->
      let scrut_ty = typ_or (TPair (ty, ty)) scrut in
      rebuild (fun scrut' -> Source.LetPair (x, y, scrut', body)) scrut_ty scrut @
      rebuild (fun body' -> Source.LetPair (x, y, scrut, body')) ty body
  | _, Source.Case (scrut, (x, l), (y, r)) ->
      let scrut_ty = typ_or (TSum (ty, ty)) scrut in
      rebuild (fun scrut' -> Source.Case (scrut', (x, l), (y, r))) scrut_ty scrut @
      rebuild (fun l' -> Source.Case (scrut, (x, l'), (y, r))) ty l @
      rebuild (fun r' -> Source.Case (scrut, (x, l), (y, r'))) ty r
  | TMu (_, _), Source.Roll (roll_ty, t) ->
      let payload_ty = typ_or ty t in
      rebuild (fun t' -> Source.Roll (roll_ty, t')) payload_ty t
  | _, Source.Unroll t ->
      let t_ty = typ_or (TMu ("a", ty)) t in
      rebuild (fun t' -> Source.Unroll t') t_ty t
  | _ -> []

let shrink specimen =
  shrink_term specimen.ty specimen.term
  |> List.filter (fun candidate -> Typecheck.is_well_typed candidate.ty candidate.term)

let comparison_for ?(fuel = 256) flags relation specimen =
  let compiled = Pipeline.compile flags specimen.ty specimen.term in
  Relation.compare_programs ~fuel relation specimen.ty specimen.term compiled.target_program

let is_failure cmp =
  match cmp.Relation.verdict with
  | Relation.Related -> false
  | Relation.Unrelated _ -> true

let shrink_failure ?(fuel = 256) flags relation specimen =
  let rec improve current =
    let current_cmp = comparison_for ~fuel flags relation current in
    let candidates =
      shrink current
      |> List.filter_map (fun candidate ->
             let cmp = comparison_for ~fuel flags relation candidate in
             if is_failure cmp then Some (candidate, cmp) else None)
    in
    match candidates with
    | (candidate, _) :: _ when Source.string_of_term candidate.term <> Source.string_of_term current.term ->
        improve candidate
    | _ -> (current, current_cmp)
  in
  improve specimen

let run_campaign ?(fuel = 256) flags relation ~count ~max_depth () =
  let specimens = sample_terms ~count ~max_depth () in
  let failures, related =
    List.fold_left
      (fun (fails, ok) specimen ->
        let cmp = comparison_for ~fuel flags relation specimen in
        match cmp.verdict with
        | Relation.Related -> (fails, ok + 1)
        | Relation.Unrelated _ ->
            let minimal, minimal_cmp = shrink_failure ~fuel flags relation specimen in
            ({ original = specimen; minimal; comparison = minimal_cmp } :: fails, ok))
      ([], 0)
      specimens
  in
  { total = List.length specimens; related; failures = List.rev failures }