module Language.Elna.RiscV.CodeGenerator
    ( 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 (Operand(..), 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
        ]
    }

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 (Vector RiscVQuadruple) -> Vector Statement
generateRiscV = flip evalState 0
    . runRiscVGenerator
    . foldlM go Vector.empty
    . HashMap.toList
  where
    go accumulator (Identifier key, 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 register))
    | IntOperand immediateOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        pure $ lui (immediateOperand1 + immediateOperand2) register
    | VariableOperand variableOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        let Store operandRegister1 = variableOperand1
            Store operandRegister2 = variableOperand2
        in pure $ Vector.singleton $ Instruction
            $ RiscV.BaseInstruction RiscV.Op
            $ RiscV.R register RiscV.ADD operandRegister1 operandRegister2 (RiscV.Funct7 0b0000000)
    | VariableOperand variableOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        addImmediateRegister variableOperand1 immediateOperand2
    | IntOperand immediateOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        addImmediateRegister variableOperand2 immediateOperand1
  where
    addImmediateRegister variableOperand immediateOperand = 
        let statements = lui immediateOperand register
            Store operandRegister = variableOperand
        in pure $ Vector.snoc statements
            $ Instruction
            $ RiscV.BaseInstruction RiscV.Op
            $ RiscV.R register RiscV.ADD register operandRegister
            $ RiscV.Funct7 0b0000000
quadruple (SubtractionQuadruple operand1 operand2 (Store register))
    | IntOperand immediateOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        pure $ lui (immediateOperand1 - immediateOperand2) register
    | VariableOperand variableOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        let Store operandRegister1 = variableOperand1
            Store operandRegister2 = variableOperand2
        in pure $ Vector.singleton $ Instruction
            $ RiscV.BaseInstruction RiscV.Op
            $ RiscV.R register RiscV.SUB operandRegister1 operandRegister2
            $ RiscV.Funct7 0b0100000
    | IntOperand immediateOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        let statements1 = lui immediateOperand1 register
            Store operandRegister2 = variableOperand2
         in pure $ Vector.snoc statements1
            $ Instruction
            $ RiscV.BaseInstruction RiscV.Op
            $ RiscV.R register RiscV.SUB register operandRegister2
            $ RiscV.Funct7 0b0100000
    | VariableOperand variableOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        let statements2 = lui (negate immediateOperand2) register
            Store operandRegister1 = variableOperand1
         in pure $ Vector.snoc statements2
            $ Instruction
            $ RiscV.BaseInstruction RiscV.Op
            $ RiscV.R register RiscV.ADD register operandRegister1
            $ RiscV.Funct7 0b0000000
quadruple (NegationQuadruple operand1 (Store register))
    | IntOperand immediateOperand1 <- operand1 =
        pure $ lui (negate immediateOperand1) register
    | VariableOperand variableOperand1 <- operand1 =
        let Store operandRegister1 = variableOperand1
         in pure $ Vector.singleton
            $ Instruction
            $ RiscV.BaseInstruction RiscV.Op
            $ RiscV.R register RiscV.SUB RiscV.Zero operandRegister1
            $ RiscV.Funct7 0b0100000
quadruple (ProductQuadruple operand1 operand2 (Store register))
    | IntOperand immediateOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        pure $ lui (immediateOperand1 * immediateOperand2) register
    | VariableOperand variableOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        let Store operandRegister1 = variableOperand1
            Store operandRegister2 = variableOperand2
        in pure $ Vector.singleton $ Instruction
            $ RiscV.BaseInstruction RiscV.Op
            $ RiscV.R register RiscV.MUL operandRegister1 operandRegister2 (RiscV.Funct7 0b0000001)
    | VariableOperand variableOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        multiplyImmediateRegister variableOperand1 immediateOperand2
    | IntOperand immediateOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        multiplyImmediateRegister variableOperand2 immediateOperand1
  where
    multiplyImmediateRegister variableOperand immediateOperand = 
        let statements = lui immediateOperand register
            Store operandRegister = variableOperand
        in pure $ Vector.snoc statements
            $ Instruction
            $ RiscV.BaseInstruction RiscV.Op
            $ RiscV.R register RiscV.MUL register operandRegister
            $ RiscV.Funct7 0b0000001
quadruple (DivisionQuadruple operand1 operand2 (Store register))
    | IntOperand immediateOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        if immediateOperand2 == 0
        then pure $ Vector.singleton
            $ Instruction (RiscV.CallInstruction "_divide_by_zero_error")
        else pure $ lui (quot immediateOperand1 immediateOperand2) register
    | VariableOperand variableOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 = do
        let Store operandRegister1 = variableOperand1
            Store operandRegister2 = variableOperand2
            divisionInstruction = Instruction
                $ RiscV.BaseInstruction RiscV.Op
                $ RiscV.R register 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 $ Vector.fromList
            [ branchInstruction
            , Instruction (RiscV.CallInstruction "_divide_by_zero_error")
            , JumpLabel branchLabel []
            , divisionInstruction
            ]
    | VariableOperand variableOperand1 <- operand1
    , IntOperand immediateOperand2 <- operand2 =
        let statements2 = lui immediateOperand2 register
            Store operandRegister1 = variableOperand1
        in pure $ Vector.snoc statements2
            $ Instruction
            $ RiscV.BaseInstruction RiscV.Op
            $ RiscV.R register RiscV.DIV operandRegister1 register
            $ RiscV.Funct7 0b0000001
    | IntOperand immediateOperand1 <- operand1
    , VariableOperand variableOperand2 <- operand2 =
        let statements1 = lui immediateOperand1 register
            Store operandRegister2 = variableOperand2
        in pure $ Vector.snoc statements1
            $ Instruction
            $ RiscV.BaseInstruction RiscV.Op
            $ RiscV.R register RiscV.DIV register operandRegister2
            $ RiscV.Funct7 0b0000001

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

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