module Language.Elna.RiscV.CodeGenerator
    ( Directive(..)
    , Statement(..)
    , generateRiscV
    , riscVConfiguration
    ) where

import Control.Monad.Trans.State (State, get, evalState, modify')
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.Int (Int32)
import Data.Word (Word32)
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import qualified Language.Elna.Architecture.RiscV as RiscV
import Language.Elna.Backend.Allocator (MachineConfiguration(..), Store(..))
import Language.Elna.Backend.Intermediate
    ( Label(..)
    , Operand(..)
    , ProcedureQuadruples(..)
    , Quadruple(..)
    )
import Language.Elna.Location (Identifier(..))
import Data.Bits (Bits(..))
import Data.Foldable (Foldable(..), foldlM)
import Data.Text (Text)
import qualified Data.Text.Lazy.Builder.Int as Text.Builder
import qualified Data.Text.Lazy.Builder as Text.Builder
import qualified Data.Text.Lazy as Text.Lazy

data Directive
    = GlobalDirective
    | FunctionDirective
    deriving (Eq, Show)

data Statement
    = Instruction RiscV.Instruction
    | JumpLabel Text [Directive]
    deriving Eq

riscVConfiguration :: MachineConfiguration RiscV.XRegister
riscVConfiguration = MachineConfiguration
    { temporaryRegisters =
        [ RiscV.T0
        , RiscV.T1
        , RiscV.T2
        , RiscV.T3
        , RiscV.T4
        , RiscV.T5
        , RiscV.T6
        ]
    }

-- | Reserved register used for calculations to save an immediate temporary.
immediateRegister :: RiscV.XRegister
immediateRegister = RiscV.A7

type RiscVStore = Store RiscV.XRegister
type RiscVQuadruple = Quadruple RiscVStore
type RiscVOperand = Operand RiscVStore

newtype RiscVGenerator a = RiscVGenerator
    { runRiscVGenerator :: State Word32 a }

instance Functor RiscVGenerator
  where
    fmap f (RiscVGenerator x) = RiscVGenerator $ f <$> x

instance Applicative RiscVGenerator
  where
    pure = RiscVGenerator . pure
    (RiscVGenerator f) <*> (RiscVGenerator x) = RiscVGenerator $ f <*> x

instance Monad RiscVGenerator
  where
    (RiscVGenerator x) >>= f = RiscVGenerator $ x >>= (runRiscVGenerator . f)

createLabel :: RiscVGenerator Text
createLabel = do
    currentCounter <- RiscVGenerator get
    RiscVGenerator $ modify' (+ 1)
    pure
        $ mappend ".A"
        $ Text.Lazy.toStrict
        $ Text.Builder.toLazyText
        $ Text.Builder.decimal currentCounter

generateRiscV :: HashMap Identifier (ProcedureQuadruples RiscVStore) -> Vector Statement
generateRiscV = flip evalState 0
    . runRiscVGenerator
    . foldlM go Vector.empty
    . HashMap.toList
  where
    go accumulator (Identifier key, ProcedureQuadruples{ quadruples = value }) =
        let code = Vector.cons (JumpLabel key [GlobalDirective, FunctionDirective])
                . fold <$> mapM quadruple value
         in (accumulator <>) <$> code

quadruple :: RiscVQuadruple -> RiscVGenerator (Vector Statement)
quadruple StartQuadruple = pure $ Vector.fromList
    [ Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.SP RiscV.ADDI RiscV.SP (negate 4))
    , Instruction (RiscV.BaseInstruction RiscV.Store $ RiscV.S 0 RiscV.SW RiscV.SP RiscV.S0)
    , Instruction (RiscV.BaseInstruction RiscV.Store $ RiscV.S 4 RiscV.SW RiscV.SP RiscV.RA)
    , Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.S0 RiscV.ADDI RiscV.SP 4)
    ]
quadruple (ParameterQuadruple operand1) =
    let (operandRegister, statements) = loadImmediateOrRegister operand1 RiscV.A0
     in pure $ mappend statements $ Vector.fromList
        [ Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.SP RiscV.ADDI RiscV.SP (negate 4))
        , Instruction (RiscV.BaseInstruction RiscV.Store $ RiscV.S 0 RiscV.SW RiscV.SP operandRegister)
        ]
