module Bril (
    Function (..), Type (..), Value (..), InstrOperation (..), FunctionArg (..), Program (..),
    Instruction (..),
    parseValue,
    parseBrilFromPath, parseBrilJSON,
    )where

import Data.Aeson (FromJSON, parseJSON, withObject, withText, (.:), (.:?), eitherDecode)
import Data.Aeson.Types (Parser, modifyFailure)
import qualified Data.Aeson (Value, Object, Key)
import Data.Aeson.KeyMap
import Data.Int (Int64)
import qualified Data.Text as T
import Data.Text.Read (Reader, signed, decimal)
import Data.SortedList
import Data.Maybe (fromMaybe, isJust)

import qualified Data.ByteString.Lazy as LB

data Program = Program { programFns :: [Function] }
    deriving (Show, Eq)

type VarName = String
type FuncName = String

data Function = Function {
        functionName :: FuncName,
        returnType :: Maybe Type,
        functionArgs :: [FunctionArg],
        instructions :: [Instruction]
    } deriving (Show, Eq)

data Type = Int | Bool
    deriving (Show, Eq)
data Value = IntValue Int64 | BoolValue Bool
    deriving (Eq, Show)

data FunctionArg = FunctionArg {
        argName :: VarName,
        argType :: Type
    } deriving (Show, Eq)

data Instruction =
    Label String
  | LiteralInstr {
        value :: Value,
        dest :: VarName
    }
  | ValueInstr {
       op :: InstrOperation,
       type_ :: Type,
       dest :: VarName,
       args :: [VarName],
       funcs :: [FuncName],
       labels :: [String]
     }
   | EffectInstr {
        op :: InstrOperation,
        args :: [VarName],
        funcs :: [FuncName],
        labels :: [String]
     }
   deriving (Show, Eq)

data InstrOperation =
    Add | Mul | Sub | Div | Eq | Lt | Gt | Le | Ge | Not | And | Or |
    Jmp | Br | Call | Ret | Id | Print | Nop | Phi
   deriving (Eq, Show)

data Arity =
    Arity Int Int Int -- min, max, labels arity
  | VarArgs
   deriving (Eq, Show)

valueOps = toSortedList ([
        "add", "mul", "sub", "div", "eq", "lt", "gt", "le", "ge", "not", "and", "or"
    ] :: [T.Text])

effectOps = toSortedList ([
        "nop", "call", "jmp", "br", "print", "lt", "gt", "le", "ge", "not", "and", "or"
    ] :: [T.Text])

arity instr
    | instr `elem` [Add, Mul, Sub, Div, Eq, Lt, Gt, Le, Ge, And, Or] = Arity 2 2 0
    | instr `elem` [Id, Not] = Arity 1 1 0
    | instr == Ret  = Arity 0 1 0
    | instr == Jmp = Arity 0 0 1
    | instr == Br = Arity 1 1 2
    | instr == Nop = Arity 0 0 0
    | otherwise = VarArgs

checkArity :: InstrOperation -> [b] -> [c] -> Bool
checkArity instr args labels =
    case arity instr of
        Arity n m l ->
            (((length args) >= n) && ((length args) <= m)) && ((length labels) == l)
        VarArgs -> True

instrFromString s
    | s == "add" = Add
    | s == "sub" = Sub
    | s == "mul" = Mul
    | s == "eq" = Eq
    | s == "lt" = Lt
    | s == "gt" = Gt
    | s == "le" = Le
    | s == "ge" = Ge
    | s == "not" = Not
    | s == "and" = And
    | s == "or" = Or
    | s == "jmp" = Jmp
    | s == "br" = Br
    | s == "ret" = Ret
    | s == "call" = Call
    | s == "print" = Print
    | s == "id" = Id
    | s == "nop" = Nop
    | otherwise = error "wrong instruction name"

optionalArrayField :: FromJSON a => Data.Aeson.Object -> Data.Aeson.Key -> Parser [a]
optionalArrayField v field =
    ((v .:? field) >>= (return . fromMaybe []))


