Require Import Coq.ZArith.ZArith.
Require Import Coq.Lists.List.
Require Import Coq.Strings.String.
Require Crypto.Util.Strings.String.
Require Import Crypto.Util.Strings.Decimal.
Require Import Crypto.Util.Strings.HexString.
Require Import Crypto.Util.Option.
Require Import Crypto.Util.Strings.Show.
Require Import Crypto.Experiments.NewPipeline.Toplevel1.
Require Import Crypto.Experiments.NewPipeline.CStringification.
Import ListNotations. Local Open Scope Z_scope. Local Open Scope string_scope.

Import CStringification.Compilers.

Module ForExtraction.
  Definition parse_neg (s : string) : string * Z
    := match s with
       | String a b
         => if Ascii.ascii_dec a "-"
            then (b, -1)
            else if Ascii.ascii_dec a "+"
                 then (b, 1)
                 else (s, 1)
       | _ => (s, 1)
       end.
  Definition parse_N (s : string) : N
    := DecimalHelpers.N.of_uint (DecimalHelpers.String.to_uint s).
  Definition parse_Z (s : string) : Z
    := let '(s, sgn) := parse_neg s in
       sgn * Z.of_N (parse_N s).
  Definition parse_nat (s : string) : nat
    := N.to_nat (parse_N s).

  Definition parse_n (n : string) : nat
    := parse_nat n.

  Definition parse_pow (s : string) : option Z
    := let '(s, sgn) := parse_neg s in
       match String.split "^" s with
       | v::nil
         => Some (sgn * parse_Z v)
       | b::e::nil
         => Some (sgn * parse_Z b ^ parse_Z e)
       | _ => None
       end.

  Definition parse_mul (s : string) : option Z
    := List.fold_right
         (fun a b => (a <- a; b <- b; Some (Z.mul a b))%option)
         (Some 1)
         (List.map parse_pow (String.split "*" s)).

  (** We take in [c] in the format [a,b;c,d;e,f;...] becoming the list
    [[(a,b), (c,d), (e, f), ...]] *)
  Definition parse_s (s : string) : option Z
    := parse_mul s.
  Definition parse_c (s : string) : option (list (Z * Z))
    := List.fold_right
         (fun ls rest
          => (rest <- rest;
                match ls with
                | a::b::nil => (a <- parse_mul a; b <- parse_mul b; Some ((a, b)::rest))
                | _ => None
                end)%option)
         (Some nil)
         (List.map (String.split ",") (String.split ";" s)).

  Definition parse_machine_wordsize (s : string) : Z
    := parse_Z s.

  Local Open Scope string_scope.
  Local Notation NewLine := (String "010" "") (only parsing).

  Definition CollectErrors
             (res : list (string * Pipeline.ErrorT (list string)) + string)
    : list (list string) + list string
    := match res with
       | inl res
         => let header := hd "" (List.map (@fst _ _) res) in
            let res :=
                List.fold_right
                  (fun '(name, res) rest
                   => match res, rest with
                      | ErrorT.Error err, rest
                        => let cur := ("In " ++ name ++ ": " ++ show false err) in
                           let rest := match rest with inl _ => nil | inr rest => rest end in
                           inr (cur :: rest)
                      | ErrorT.Success v, inr ls => inr ls
                      | ErrorT.Success v, inl ls
                        => inl (v :: ls)
                      end)
                  (inl nil)
                  res in
            match res with
            | inl ls => inl ls
            | inr err => inr (header::err)
            end
       | inr res
         => inr (res::nil)
       end.

  Module UnsaturatedSolinas.
    Definition PipelineLines
               (n : string)
               (s : string)
               (c : string)
               (machine_wordsize : string)
    : list (string * Pipeline.ErrorT (list string)) + string
      := let str_n := n in
         let n : nat := parse_n n in
         let str_machine_wordsize := machine_wordsize in
         let str_c := c in
         let str_s := s in
         let machine_wordsize := parse_machine_wordsize machine_wordsize in
         match parse_s s, parse_c c with
         | None, None
           => inr ("Could not parse s (" ++ s ++ ") nor c (" ++ c ++ ")")
         | None, _
           => inr ("Could not parse s (" ++ s ++ ")")
         | _, None
           => inr ("Could not parse c (" ++ c ++ ")")
         | Some s, Some c
           => let header :=
                  ((["/* Autogenerated */";
                       "/* n = " ++ show false n ++ " (from """ ++ str_n ++ """) */";
                       "/* s = " ++ Hex.show_Z false s ++ " (from """ ++ str_s ++ """) */";
                       "/* c = " ++ show false c ++ " (from """ ++ str_c ++ """) */";
                       "/* machine_wordsize = " ++ show false machine_wordsize ++ " (from """ ++ str_machine_wordsize ++ """) */";
                       ""]%string)
                     ++ ToString.C.String.typedef_header
                     ++ [""])%list in
             inl
               ([("check_args" ++ NewLine ++ String.concat NewLine header,
                  UnsaturatedSolinas.check_args
                    n s c machine_wordsize
                    (ErrorT.Success header))%string]
                  ++ UnsaturatedSolinas.Synthesize n s c machine_wordsize "fe")%list
         end.

    Definition ProcessedLines
               (n : string)
               (s : string)
               (c : string)
               (machine_wordsize : string)
      : list string + string
      := match CollectErrors (PipelineLines n s c machine_wordsize) with
         | inl ls
           => inl
               (List.map (fun s => String.concat NewLine s ++ NewLine ++ NewLine)
                         ls)
         | inr ls
           => inr (String.concat
                    (NewLine ++ NewLine)
                    ls)
         end.

    Definition Pipeline
               {A}
               (n : string)
               (s : string)
               (c : string)
               (machine_wordsize : string)
               (success : list string -> A)
               (error : string -> A)
      : A
      := match ProcessedLines n s c machine_wordsize with
         | inl s => success s
         | inr s => error s
         end.

    Definition PipelineMain
               {A}
               (argv : list string)
               (success : list string -> A)
               (error : string -> A)
      : A
      := match argv with
         | _::n::s::c::machine_wordsize::nil
           => Pipeline
               n s c machine_wordsize
               success
               error
         | nil => error "empty argv"
         | prog::args
           => error ("Expected arguments n, s, c, machine_wordsize, got " ++ show false (List.length args) ++ " arguments in " ++ prog)
         end.
  End UnsaturatedSolinas.

  Module SaturatedSolinas.
    Definition PipelineLines
               (s : string)
               (c : string)
               (machine_wordsize : string)
    : list (string * Pipeline.ErrorT (list string)) + string
      := let str_machine_wordsize := machine_wordsize in
         let str_c := c in
         let str_s := s in
         let machine_wordsize := parse_machine_wordsize machine_wordsize in
         match parse_s s, parse_c c with
         | None, None
           => inr ("Could not parse s (" ++ s ++ ") nor c (" ++ c ++ ")")
         | None, _
           => inr ("Could not parse s (" ++ s ++ ")")
         | _, None
           => inr ("Could not parse c (" ++ c ++ ")")
         | Some s, Some c
           => let header :=
                  ((["/* Autogenerated */";
                       "/* s = " ++ Hex.show_Z false s ++ " (from """ ++ str_s ++ """) */";
                       "/* c = " ++ show false c ++ " (from """ ++ str_c ++ """) */";
                       "/* machine_wordsize = " ++ show false machine_wordsize ++ " (from """ ++ str_machine_wordsize ++ """) */";
                       ""]%string)
                     ++ ToString.C.String.typedef_header
                     ++ [""])%list in
              inl
                ([("check_args" ++ NewLine ++ String.concat NewLine header,
                   SaturatedSolinas.check_args
                     s c machine_wordsize
                     (ErrorT.Success header))%string]
                   ++ SaturatedSolinas.Synthesize s c machine_wordsize "fe")%list
         end.

    Definition ProcessedLines
               (s : string)
               (c : string)
               (machine_wordsize : string)
      : list string + string
      := match CollectErrors (PipelineLines s c machine_wordsize) with
         | inl ls
           => inl
               (List.map (fun s => String.concat NewLine s ++ NewLine ++ NewLine)
                         ls)
         | inr ls
           => inr (String.concat
                    (NewLine ++ NewLine)
                    ls)
         end.

    Definition Pipeline
               {A}
               (s : string)
               (c : string)
               (machine_wordsize : string)
               (success : list string -> A)
               (error : string -> A)
      : A
      := match ProcessedLines s c machine_wordsize with
         | inl s => success s
         | inr s => error s
         end.

    Definition PipelineMain
               {A}
               (argv : list string)
               (success : list string -> A)
               (error : string -> A)
      : A
      := match argv with
         | _::s::c::machine_wordsize::nil
           => Pipeline
               s c machine_wordsize
               success
               error
         | nil => error "empty argv"
         | prog::args
           => error ("Expected arguments s, c, machine_wordsize, got " ++ show false (List.length args) ++ " arguments in " ++ prog)
         end.
  End SaturatedSolinas.
End ForExtraction.