quadruple (CallQuadruple callName numberOfArguments) = pure $ Vector.fromList
    [ Instruction (RiscV.CallInstruction callName)
    , Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.SP RiscV.ADDI RiscV.SP (numberOfArguments * 4))
    ]
quadruple StopQuadruple = pure $ Vector.fromList
    [ Instruction (RiscV.BaseInstruction RiscV.Load $ RiscV.I RiscV.S0 RiscV.LW RiscV.SP 0)
    , Instruction (RiscV.BaseInstruction RiscV.Load $ RiscV.I RiscV.RA RiscV.LW RiscV.SP 4)
    , Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.SP RiscV.ADDI RiscV.SP 4)
    , Instruction (RiscV.BaseInstruction RiscV.Jalr $ RiscV.I RiscV.RA RiscV.JALR RiscV.Zero 0)
    ]
quadruple (AddQuadruple operand1 operand2 store) =
    commutativeBinary (+) RiscV.ADD (RiscV.Funct7 0b0000000) (operand1, operand2) store
quadruple (ProductQuadruple operand1 operand2 store) =
    commutativeBinary (*) RiscV.MUL (RiscV.Funct7 0b0000001) (operand1, operand2) store
quadruple (SubtractionQuadruple operand1 operand2 store)
    | IntOperand immediateOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        let (storeRegister, storeStatements) = storeToStore store
         in pure $ lui (immediateOperand1 - immediateOperand2) storeRegister <> storeStatements
    | VariableOperand variableOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        let (storeRegister, storeStatements) = storeToStore store
            (operandRegister1, statements1) = loadFromStore variableOperand1
            (operandRegister2, statements2) = loadFromStore variableOperand2
            instruction = Instruction
                $ RiscV.BaseInstruction RiscV.Op
                $ RiscV.R storeRegister RiscV.SUB operandRegister1 operandRegister2
                $ RiscV.Funct7 0b0100000
        in pure $ statements1 <> statements2 <> Vector.cons instruction storeStatements
    | IntOperand immediateOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        let (storeRegister, storeStatements) = storeToStore store
            statements1 = lui immediateOperand1 storeRegister
            (operandRegister2, statements2) = loadFromStore variableOperand2
            instruction = Instruction
                $ RiscV.BaseInstruction RiscV.Op
                $ RiscV.R storeRegister RiscV.SUB storeRegister operandRegister2
                $ RiscV.Funct7 0b0100000
         in pure $ statements1 <> statements2 <> Vector.cons instruction storeStatements
    | VariableOperand variableOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        let (storeRegister, storeStatements) = storeToStore store
            statements2 = lui (negate immediateOperand2) storeRegister
            (operandRegister1, statements1) = loadFromStore variableOperand1
            instruction = Instruction
                $ RiscV.BaseInstruction RiscV.Op
                $ RiscV.R storeRegister RiscV.ADD storeRegister operandRegister1
                $ RiscV.Funct7 0b0000000
         in pure $ statements1 <> statements2 <> Vector.cons instruction storeStatements
quadruple (NegationQuadruple operand1 store)
    | IntOperand immediateOperand1 <- operand1 =
        let (storeRegister, storeStatements) = storeToStore store
         in pure $ lui (negate immediateOperand1) storeRegister <> storeStatements
    | VariableOperand variableOperand1 <- operand1 =
        let (storeRegister, storeStatements) = storeToStore store
            (operandRegister1, statements1) = loadFromStore variableOperand1
            instruction = Instruction
                $ RiscV.BaseInstruction RiscV.Op
                $ RiscV.R storeRegister RiscV.SUB RiscV.Zero operandRegister1
                $ RiscV.Funct7 0b0100000
         in pure $ statements1 <> Vector.cons instruction storeStatements
