-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | This module contains everything related to compilation from Indigo to Lorentz,
-- including plain Indigo code, as well as Indigo contracts.

module Indigo.Compilation
  ( compileIndigo
  , IndigoWithParams
  , IndigoContract
  , compileIndigoContract

  , Ops
  , HasSideEffects
  , operationsVar
  , HasStorage
  , storageVar
  ) where

import qualified Data.Map as M
import Data.Reflection (give)
import qualified Data.Set as S
import Data.Singletons (SingI(..))
import Data.Typeable ((:~:)(..), eqT)
import Data.Vinyl.Core (RMap(..))

import qualified Indigo.Backend as B
import Indigo.Compilation.Lambda
import Indigo.Compilation.Params
import Indigo.Frontend.Program (IndigoM(..), Program(..))
import Indigo.Frontend.Statement
import Indigo.Internal hiding (SetField, return, (>>), (>>=))
import qualified Indigo.Internal as I
import Indigo.Lorentz
import Indigo.Prelude
import qualified Lorentz.Instr as L
import qualified Lorentz.Macro as L
import Util.Peano

-- | Iteration over Indigo freer monad
compileIndigoM
  :: forall inp a .
    (forall x anyInp . StatementF IndigoM x -> SomeIndigoState anyInp x)
  -> IndigoM a
  -> SomeIndigoState inp a
compileIndigoM _ (IndigoM (Done a)) = returnSIS a
compileIndigoM interp (IndigoM (Instr i)) = interp i
compileIndigoM interp (IndigoM (Bind instr cont)) =
  compileIndigoM interp (IndigoM instr) `bindSIS` (compileIndigoM interp . IndigoM . cont)

