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

import Data.ByteString (ByteString)
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.Int (Int32)
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import qualified Data.Text.Encoding as Text.Encoding
import Language.Elna.Allocator (MachineConfiguration(..), Store(..))
import Language.Elna.Location (Identifier(..))
import Language.Elna.Intermediate (Operand(..), Quadruple(..))
import qualified Language.Elna.Architecture.RiscV as RiscV
import Data.Bits (Bits(..))

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

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

riscVConfiguration :: MachineConfiguration RiscV.XRegister
riscVConfiguration = MachineConfiguration
    { temporaryRegister = RiscV.T0
    }

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

generateRiscV :: HashMap Identifier (Vector RiscVQuadruple) -> Vector Statement
generateRiscV = HashMap.foldlWithKey' go Vector.empty
  where
    go accumulator (Identifier key) value =
        let code = Vector.cons (JumpLabel (Text.Encoding.encodeUtf8 key) [GlobalDirective, FunctionDirective])
                $ Vector.foldMap quadruple value
         in accumulator <> code

quadruple :: RiscVQuadruple -> Vector Statement
quadruple StartQuadruple = 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 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) = Vector.fromList
    [ Instruction (RiscV.CallInstruction callName)
    , Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.SP RiscV.ADDI RiscV.SP (numberOfArguments * 4))
    ]
quadruple StopQuadruple = 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)) =
    let (operandRegister1, statements1) = loadImmediateOrRegister operand1 RiscV.A0
        (operandRegister2, statements2) = loadImmediateOrRegister operand2 RiscV.A1
     in Vector.snoc (statements1 <> statements2)
        $ Instruction
        $ RiscV.BaseInstruction RiscV.Op
        $ RiscV.R register RiscV.ADD operandRegister1 operandRegister2 (RiscV.Funct7 0b0000000)
quadruple (SubtractionQuadruple operand1 operand2 (Store register)) =
    let (operandRegister1, statements1) = loadImmediateOrRegister operand1 RiscV.A0
        (operandRegister2, statements2) = loadImmediateOrRegister operand2 RiscV.A1
     in Vector.snoc (statements1 <> statements2)
        $ Instruction
        $ RiscV.BaseInstruction RiscV.Op
        $ RiscV.R register RiscV.SUB operandRegister1 operandRegister2 (RiscV.Funct7 0b0100000)

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