quadruple (DivisionQuadruple operand1 operand2 store)
    | IntOperand immediateOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        if immediateOperand2 == 0
        then pure $ Vector.singleton
            $ Instruction (RiscV.CallInstruction "_divide_by_zero_error")
        else
            let (storeRegister, storeStatements) = storeToStore store
            in pure $ lui (quot immediateOperand1 immediateOperand2) storeRegister <> storeStatements
    | VariableOperand variableOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 = do
        let (storeRegister, storeStatements) = storeToStore store
            (operandRegister1, statements1) = loadFromStore variableOperand1
            (operandRegister2, statements2) = loadFromStore variableOperand2
            divisionInstruction = Instruction
                $ RiscV.BaseInstruction RiscV.Op
                $ RiscV.R storeRegister RiscV.DIV operandRegister1 operandRegister2 (RiscV.Funct7 0b0000001)
        branchLabel <- createLabel
        let branchInstruction = Instruction
                $ RiscV.RelocatableInstruction RiscV.Branch
                $ RiscV.RBranch branchLabel RiscV.BNE RiscV.Zero operandRegister2
        pure $ statements1 <> statements2 <> Vector.fromList
            [ branchInstruction
            , Instruction (RiscV.CallInstruction "_divide_by_zero_error")
            , JumpLabel branchLabel []
            , divisionInstruction
            ] <> storeStatements
    | VariableOperand variableOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        let (storeRegister, storeStatements) = storeToStore store
            statements2 = lui immediateOperand2 storeRegister
            (operandRegister1, statements1) = loadFromStore variableOperand1
            operationInstruction
                | immediateOperand2 == 0 =
                    RiscV.CallInstruction "_divide_by_zero_error"
                | otherwise = RiscV.BaseInstruction RiscV.Op
                    $ RiscV.R storeRegister RiscV.DIV operandRegister1 storeRegister
                    $ RiscV.Funct7 0b0000001
         in pure $ statements1 <> statements2
            <> Vector.cons (Instruction operationInstruction) storeStatements
    | IntOperand immediateOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 = do
        let (storeRegister, storeStatements) = storeToStore store
            statements1 = lui immediateOperand1 storeRegister
            (operandRegister2, statements2) = loadFromStore variableOperand2
            divisionInstruction = Instruction
                $ RiscV.BaseInstruction RiscV.Op
                $ RiscV.R storeRegister RiscV.DIV storeRegister operandRegister2 (RiscV.Funct7 0b0000001)
        branchLabel <- createLabel
        let branchInstruction = Instruction
                $ RiscV.RelocatableInstruction RiscV.Branch
                $ RiscV.RBranch branchLabel RiscV.BNE RiscV.Zero operandRegister2
        pure $ statements1 <> statements2 <> Vector.fromList 
            [ branchInstruction
            , Instruction (RiscV.CallInstruction "_divide_by_zero_error")
            , JumpLabel branchLabel []
            , divisionInstruction
            ] <> storeStatements
quadruple (LabelQuadruple (Label label)) = pure $ Vector.singleton $ JumpLabel label mempty
quadruple (GoToQuadruple label) = pure $ Vector.singleton $ unconditionalJal label
quadruple (EqualQuadruple operand1 operand2 goToLabel) =
    commutativeComparison (==) RiscV.BEQ (operand1, operand2) goToLabel
quadruple (NonEqualQuadruple operand1 operand2 goToLabel) =
    commutativeComparison (/=) RiscV.BNE (operand1, operand2) goToLabel
quadruple (LessQuadruple operand1 operand2 goToLabel) =
    lessThan (operand1, operand2) goToLabel
quadruple (GreaterQuadruple operand1 operand2 goToLabel) =
    lessThan (operand2, operand1) goToLabel
quadruple (LessOrEqualQuadruple operand1 operand2 goToLabel) =
    lessOrEqualThan (operand1, operand2) goToLabel
quadruple (GreaterOrEqualQuadruple operand1 operand2 goToLabel) =
    lessOrEqualThan (operand2, operand1) goToLabel
quadruple (AssignQuadruple operand1 store)
    | IntOperand immediateOperand1 <- operand1 =
        let (storeRegister, storeStatements) = storeToStore store
         in pure $ lui immediateOperand1 storeRegister <> storeStatements
    | VariableOperand variableOperand1 <- operand1 =
        let (operandRegister1, statements1) = loadFromStore variableOperand1
            (storeRegister, storeStatements) = storeToStore store
            instruction = Instruction
                $ RiscV.BaseInstruction RiscV.OpImm
                $ RiscV.I storeRegister RiscV.ADDI operandRegister1 0
         in pure $ statements1 <> Vector.cons instruction storeStatements 

unconditionalJal :: Label -> Statement
unconditionalJal (Label goToLabel) = Instruction
    $ RiscV.RelocatableInstruction RiscV.Jal
    $ RiscV.RJal RiscV.Zero goToLabel

loadImmediateOrRegister :: RiscVOperand -> RiscV.XRegister -> (RiscV.XRegister, Vector Statement)
loadImmediateOrRegister (IntOperand intValue) targetRegister =
    (targetRegister, lui intValue targetRegister)
