{-# LANGUAGE ScopedTypeVariables #-}

module AST where

-- this file mostly just contains the operations neccisary for basic operations
-- on terms, like binding, and "compilation" to "bytecode"

import           Data.Serialize
import           Text.Megaparsec    (Pos, SourcePos (..), unPos, unsafePos)

import           Data.Text          (Text, length, pack)
import           Data.Text.Encoding
import           Data.Version
import           Data.Word          (Word64)
import           GHC.Generics

-- represent a term with variable name type v, and variable binding index type i that has metadata m
-- associated to each node and leaf
data Term m v i = Variable v i m
                | Abstraction v m (Term m v i) (Term m v i)
                | Application m (Term m v i) (Term m v i)
                | Type m
  deriving (Generic)

-- kinds are needed for the representation of the types of terms whose tail
-- is not a variable, but Type

data Kind m v i = Function v m (Term m v i) (Kind m v i)
                | Kind m
  deriving (Generic)

-- pseudoterms are either terms or kinds, and are needed as the return types of
-- many of the functions in the typing algorithm

data PseudoTerm m v i = Term (Term m v i)
                      | Kind' (Kind m v i)

-- compiled files have a special header that encodes much of the information
-- about the version of the compiler needed

data DecompiledFileHeader =
       DecompiledFileHeader
         { compiledFileName   :: Text
         , interpreterVersion :: Version
         , extensionsUsed     :: [Text]
         }

-- this is the actual file representation of the file after compilation

data DecompiledFile m v i =
       DecompiledFile
         { decompiledFileHeader :: DecompiledFileHeader
         , decompiledTerm       :: Term m v i
         }

instance (Serialize m, Serialize v, Serialize i) =>
         Serialize (Term m v i)

instance (Serialize m, Serialize v, Serialize i) =>
         Serialize (Kind m v i)

  -- some orphan instances
instance Serialize Version where
  put vers = putListOf (put :: Putter Int) (versionBranch vers)
  get = do
    branchList :: [Int] <- getListOf get
    return Version { versionBranch = branchList, versionTags = [] }

instance Serialize Text where
  put text' = do
    putWord64le (fromIntegral (Data.Text.length text') :: Word64)
    put (encodeUtf8 text')
  get = do
    leng <- getWord64le
    bs <- getByteString (fromIntegral leng :: Int)
    return (decodeUtf8 bs)

instance Serialize SourcePos where
  put srcPos = do
    put (sourceName srcPos)
    put (sourceLine srcPos)
    put (sourceColumn srcPos)
  get = do
    str :: String <- get
    pos1 :: Pos <- get
    pos2 :: Pos <- get
    return $ SourcePos str pos1 pos2

instance Serialize Pos where
  put pos = putWord64le (fromIntegral (unPos pos) :: Word64)
  get = fmap (unsafePos . fromIntegral) getWord64le

instance Serialize DecompiledFileHeader where
  put decompHeader = do
    put (compiledFileName decompHeader)
    put (interpreterVersion decompHeader)
    put (extensionsUsed decompHeader)
  get = do
    fname :: Text <- get
    iversion :: Version <- get
    xtensions :: [Text] <- get
    return
      DecompiledFileHeader
        { compiledFileName = fname
        , interpreterVersion = iversion
        , extensionsUsed = xtensions
        }

instance (Serialize m, Serialize v, Serialize i) =>
         Serialize (DecompiledFile m v i) where
  put decomp = do
    put (decompiledFileHeader decomp)
    put (decompiledTerm decomp)
  get = do
    fileHeader :: DecompiledFileHeader <- get
    fileBody :: Term m v i <- get
    return DecompiledFile { decompiledFileHeader = fileHeader, decompiledTerm = fileBody }

-- transpile is used to change the metadata, variable, and index datatypes of
-- the file

transpile :: (m -> m') -> (v -> v') -> (i -> i') -> DecompiledFile m v i -> DecompiledFile m' v' i'
transpile fm fv fi decompFile@DecompiledFile { decompiledTerm = fileBody } =
  let transpile' t =
        case t of
          Type m -> Type (fm m)
          Variable v i m -> Variable (fv v) (fi i) (fm m)
          Application m a f -> Application (fm m) (transpile' a) (transpile' f)
          Abstraction v m t b ->
            Abstraction (fv v) (fm m) (transpile' t) (transpile' b)
  in decompFile { decompiledTerm = transpile' fileBody }

getNodeMetadata :: PseudoTerm m v i -> m
getNodeMetadata node =
  case node of
    Term term' ->
      case term' of
        Type metadata              -> metadata
        Variable _ _ metadata      -> metadata
        Abstraction _ metadata _ _ -> metadata
        Application metadata _ _   -> metadata
    Kind' kind ->
      case kind of
        Kind metadata           -> metadata
        Function _ metadata _ _ -> metadata

bindTerm :: (Eq v, Eq i, Enum i) => v -> i -> Term m v i -> Term m v i
bindTerm variable index subject =
  case subject of
    Type metadata -> Type metadata
    Application metadata argument function ->
      Application metadata (bindTerm variable index argument) (bindTerm variable index function)
    Abstraction variable' metadata parameter body ->
      Abstraction variable' metadata (bindTerm variable index parameter)
        (bindTerm variable (succ index) body)
    Variable variable' index' metadata
      | variable == variable' && index' == toEnum 0 -> Variable variable' index metadata
      | otherwise -> subject

bindContext :: (Eq i, Enum i, Show i) => [Term m Text i] -> i -> Term m Text i -> Term m Text i
bindContext context index term =
  case context of
    []     -> term
    _:rest -> bindContext rest (succ index) (bindTerm (pack ("$" ++ show index)) index term)