module Language.Elna.CodeGenerator
    ( Statement(..)
    , generateCode
    ) 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.Location (Identifier(..))
import Language.Elna.Intermediate (Operand(..), Quadruple(..), Variable(..))
import qualified Language.Elna.Architecture.RiscV as RiscV
import Language.Elna.SymbolTable (SymbolTable)
import Data.Bits (Bits(..))

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

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

generateCode :: SymbolTable -> HashMap Identifier (Vector Quadruple) -> Vector Statement
generateCode _ = 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 :: Quadruple -> 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 (TempVariable _)) =
    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 RiscV.T0 RiscV.ADD operandRegister1 operandRegister2 (RiscV.Funct7 0b0000000)
quadruple (SubtractionQuadruple operand1 operand2 (TempVariable _)) =
    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 RiscV.T0 RiscV.SUB operandRegister1 operandRegister2 (RiscV.Funct7 0b0100000)

loadImmediateOrRegister :: Operand -> RiscV.XRegister -> (RiscV.XRegister, Vector Statement)
loadImmediateOrRegister (IntOperand intValue) targetRegister =
    (targetRegister, lui intValue targetRegister)
loadImmediateOrRegister (VariableOperand _) _ = (RiscV.T0, 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