loadImmediateOrRegister (VariableOperand store) _ = loadFromStore store

lui :: Int32 -> RiscV.XRegister -> Vector Statement
lui intValue targetRegister
    | intValue >= -2048
    , intValue <= 2047 = Vector.singleton
        $ Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I targetRegister RiscV.ADDI RiscV.Zero lo)
    | intValue .&. 0x800 /= 0 = Vector.fromList
        [ Instruction (RiscV.BaseInstruction RiscV.Lui $ RiscV.U targetRegister $ fromIntegral $ succ hi)
        , Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I targetRegister RiscV.ADDI targetRegister lo)
        ]
    | otherwise = Vector.fromList
        [ Instruction (RiscV.BaseInstruction RiscV.Lui $ RiscV.U targetRegister $ fromIntegral hi)
        , Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I targetRegister RiscV.ADDI targetRegister lo)
        ]
  where
    hi = intValue `shiftR` 12
    lo = fromIntegral intValue

commutativeBinary
    :: (Int32 -> Int32 -> Int32)
    -> RiscV.Funct3
    -> RiscV.Funct7
    -> (Operand RiscVStore, Operand RiscVStore)
    -> Store RiscV.XRegister
    -> RiscVGenerator (Vector Statement)
commutativeBinary immediateOperation funct3 funct7 (operand1, operand2) store
    | IntOperand immediateOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        let (storeRegister, storeStatements) = storeToStore store
            immediateOperation' = immediateOperation immediateOperand1 immediateOperand2
         in pure $ lui immediateOperation' storeRegister <> storeStatements
    | VariableOperand variableOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        let (operandRegister1, statements1) = loadFromStore variableOperand1
            (operandRegister2, statements2) = loadFromStore variableOperand2
            (storeRegister, storeStatements) = storeToStore store
            instruction = Instruction $ RiscV.BaseInstruction RiscV.Op
                $ RiscV.R storeRegister funct3 operandRegister1 operandRegister2 funct7
         in pure $ statements1 <> statements2
            <> Vector.cons instruction storeStatements
    | VariableOperand variableOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        commutativeImmediateRegister variableOperand1 immediateOperand2
    | IntOperand immediateOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        commutativeImmediateRegister variableOperand2 immediateOperand1
  where
    commutativeImmediateRegister variableOperand immediateOperand = 
        let (storeRegister, storeStatements) = storeToStore store
            immediateStatements = lui immediateOperand storeRegister
            (operandRegister, registerStatements) = loadFromStore variableOperand
            instruction = Instruction
                $ RiscV.BaseInstruction RiscV.Op
                $ RiscV.R storeRegister funct3 storeRegister operandRegister funct7
        in pure $ immediateStatements <> registerStatements
            <> Vector.cons instruction storeStatements

commutativeComparison
    :: (Int32 -> Int32 -> Bool)
    -> RiscV.Funct3
    -> (Operand RiscVStore, Operand RiscVStore)
    -> Label
    -> RiscVGenerator (Vector Statement)
commutativeComparison immediateOperation funct3 (operand1, operand2) goToLabel
    | IntOperand immediateOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        if immediateOperation immediateOperand1 immediateOperand2
        then pure $ Vector.singleton $ unconditionalJal goToLabel
        else pure Vector.empty
    | VariableOperand variableOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 = do
        let (operandRegister1, statements1) = loadFromStore variableOperand1
            (operandRegister2, statements2) = loadFromStore variableOperand2
            Label goToLabel' = goToLabel
        pure $ Vector.snoc (statements1 <> statements2)
            $ Instruction
            $ RiscV.RelocatableInstruction RiscV.Branch
            $ RiscV.RBranch goToLabel' funct3 operandRegister1 operandRegister2
    | VariableOperand variableOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        compareImmediateRegister variableOperand1 immediateOperand2
    | IntOperand immediateOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        compareImmediateRegister variableOperand2 immediateOperand1
  where
    compareImmediateRegister variableOperand immediateOperand = 
        let immediateStatements = lui immediateOperand immediateRegister
            (operandRegister, registerStatements) = loadFromStore variableOperand
            Label goToLabel' = goToLabel
        in pure $ Vector.snoc (immediateStatements <> registerStatements)
            $ Instruction
            $ RiscV.RelocatableInstruction RiscV.Branch
            $ RiscV.RBranch goToLabel' funct3 operandRegister immediateRegister