instance FromJSON Program  where
    parseJSON = withObject "Program" $ \v -> Program
        <$> v .: "functions"

instance FromJSON Function where
    parseJSON = withObject "Function" $ \v -> Function
        <$> v .: "name"
        <*> v .:? "type"
        <*> v `optionalArrayField` "args"
        <*> v .: "instrs"

instance FromJSON Instruction where
    parseJSON = withObject "Instruction" $ \v ->
        if (Data.Aeson.KeyMap.member "label" v)
            then Label <$> v .: "label"
            else do
                op <- (v .: "op") :: Parser T.Text
                parseInstr op v
     where
        parseInstr op
            | op == "const" = parseLiteralInstr
            | op `elemOrd` valueOps = parseValueInstr op
            | op `elemOrd` effectOps = parseEffectInstr op
            | otherwise = const $ fail ("Invalid op: " ++ (T.unpack op))
        parseLiteralInstr v = do
            type_ <- (v .: "type") :: Parser Type
            LiteralInstr
                <$> ((v .: "value") >>= (parseValue type_))
                <*> v .: "dest"
        parseValueInstr op v = do
            dest  <- (v .: "dest") :: Parser VarName
            args  <- (v `optionalArrayField` "args") :: Parser [VarName]
            labels  <- (v `optionalArrayField` "labels") :: Parser [VarName]
            funcs  <- (v `optionalArrayField` "funcs") :: Parser [VarName]
            let instr = instrFromString op
            if not (checkArity instr args labels)
                then fail $ unexpectedArityError op args labels
                else ValueInstr
                    <$> return instr
                    <*> v .: "type"
                    <*> v .: "dest"
                    <*> return args
                    <*> return funcs
                    <*> return labels
        parseEffectInstr op v = do
            args  <- (v `optionalArrayField` "args") :: Parser [VarName]
            labels  <- (v `optionalArrayField` "labels") :: Parser [VarName]
            funcs  <- (v `optionalArrayField` "funcs") :: Parser [VarName]
            dest <- (v .:? "dest") :: Parser (Maybe VarName)
            if (op == "call") && ((length funcs) /= 1)
                then fail $ "call op must have exactly one entry in \"funcs\""
                else if (isJust dest) && (op == "call")
                    then parseValueInstr op v
                    else do
                        let instr = instrFromString op
                        if not (checkArity instr args labels)
                            then fail $ unexpectedArityError op args labels
                            else EffectInstr
                                <$> return instr
                                <*> return args
                                <*> return funcs
                                <*> return labels

        unexpectedArityError op args labels =
            "wrong arity for " ++ (T.unpack op) ++ ": got " ++ (show (length args)) ++ ", " ++ (show (length labels))

instance FromJSON FunctionArg where
    parseJSON = withObject "FunctionArg" $ \v -> FunctionArg
        <$> v .: "name"
        <*> v .: "type"

instance FromJSON Type where
    parseJSON = withText "Type" $ \v ->
        case v of
             "int" -> return Int
             "bool" -> return Bool
             _ -> fail "wrong type"

parseValue :: Type -> Data.Aeson.Value -> Parser Value
parseValue Bool =
    withText "BoolValue" $ \v ->
        case v of
             "true" -> return $ BoolValue True
             "false" -> return $ BoolValue False
             _ -> fail $ "invalid bool literal: " ++ (T.unpack v)
parseValue Int =
    withText "IntValue" $ \v ->
        case ((signed decimal :: Reader Int64) v) of
            Left x -> fail $ "invalid int literal: " ++ x
            Right (val, t) ->
                if (T.null t)
                    then return $ IntValue val
                    else fail $ "invalid int literal: " ++ (T.unpack v)
            

parseBrilJSON text = (eitherDecode text :: Either String Program)

parseBrilFromPath path = do
    contents <- LB.readFile path
    return $ parseBrilJSON contents