content: {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  -- | Keelung is a DSL for building zero-knowledge proofs
  module Keelung
    ( module Keelung.Syntax,
      module Keelung.Field,
      module Keelung.Heap,
      module Keelung.Monad,
      module Keelung.Data.Bits,
      module Keelung.CircuitFormat,
      keelung,
      -- interpret
      interpret,
      --interpretData,
      interpretEither,
      -- compile
      compile,
      compileO0,
      compileWithOpts,
      rtsoptProf,
      rtsoptMemory,
      -- solve R1CS
      solveOutput,
      --solveOutputData,
      solveOutputEither,
      -- witness generation
      witness,
      witness',
      -- proof generation
      prove,
      prove',
      -- proof verification
      verify,
      verify',
      -- genCircuit,
      genCircuitBin,
      genWtns,
      -- genCircuitDefault,
      genInputs,
      genInputsDefault,
      Encode,
      GaloisField,
      keelungVersion,
      proveData,
      interpretData,
      interpretDataEither,
      solveOutputDataEither,
      solveOutputData,
      stringToJsonInputs,
      jsonFileInputs,
      jsonInputs,
      SnarkjsBackend(..),
      setupSnarkjsPtau,
      checkWtnsBySnarkjs,
      genZkeyBySnarkjs,
      proveBySnarkjs,
    )
  where
  
  import Control.Monad.Except
  import Data.ByteString.Char8 qualified as BS
  import Data.Field.Galois (GaloisField)
  import Data.List (intercalate)
  import Data.Serialize (Serialize)
  import Data.Serialize qualified as Serialize
  import Data.String (IsString (fromString))
  import Keelung.Constraint.R1CS (R1CS)
  import Keelung.Data.Bits
  import Keelung.Error
  import Keelung.Field
  import Keelung.Heap
  import Keelung.Monad
  import Keelung.Options
  import Keelung.Syntax
  import Keelung.Syntax.Encode
  import Keelung.CircuitFormat
  import System.Directory qualified as Path
  --import System.FilePath qualified as FP
  import System.IO.Error qualified as IO
  import System.Info qualified
  import System.Process qualified as Process
  import Text.Read (readMaybe)
  import Data.Aeson
  import GHC.IO.Exception (ExitCode(..))
  
  -- | IMPORTANT: The compatibale compiler version of this library, Make sure it's updated and matched accordingly.
  keelungCompilerVersion :: (Int, Int, Int)
  keelungCompilerVersion = (0, 26, 1)
  
  -- | The version of this library in String
  keelungVersion :: String
  keelungVersion = (\(a,b,c) -> intercalate "." [show a, show b, show c]) keelungCompilerVersion
  
  --------------------------------------------------------------------------------
  
  -- | Entry point for the Keelung command line interface
  keelung :: (Encode t) => Comp t -> IO ()
  keelung program = do
    -- replace with beefier option parser
    options <- getOptions
    case options of
      Compile fieldType -> compile fieldType program >>= printResult
      Interpret fieldType publicInputs privateInputs -> interpret fieldType program publicInputs privateInputs >>= print
      Witness fieldType publicInputs privateInputs outputFilePath -> witness' outputFilePath fieldType program publicInputs privateInputs >>= print
      Prove fieldType publicInputs privateInputs circuitPath witnessPath proofPath ->
        prove' circuitPath witnessPath "aurora/parameter.json" proofPath fieldType program publicInputs privateInputs
      Verify circuitPath witnessPath proofPath ->
        verify' circuitPath witnessPath "aurora/parameter.json" proofPath
      Version -> putStrLn keelungVersion
    where
      printResult (Left err) = print err
      printResult (Right result) = print result
  
  --------------------------------------------------------------------------------
  
  -- | Compile a program to a 'R1CS' constraint system.
  compile :: (Encode t) => FieldType -> Comp t -> IO (Either Error (R1CS Integer))
  compile = compileWithOpts 1 [] []
  
  -- | Compile a program to a 'R1CS' constraint system with optimization level 0.
  compileO0 :: (Encode t) => FieldType -> Comp t -> IO (Either Error (R1CS Integer))
  compileO0 = compileWithOpts 0 [] []
  
  -- | Compile a program to a 'R1CS' constraint system with optimization level and RTS options as arguments.
  compileWithOpts :: (Encode t) => Int -> [String] -> [String] -> FieldType -> Comp t -> IO (Either Error (R1CS Integer))
  compileWithOpts level opts rtsopts fieldType prog = runM $ do
    elab <- liftEither (elaborateAndEncode prog)
    let options = "protocol" : optOptimize level : opts <> ["+RTS"] <> rtsopts <> ["-RTS"]
    callKeelungc options (fieldType, elab) :: M (R1CS Integer)
    where
      optOptimize :: Int -> String
      optOptimize i = "O" <> show i
  
  -- | Default RTS options for profiling
  rtsoptProf :: [String]
  rtsoptProf = ["-p"]
  
  -- | Helper function for compiling RTS options
  -- Memory size in GB for RTS options -M, -H and in MB for -A
  -- Try to increase if keelungc produces segmentation fault.
  -- https://downloads.haskell.org/ghc/latest/docs/users_guide/runtime_control.html
  rtsoptMemory :: Int -> Int -> Int -> [String]
  rtsoptMemory m h a = ["-M" <> show m <> "G", "-H" <> show h <> "G", "-A" <> show a <> "M"]
  
  --------------------------------------------------------------------------------
  
  stringToJsonInputs :: (FromJSON t, Inputable t) => (t -> t -> IO ()) -> BS.ByteString -> BS.ByteString -> IO ()
  stringToJsonInputs (f :: t -> t -> IO ()) a b = let a' = decodeStrict a :: Maybe t
                                                      b' = decodeStrict b :: Maybe t
                                                  in case (a', b') of
                                                   (Just i, Just j) -> f i j
                                                   (_, _) -> error "failed to decode JSON"
  
  jsonFileInputs :: (FromJSON t, Inputable t) => (t -> t -> IO ()) -> FilePath -> FilePath -> IO ()
  jsonFileInputs f a b = do a' <- readFile a
                            b' <- readFile b
                            stringToJsonInputs f (BS.pack a') (BS.pack b')
  
  -- Default method of JSON inputs
  jsonInputs :: (FromJSON t, Inputable t) => (t -> t -> IO ()) -> IO ()
  jsonInputs f = jsonFileInputs f "public.json" "private.json"
  
  -- | Generate a proof given circuit, inputs (witness), paratemer, and proof
  prove' ::
    (Encode t) =>
    FilePath ->
    FilePath ->
    FilePath ->
    FilePath ->
    FieldType ->
    Comp t ->
    [Integer] ->
    [Integer] ->
    IO ()
  prove' circuitPath witnessPath paramPath proofPath fieldType prog publicInput privateInput = do
    result <- runM $ do
      (cmd, args) <- findAuroraProver
      _ <- genCircuit circuitPath fieldType prog
      _ <- genWitness_ witnessPath fieldType prog publicInput privateInput -- Should generate public as well as private inputs
      -- _ <- genWitness_ witnessPath fieldType prog publicInput privateInput -- Should generate public as well as private inputs
      genParameters paramPath
      -- genInputs inputs publicInput -- Should generate public inputs only for verifier
      lift $ do
        let arguments =
              args
                ++ [ "--r1cs_filepath",
                     circuitPath,
                     "--input_filepath",
                     witnessPath,
                     "--parameter_filepath",
                     paramPath,
                     "--output_filepath",
                     proofPath
                   ]
        -- print $ "Running: " ++ cmd ++ " " ++ unwords arguments
        msg <- Process.readProcess cmd arguments mempty
        return (proofPath, msg)
    case result of
      Left err -> print err
      Right (_, msg) -> putStr msg
  
  prove :: (Encode t) => FieldType -> Comp t -> [Integer] -> [Integer] -> IO ()
  prove f p i o = do
    Path.createDirectoryIfMissing True "aurora"
    prove' "aurora/circuit.jsonl" "aurora/witness.jsonl" "aurora/parameter.json" "aurora/proof" f p i o
  
  proveData :: (Encode t, Inputable a1, Inputable a2) => FieldType -> Comp t -> a1 -> a2 -> IO ()
  proveData f p i o = prove f p (toInts i) (toInts o)
  
  -- | Generate and verify a proof given circuit, witness, paratemer, and proof
  verify' :: FilePath -> FilePath -> FilePath -> FilePath -> IO ()
  verify' circuitPath witnessPath paramPath proofPath = do
    result <- runM $ do
      (cmd, args) <- findAuroraVerifier
      genParameters proofPath
      lift $ do
        let arguments =
              args
                ++ [ "--r1cs_filepath",
                     circuitPath,
                     "--input_filepath",
                     witnessPath,
                     "--parameter_filepath",
                     paramPath,
                     "--proof_filepath",
                     proofPath
                   ]
        Process.readProcess cmd arguments mempty
    case result of
      Left err -> print err
      Right msg -> putStr msg
  
  -- TODO: Verify inputs.jsonl instead
  verify :: IO ()
  verify = verify' "aurora/circuit.jsonl" "aurora/witness.jsonl" "aurora/parameter.json" "aurora/proof"
  
  -- | Compile a program as R1CS and write it to circuit.jsonl.
  genCircuit :: (Encode t) => FilePath -> FieldType -> Comp t -> M (R1CS Integer)
  genCircuit filePath fieldType prog = do
    elab <- liftEither (elaborateAndEncode prog)
    r1cs <- callKeelungc ["protocol", "toJSON", "--filepath", filePath] (fieldType, elab) :: M (R1CS Integer)
    liftIO $ putStrLn $ "Generated circuit file at: " <> filePath
    return r1cs
  
  genCircuitBin :: (Encode t) => FilePath -> FieldType -> Comp t -> IO (Either Error String)
  genCircuitBin filePath fieldType prog = runM $ do
    elab <- liftEither (elaborateAndEncode prog)
    _ <- callKeelungc ["protocol", "genCircuitBin", "--filepath", filePath] (fieldType, elab) :: M (R1CS Integer)
    liftIO $ (\s -> putStrLn $ "Generated binary circuit file at: " <> s) =<< Path.makeAbsolute filePath
    return "Success"
  
  -- genCircuitDefault :: Encode t => FieldType -> Comp t -> M (R1CS Integer)
  -- genCircuitDefault = genCircuit "aurora/circuit.jsonl"
  
  -- | Generate witnesses for a program with inputs and write them to witness.jsonl.
  genWitness_ :: (Encode t) => FilePath -> FieldType -> Comp t -> [Integer] -> [Integer] -> M [Integer]
  genWitness_ filePath fieldType prog publicInput privateInput = do
    elab <- liftEither (elaborateAndEncode prog)
    output <- callKeelungc ["protocol", "genWitness", "--filepath", filePath] (fieldType, elab, publicInput, privateInput)
    liftIO $ putStrLn $ "Generated witness file at: " <> filePath
    return output
  
  -- | Generate witnesses for a program with inputs and write them to witness.jsonl.
  genWtns :: (Encode t) => FilePath -> FieldType -> Comp t -> [Integer] -> [Integer] -> IO (Either Error String)
  genWtns filePath fieldType prog publicInput privateInput = runM $ do
    elab <- liftEither (elaborateAndEncode prog)
    _ <- callKeelungc ["protocol", "genWtns", "--filepath", filePath] (fieldType, elab, publicInput, privateInput) :: M [Integer]
    liftIO $ (\s -> putStrLn $ "Generated wtns file at: " <> s) =<< Path.makeAbsolute filePath
    return "Success"
  
  -- | Generate parameters for a program and write them to parameter.json.
  genParameters :: FilePath -> M ()
  genParameters filePath = do
    lift $ BS.writeFile filePath "{\"security_level\": 128, \"heuristic_ldt_reducer_soundness\": true, \"heuristic_fri_soundness\": true, \"bcs_hash_type\": \"blake2b_type\", \"make_zk\": false, \"parallel\": true, \"field_size\": 181, \"is_multiplicative\": true}"
    liftIO $ putStrLn $ "Generated parameter file at: " <> filePath
  
  -- | For generating witness.jsonl
  witness' :: (Encode t) => FilePath -> FieldType -> Comp t -> [Integer] -> [Integer] -> IO [Integer]
  witness' fp fieldType prog publicInput privateInput = runM (genWitness_ fp fieldType prog publicInput privateInput) >>= printErrorInstead
  
  witness :: (Encode t) => FieldType -> Comp t -> [Integer] -> [Integer] -> IO [Integer]
  witness = witness' "aurora/witness.jsonl"
  
  -- | For generating inputs.jsonl
  genInputs :: FilePath -> [Integer] -> M ()
  genInputs fp inputs = do
    let inputs' = intercalate "," $ map ((\x -> "\"" ++ x ++ "\"") . show) inputs
    lift $ BS.writeFile fp $ fromString $ "{\"inputs\":[" ++ inputs' ++ "]}"
  
  genInputsDefault :: [Integer] -> M ()
  genInputsDefault = genInputs "inputs.jsonl"
  
  --------------------------------------------------------------------------------
  
  printErrorInstead :: (Show e) => Either e [a] -> IO [a]
  printErrorInstead (Left err) = do
    print err
    return []
  printErrorInstead (Right values) = return values
  
  -- | Interpret a program with public and private inputs
  interpret :: (Encode t) => FieldType -> Comp t -> [Integer] -> [Integer] -> IO [Integer]
  interpret fieldType prog publicInput privateInput = interpretEither fieldType prog publicInput privateInput >>= printErrorInstead
  
  -- | Interpret a program where public and private inputs are provided as datatype encodable to JSON.
  interpretData :: (Encode t, Inputable a, Inputable b) => FieldType -> Comp t -> a -> b -> IO [Integer]
  interpretData fieldType prog pub prv = interpret fieldType prog (toInts pub) (toInts prv)
  
  -- | Interpret a program with public and private inputs
  interpretEither :: (Encode t) => FieldType -> Comp t -> [Integer] -> [Integer] -> IO (Either Error [Integer])
  interpretEither fieldType prog publicInput privateInput =
    runM
      ( do
          elab <- liftEither (elaborateAndEncode prog)
          callKeelungc ["protocol", "interpret"] (fieldType, elab, publicInput, privateInput)
      )
  interpretDataEither :: (Encode t, Inputable a, Inputable b) => FieldType -> Comp t -> a -> b -> IO (Either Error [Integer])
  interpretDataEither fieldType prog publicInput privateInput = interpretEither fieldType prog (toInts publicInput) (toInts privateInput)
  
  --------------------------------------------------------------------------------
  
  -- | Solves the R1CS of a Keelung program with given inputs and outputs the result
  solveOutput :: (Encode t) => FieldType -> Comp t -> [Integer] -> [Integer] -> IO [Integer]
  solveOutput fieldType prog publicInput privateInput = solveOutputEither fieldType prog publicInput privateInput >>= printErrorInstead
  
  solveOutputData :: (Encode t, Inputable d, Inputable e) => FieldType -> Comp t -> d -> e -> IO [Integer]
  solveOutputData fieldType prog pub prv = solveOutput fieldType prog (toInts pub) (toInts prv)
  
  -- | Solves the R1CS of a Keelung program with given inputs and outputs the result
  solveOutputEither :: (Encode t) => FieldType -> Comp t -> [Integer] -> [Integer] -> IO (Either Error [Integer])
  solveOutputEither fieldType prog publicInput privateInput =
    runM
      ( do
          elab <- liftEither (elaborateAndEncode prog)
          callKeelungc ["protocol", "solve"] (fieldType, elab, publicInput, privateInput)
      )
  
  solveOutputDataEither :: (Encode t, Inputable d, Inputable e) => FieldType -> Comp t -> d -> e -> IO (Either Error [Integer])
  solveOutputDataEither fieldType prog pub prv = solveOutputEither fieldType prog (toInts pub) (toInts prv)
  
  setupSnarkjsPtau :: Integer -> String -> IO (Either Error String)
  setupSnarkjsPtau size entropy = runM $ do
    let p = show size
    lift $ Path.createDirectoryIfMissing True "snarkjs"
    fp <- lift $ Path.makeAbsolute "snarkjs"
    lift $ putStrLn "These generated files are for testings only and not production-ready.\nGenerating powersoftau files..."
    let pot :: String -> String -> String
        pot s i = fp <> "/pot" <> s <> "_" <> i <> ".ptau"
    _ <- putStrLn <$> callSnarkjs ["powersoftau", "new", "bn128", p, pot p "0"] 
    _ <- putStrLn <$> callSnarkjs ["powersoftau", "contribute", pot p "0", pot p "1", "-n=\"First Contribution\"", "-e=" <> entropy]
    _ <- putStrLn <$> callSnarkjs ["powersoftau", "beacon", pot p "1", pot p "beacon", "0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f", "10", "-n\"Beacon\"", "-e=" <> entropy]
    _ <- putStrLn <$> callSnarkjs ["powersoftau", "prepare", "phase2", pot p "beacon", pot p "final"]
    return $ "Generated .ptau file: " <> pot p "final"
    
  checkWtnsBySnarkjs :: FilePath -> FilePath -> IO (Either Error String)
  checkWtnsBySnarkjs r1cs wtns = runM $ callSnarkjs ["wtns", "check", r1cs, wtns] >> return "Success"
  
  data SnarkjsBackend = Groth16 | PLONK
  
  genZkeyBySnarkjs :: SnarkjsBackend -> FilePath -> FilePath -> IO (Either Error String)
  genZkeyBySnarkjs b r1cs ptau = runM $ do
    let backend = case b of Groth16 -> "groth16";  PLONK -> "plonk";
    fp <- lift $ Path.makeAbsolute r1cs
    let parent = take (length fp - length r1cs) fp
        name = take (length r1cs - 5) r1cs
    _ <- callSnarkjs [backend, "setup", r1cs, ptau, parent <> name <> ".zkey"]
    return $ "Generated .zkey file: " <> parent <> name <> ".zkey"
                               
  proveBySnarkjs :: SnarkjsBackend -> String -> String -> IO (Either Error String)
  proveBySnarkjs b zkey wtns = runM $ do
    let backend = case b of Groth16 -> "groth16";  PLONK -> "plonk";
    callSnarkjs [backend, "prove", zkey, wtns]
  
  --------------------------------------------------------------------------------
  
  -- | Internal function for handling data serialization
  callKeelungc :: (Serialize a, Serialize b) => [String] -> a -> M b
  callKeelungc args' payload = do
    (cmd, args) <- findKeelungc
    version <- readKeelungVersion cmd args
    checkCompilerVersion version
    blob <- lift $ Process.readProcess cmd (args ++ args') (BS.unpack $ Serialize.encode payload)
    let result = Serialize.decode (BS.pack blob)
    case result of
      Left err -> throwError (DecodeError err)
      Right (Left err) -> throwError (CompileError err)
      Right (Right x) -> return x
  
  callSnarkjs :: [String] -> M String
  callSnarkjs args' = do
    (cmd, args) <- findSnarkjs
    -- TODO: check for Snarkjs' version
    (code, out, err) <- lift $ Process.readProcessWithExitCode cmd (args ++ args') [] 
    if code == ExitSuccess || code == ExitFailure 99 then
      return out
    else
      throwError (SnarkjsError (err <> "\n" <> out))
  
  -- | Locate the Keelung compiler
  --      1. see if "keelungc" is in PATH
  --      2. if not, try to run "docker run banacorn/keelung"
  --   Returns the command and arguments to run when found
  findKeelungc :: M (String, [String])
  findKeelungc = do
    keelungcExists <- checkCmd "keelungc"
    if keelungcExists
      then return ("keelungc", [])
      else do
        dockerExists <- checkCmd "docker"
        if dockerExists
          then do
            lift $ Path.createDirectoryIfMissing True "aurora"
            filepath <- lift $ Path.makeAbsolute "aurora"
            case System.Info.arch of
              "x86_64" -> return ("docker", ["run", "-i", "--volume", filepath ++ ":/aurora", "btqag/keelungc"])
              -- insert "--platform=linux/amd64" when we are not on a x86 machine
              _ -> return ("docker", ["run", "-i", "--platform=linux/amd64", "--volume", filepath ++ ":/aurora", "btqag/keelungc"])
          else throwError CannotLocateKeelungC
  
  findAuroraProver :: M (String, [String])
  findAuroraProver = do
    auroraExists <- checkCmd "aurora_prove"
    if auroraExists
      then return ("aurora_prove", [])
      else do
        dockerExists <- checkCmd "docker"
        -- let currentDirector = Path.getCurrentDirectory
        if dockerExists
          then do
            lift $ Path.createDirectoryIfMissing True "aurora"
            filepath <- lift $ Path.makeAbsolute "aurora"
            case System.Info.arch of
              "x86_64" -> return ("docker", ["run", "-i", "--volume", filepath ++ ":/aurora", "btqag/aurora-prove"])
              -- insert "--platform=linux/amd64" when we are not on a x86 machine
              _ -> return ("docker", ["run", "-i", "--platform=linux/amd64", "--volume", filepath ++ ":/aurora", "btqag/aurora-prove"])
          else throwError CannotLocateProver
  
  findAuroraVerifier :: M (String, [String])
  findAuroraVerifier = do
    auroraExists <- checkCmd "aurora_verify"
    if auroraExists
      then return ("aurora_verify", [])
      else do
        dockerExists <- checkCmd "docker"
        if dockerExists
          then do
            lift $ Path.createDirectoryIfMissing True "aurora"
            filepath <- lift $ Path.makeAbsolute "aurora"
            case System.Info.arch of
              "x86_64" -> return ("docker", ["run", "-i", "--volume", filepath ++ ":/aurora", "btqag/aurora-verify"])
              -- insert "--platform=linux/amd64" when we are not on a x86 machine
              _ -> return ("docker", ["run", "-i", "--platform=linux/amd64", "--volume", filepath ++ ":/aurora", "btqag/aurora-verify"])
          else throwError CannotLocateVerifier
    
  
  findSnarkjs :: M (String, [String])
  findSnarkjs = do
    snarkJsExists <- checkCmd "snarkjs"
    if snarkJsExists then
      return ("snarkjs", [])
    else
      throwError CannotLocateSnarkjs
  
  -- | Check the version of the Keelung compiler
  readKeelungVersion :: FilePath -> [String] -> M (Int, Int, Int)
  readKeelungVersion cmd args = do
    -- trying to read version with `keelungc --version`
    rawString <-
      catchIOError
        CannotReadVersionError
        (Process.readProcess cmd (args ++ ["--version"]) mempty)
    -- parse see if the version number is well-formed
    let parseResult = case splitAt 9 rawString of
          ("Keelung v", versionString) -> parseVersion versionString
          _ -> Nothing
    -- throws CannotReadVersionError if it's not well-formed
    case parseResult of
      Nothing -> throwError CannotReadVersionError
      Just x -> return x
    where
      parseVersion :: String -> Maybe (Int, Int, Int)
      parseVersion versionString = do
        (major, minor, patch) <- case span (/= '.') versionString of
          (m, '.' : rest) -> case span (/= '.') rest of
            (n, '.' : p) -> Just (m, n, p)
            _ -> Nothing
          _ -> Nothing
        (,,) <$> readMaybe major <*> readMaybe minor <*> readMaybe patch
  
  -- | Check if the compiler version matches the version of this library
  --   patch number can be different
  checkCompilerVersion :: (Int, Int, Int) -> M ()
  checkCompilerVersion (major, minor, patch) = do
    let (compilerMajor, compilerMinor, compilerPatch) = keelungCompilerVersion
    if (major, minor) == (compilerMajor, compilerMinor) && patch >= compilerPatch
      then return ()
      else throwError (VersionMismatchError major minor patch)
  
  --------------------------------------------------------------------------------
  
  -- | Check if a command exists
  checkCmd :: String -> M Bool
  checkCmd cmd =
    lift $
      IO.catchIOError
        (Process.readProcess whichCmd [cmd] mempty >> return True)
        (\_ -> return False)
    where
      -- decide the command for locating executables
      whichCmd :: String
      whichCmd = case System.Info.os of
        "mingw32" -> "where" -- Windows uses "where"
        _ -> "which" -- Unix uses "which"
  
  --------------------------------------------------------------------------------
  
  type M = ExceptT Error IO
  
  runM :: M a -> IO (Either Error a)
  runM = runExceptT
  
  -- | Handle 'IO' Exceptions in the 'M' Monad
  catchIOError :: Error -> IO a -> M a
  catchIOError err f = lift (IO.catchIOError (Right <$> f) (const (return (Left err)))) >>= liftEither
  
  --------------------------------------------------------------------------------
  
  instance (Encode a) => Show (Comp a) where
    show prog = case elaborateAndEncode prog of
      Left err -> show err
      Right elaborated -> show elaborated
  