lessThan :: (Operand RiscVStore, Operand RiscVStore) -> Label -> RiscVGenerator (Vector Statement)
lessThan (operand1, operand2) goToLabel
    | IntOperand immediateOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        if immediateOperand1 < immediateOperand2
        then pure $ Vector.singleton $ unconditionalJal goToLabel
        else pure Vector.empty
    | VariableOperand variableOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 = do
        let (operandRegister1, statements1) = loadFromStore variableOperand1
            (operandRegister2, statements2) = loadFromStore variableOperand2
            Label goToLabel' = goToLabel
        pure $ Vector.snoc (statements1 <> statements2)
            $ Instruction
            $ RiscV.RelocatableInstruction RiscV.Branch
            $ RiscV.RBranch goToLabel' RiscV.BLT operandRegister1 operandRegister2
    | VariableOperand variableOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        let statements2 = lui immediateOperand2 immediateRegister
            (operandRegister1, statements1) = loadFromStore variableOperand1
            Label goToLabel' = goToLabel
        in pure $ Vector.snoc (statements1 <> statements2)
            $ Instruction
            $ RiscV.RelocatableInstruction RiscV.Branch
            $ RiscV.RBranch goToLabel' RiscV.BLT operandRegister1 immediateRegister
    | IntOperand immediateOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        let statements1 = lui immediateOperand1 immediateRegister
            (operandRegister2, statements2) = loadFromStore variableOperand2
            Label goToLabel' = goToLabel
        in pure $ Vector.snoc (statements1 <> statements2)
            $ Instruction
            $ RiscV.RelocatableInstruction RiscV.Branch
            $ RiscV.RBranch goToLabel' RiscV.BLT immediateRegister operandRegister2

lessOrEqualThan :: (Operand RiscVStore, Operand RiscVStore) -> Label -> RiscVGenerator (Vector Statement)
lessOrEqualThan (operand1, operand2) goToLabel
    | IntOperand immediateOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        if immediateOperand1 <= immediateOperand2
        then pure $ Vector.singleton $ unconditionalJal goToLabel
        else pure Vector.empty
    | VariableOperand variableOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 = do
        let (operandRegister1, statements1) = loadFromStore variableOperand1
            (operandRegister2, statements2) = loadFromStore variableOperand2
            Label goToLabel' = goToLabel
        pure $ Vector.snoc (statements1 <> statements2)
            $ Instruction
            $ RiscV.RelocatableInstruction RiscV.Branch
            $ RiscV.RBranch goToLabel' RiscV.BGE operandRegister2 operandRegister1
    | VariableOperand variableOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        let statements2 = lui immediateOperand2 immediateRegister
            (operandRegister1, statements1) = loadFromStore variableOperand1
            Label goToLabel' = goToLabel
        in pure $ Vector.snoc (statements1 <> statements2)
            $ Instruction
            $ RiscV.RelocatableInstruction RiscV.Branch
            $ RiscV.RBranch goToLabel' RiscV.BGE immediateRegister operandRegister1
    | IntOperand immediateOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        let statements1 = lui immediateOperand1 immediateRegister
            (operandRegister2, statements2) = loadFromStore variableOperand2
            Label goToLabel' = goToLabel
        in pure $ Vector.snoc (statements1 <> statements2)
            $ Instruction
            $ RiscV.RelocatableInstruction RiscV.Branch
            $ RiscV.RBranch goToLabel' RiscV.BGE operandRegister2 immediateRegister

loadFromStore :: RiscVStore -> (RiscV.XRegister, Vector Statement)
loadFromStore (RegisterStore register) = (register, mempty)
loadFromStore (StackStore offset register) =
    let loadInstruction = Instruction
            $ RiscV.BaseInstruction RiscV.Load
            $ RiscV.I register RiscV.LW RiscV.SP offset
     in (register, Vector.singleton loadInstruction)

storeToStore :: RiscVStore -> (RiscV.XRegister, Vector Statement)
storeToStore (RegisterStore register) = (register, mempty)
storeToStore (StackStore offset register) =
    let storeInstruction = Instruction
            $ RiscV.BaseInstruction RiscV.Store
            $ RiscV.S offset RiscV.SW RiscV.SP register
     in (register, Vector.singleton storeInstruction)