-- | Convert frontend Freer to 'IndigoState'.
--
-- First of all, this function generates the definitions of
-- lambdas, creates the variables that refer to them
-- and calls them in the places where they are used.
-- This happens only for those lambdas that are called
-- at least twice, those that are used only once will be
-- inlined instead.
--
-- After that the generation of the body code starts.
simpleCompileIndigoM :: forall inp a . IndigoM a -> SomeIndigoState inp a
simpleCompileIndigoM indigoM =
  let lambdas = S.toList (collectLambdas indigoM) in
  forMSIS lambdas defineLambda
  `bindSIS`
    (\defined ->
      let definedLambdas = M.fromList $ map (\l -> (_clName l, l)) defined
      in compileBody definedLambdas indigoM
    )
  where
    compileBody definedLambdas = compileIndigoM (usingReader definedLambdas . compileSt)

    compileSt :: StatementF IndigoM x -> Reader (Map String CompiledLambda) (SomeIndigoState anyInp x)
    compileSt (LiftIndigoState cd) = pure cd
    compileSt (NewVar ex) = pure $ toSIS (B.newVar ex)
    compileSt (SetVar v ex) = pure $ toSIS (B.setVar v ex)
    compileSt (SetField v fName ex) = pure $ toSIS (B.setField v fName ex)
    compileSt (VarModification act var ex) = pure $ toSIS (B.updateVar act var ex)

    compileSt (LambdaPure1Call lName (body :: (Var arg -> IndigoM res)) argm) =
      execGenericLambda @'[] @res (B.executeLambdaPure1 @res) lName body argm

    compileSt (Lambda1Call (_ :: Proxy st) lName (body :: (Var arg -> IndigoM res)) argm) =
      execGenericLambda @'[st] @res (B.executeLambda1 @st @res) lName body argm

    compileSt (LambdaEff1Call (_ :: Proxy st) lName (body :: (Var arg -> IndigoM res)) argm) =
      execGenericLambda @'[st, Ops] @res (B.executeLambdaEff1 @st @res) lName body argm

    compileSt (Scope cd) = do
      definedLambdas <- ask
      pure $ withSIS (compileBody definedLambdas cd) (toSIS . B.scope)
    compileSt (If ex tb fb) = do
      definedLambdas <- ask
      pure $ withSIS (compileBody definedLambdas tb) $ \tb' ->
        withSIS (compileBody definedLambdas fb) $ \fb' ->
         toSIS (B.if_ ex tb' fb')
    compileSt (IfSome ex tb fb) = do
      definedLambdas <- ask
      pure $ withSIS1 (compileBody definedLambdas . tb) $ \tb' ->
        withSIS (compileBody definedLambdas fb) $ \fb' ->
          toSIS (B.ifSome ex tb' fb')
    compileSt (IfRight ex rb lb) = do
      definedLambdas <- ask
      pure $ withSIS1 (compileBody definedLambdas . rb) $ \rb' ->
        withSIS1 (compileBody definedLambdas . lb) $ \lb' ->
          toSIS (B.ifRight ex rb' lb')
    compileSt (IfCons ex tb fb) = do
      definedLambdas <- ask
      pure $ withSIS2 (\x y -> compileBody definedLambdas $ tb x y) $ \tb' ->
        withSIS (compileBody definedLambdas fb) $ \fb' ->
          toSIS (B.ifCons ex tb' fb')
    compileSt (Case grd clauses) = do
      definedLambdas <- ask
      pure $ toSIS $ B.caseRec grd (rmapClauses definedLambdas clauses)
    compileSt (EntryCase proxy grd clauses) = do
      definedLambdas <- ask
      pure $ toSIS $ B.entryCaseRec proxy grd (rmapClauses definedLambdas clauses)
    compileSt (EntryCaseSimple grd clauses) = do
      definedLambdas <- ask
      pure $ toSIS $ B.entryCaseSimpleRec grd (rmapClauses definedLambdas clauses)

    compileSt (While ex body) = do
      definedLambdas <- ask
      pure $ withSIS (compileBody definedLambdas body) $ \bd -> toSIS (B.while ex bd)
    compileSt (WhileLeft ex lb) = do
      definedLambdas <- ask
      pure $
        withSIS1 (compileBody definedLambdas . lb) $ \lb' -> do
          toSIS (B.whileLeft ex lb')
    compileSt (ForEach e body) = do
      definedLambdas <- ask
      pure $ withSIS1 (compileBody definedLambdas . body) $ \bd -> toSIS (B.forEach e bd)

    compileSt (ContractName cName contr) = do
      definedLambdas <- ask
      pure $ withSIS (compileBody definedLambdas contr) $ toSIS . B.contractName cName
    compileSt (DocGroup gr ii) = do
      definedLambdas <- ask
      pure $ withSIS (compileBody definedLambdas ii) $ toSIS . B.docGroup gr
    compileSt (ContractGeneral contr) = do
      definedLambdas <- ask
      pure $ withSIS (compileBody definedLambdas contr) (toSIS . B.contractGeneral)
    compileSt (FinalizeParamCallingDoc entrypoint param) = do
      definedLambdas <- ask
      pure $ withSIS1 (compileBody definedLambdas . entrypoint)
        (\bd -> toSIS $ B.finalizeParamCallingDoc bd param)

    compileSt (TransferTokens expar exm exc) = pure $ toSIS (B.transferTokens expar exm exc)
    compileSt (SetDelegate kh) = pure $ toSIS (B.setDelegate kh)
    compileSt (CreateContract lCtr ek em es) = pure $ toSIS $
      I.iget I.>>= \(MetaData s _) ->
        ternaryOp ek em es (L.createContract lCtr
                                              # varActionOperation (NoRef :& s))
        I.>> makeTopVar
    compileSt (ContractCalling (_ :: Proxy cp) ref addr) = pure $ toSIS $ B.contractCalling @cp ref addr

    compileSt (FailWith ex) = pure $ toSIS $ B.failWith ex
    compileSt (Assert err expr) = pure $ toSIS $ B.assert err expr
    compileSt (FailCustom l expr) = pure $ toSIS $ B.failCustom l expr

    rmapClauses:: forall ret cs . RMap cs
       => Map String CompiledLambda
       -> Rec (IndigoMCaseClauseL IndigoM ret) cs
       -> Rec (B.IndigoCaseClauseL ret) cs
    rmapClauses definedLambdas = rmap (\(OneFieldIndigoMCaseClauseL cName clause) ->
      cName /-> (\v -> B.IndigoAnyOut $ compileBody definedLambdas $ clause v))

    forMSIS :: [r] -> (forall someInp . r -> SomeIndigoState someInp v) -> SomeIndigoState someInp1 [v]
    forMSIS [] _ = returnSIS []
    forMSIS (x : xs) f = f x `bindSIS` (\what -> (what :) <$> forMSIS xs f)

    defineLambda :: Lambda1Def -> SomeIndigoState someOut CompiledLambda
    defineLambda (LambdaPure1Def (_ :: Proxy (_s, arg, res)) lName fun) =
      defineGenericLambda @'[] B.initMetaDataPure B.createLambdaPure1 lName fun
    defineLambda (Lambda1Def (_ :: Proxy (st, arg, res)) lName fun) =
      defineGenericLambda @'[st] B.initMetaData B.createLambda1 lName fun
    defineLambda (LambdaEff1Def (_ :: Proxy (st, arg, res)) lName fun) =
      defineGenericLambda @'[st, Ops] B.initMetaDataEff B.createLambdaEff1 lName fun

    defineGenericLambda
      :: forall extra res arg someOut .
      (Typeable arg, Typeable res, Typeable extra)
      => (Var arg, MetaData (arg & extra))
      -> (forall inpt out . B.LambdaCreator extra arg res inpt out)
      -> String
      -> (Var arg -> IndigoM res)
      -> SomeIndigoState someOut CompiledLambda
    defineGenericLambda (varArg, initMd) lambdaCreator lName fun = do
      runSIS
        (simpleCompileIndigoM $ fun varArg) initMd
        (\gc -> toSIS $ lambdaCreator (\_v -> IndigoState $ \_md -> gc))
      `bindSIS`
      (returnSIS . CompiledLambda (Proxy @res) lName)

    execGenericLambda
      :: forall extra res arg someOut .
         (Typeable extra, KnownValue arg, Typeable res, B.ScopeCodeGen res)
      => (forall inpt . B.LambdaExecutor extra arg res inpt)
      -> String
      -> (Var arg -> IndigoM res)
      -> Expr arg
      -> Reader (Map String CompiledLambda) (SomeIndigoState someOut (B.RetVars res))
    execGenericLambda executor lName (body :: (Var arg -> IndigoM res)) (argm :: Expr arg) = do
      compiled <- ask
      let maybeToRight' = flip maybeToRight
      -- This code seems to be pretty unsafe, but it works almost inevitably
      pure $ either (error . fromString) id $ do
        case M.lookup lName compiled of
          Nothing -> Right $
            -- Just inline lambda without calling Lorentz lambda
            withSIS1 (compileBody compiled . body)
              (\bd -> toSIS $ B.newVar argm I.>>= (B.scope @res . bd))
          Just compLam -> case compLam of
            CompiledLambda (_ :: Proxy res1) _ (varF :: Var (B.Lambda1Generic extra1 arg1 res1)) -> do
              Refl <- maybeToRight' (eqT @res @res1) ("unexpected result type of " ++ lName ++ " lambda didn't match")
              Refl <- maybeToRight' (eqT @arg @arg1) ("unexpected argument type of " ++ lName ++ " lambda didn't match")
              Refl <- maybeToRight' (eqT @extra @extra1) ("unexpected storage type of " ++ lName ++ " lambda didn't match")
              pure $ toSIS (executor varF argm)

-- | Compile Indigo code to Lorentz.
--
-- Note: it is necessary to specify the number of parameters (using the first
-- type variable) of the Indigo function. Also, these should be on the top of
-- the input stack in inverse order (see 'IndigoWithParams').
compileIndigo
  :: forall n inp a.
     ( SingI (ToPeano n), Default (MetaData inp)
     , AreIndigoParams (ToPeano n) inp, KnownValue a
     )
  => IndigoWithParams (ToPeano n) inp a
  -> inp :-> inp
compileIndigo paramCode =
  runSIS (simpleCompileIndigoM code) md cleanGenCode
  where
    (code, md) = fromIndigoWithParams @inp @_ @a paramCode def (sing @(ToPeano n))

-- | Type of a contract that can be compiled to Lorentz with 'compileIndigoContract'.
type IndigoContract param st =
  (HasStorage st, HasSideEffects) => Var param -> IndigoM ()

-- | Compile Indigo code to Lorentz contract.
-- Drop elements from the stack to return only @[Operation]@ and @storage@.
compileIndigoContract
  :: forall param st .
  ( KnownValue param
  , IsObject st
  )
  => IndigoContract param st
  -> ContractCode param st
compileIndigoContract code =
  let  (varOps, opsMd) = pushRefMd emptyMetadata
       mdSt = pushNoRefMd opsMd in
  -- Decompose storage value first, run contract and then compose it back.
  runSIS (deepDecomposeCompose @st) mdSt $ \(GenCode varSt decomposedMd decomposeSt composeSt) ->
    let (varParam, initMd) = pushRefMd decomposedMd
        everythingGiven = (give @(Var Ops) varOps $ give @(Var st) varSt code) varParam
        indigoCode = runSIS (simpleCompileIndigoM everythingGiven) initMd cleanGenCode in
    L.nil # L.swap # L.unpair #
    L.dip decomposeSt # -- decompose storage
    indigoCode # -- run indigo code
    L.drop # -- drop param
    composeSt # -- compose storage back
    L.swap # L.pair