{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import Text.PrettyPrint.HughesPJ

import Data.Maybe

import Data.Char

import Text.Parsec.Prim hiding (State)

import Control.Monad.State

import System.IO

import System.Directory

import Control.Monad.IO.Class

import PascalPreprocessor

import Control.Exception

import System.IO.Error

import qualified Data.Map as Map

import qualified Data.Set as Set

import Data.List (find)

import Numeric

import PascalParser(pascalUnit)

import PascalUnitSyntaxTree

data InsertOption =

    IOInsert

    | IOInsertWithType Doc

    | IOLookup

    | IOLookupLast

    | IOLookupFunction Int

    | IODeferred

data Record = Record

    {

        lcaseId :: String,

        baseType :: BaseType,

        typeDecl :: Doc

    }

    deriving Show

type Records = Map.Map String [Record]

data RenderState = RenderState

    {

        currentScope :: Records,

        lastIdentifier :: String,

        lastType :: BaseType,

        lastIdTypeDecl :: Doc,

        stringConsts :: [(String, String)],

        uniqCounter :: Int,

        toMangle :: Set.Set String,

        currentUnit :: String,

        currentFunctionResult :: String,

        namespaces :: Map.Map String Records

    }

rec2Records = map (\(a, b) -> Record a b empty)

emptyState = RenderState Map.empty "" BTUnknown empty [] 0 Set.empty "" ""

getUniq :: State RenderState Int

getUniq = do

    i <- gets uniqCounter

    modify(\s -> s{uniqCounter = uniqCounter s + 1})

    return i

addStringConst :: String -> State RenderState Doc

addStringConst str = do

    strs <- gets stringConsts

    let a = find ((==) str . snd) strs

    if isJust a then

        do

        modify (\s -> s{lastType = BTString})

        return . text . fst . fromJust $ a

    else

        do

        i <- getUniq

        let sn = "__str" ++ show i

        modify (\s -> s{lastType = BTString, stringConsts = (sn, str) : strs})

        return $ text sn

escapeStr :: String -> String

escapeStr = foldr escapeChar []

escapeChar :: Char -> ShowS

escapeChar '"' s = "\\\"" ++ s

escapeChar '\\' s = "\\\\" ++ s

escapeChar a s = a : s

strInit :: String -> Doc

strInit a = text "STRINIT" <> parens (doubleQuotes (text $ escapeStr a))

renderStringConsts :: State RenderState Doc

renderStringConsts = liftM (vcat . map (\(a, b) -> text "static const string255" <+> (text a) <+> text "=" <+> strInit b <> semi))

    $ gets stringConsts

docToLower :: Doc -> Doc

docToLower = text . map toLower . render

pas2C :: String -> String -> String -> IO ()

pas2C fn inputPath outputPath = do

    s <- flip execStateT initState $ f fn

    renderCFiles s outputPath

    where

    printLn = liftIO . hPutStrLn stdout

    print = liftIO . hPutStr stdout

    initState = Map.empty

    f :: String -> StateT (Map.Map String PascalUnit) IO ()

    f fileName = do

        processed <- gets $ Map.member fileName

        unless processed $ do

            print ("Preprocessing '" ++ fileName ++ ".pas'... ")

            fc' <- liftIO

                $ tryJust (guard . isDoesNotExistError)

            case fc' of

                (Left a) -> do

                    modify (Map.insert fileName (System []))

                    printLn "doesn't exist"

                (Right fc) -> do

                    print "ok, parsing... "

                    let ptree = parse pascalUnit fileName fc

                    case ptree of

                         (Left a) -> do

                            liftIO $ writeFile (outputPath ++ "preprocess.out") fc

                            printLn $ show a ++ "\nsee preprocess.out for preprocessed source"

                            fail "stop"

                         (Right a) -> do

                            printLn "ok"

                            modify (Map.insert fileName a)

                            mapM_ f (usesFiles a)

renderCFiles :: Map.Map String PascalUnit -> String -> IO ()

renderCFiles units outputPath = do

    let u = Map.toList units

    let nss = Map.map (toNamespace nss) units

    --hPutStrLn stderr $ "Units: " ++ (show . Map.keys . Map.filter (not . Map.null) $ nss)

    --writeFile "pas2c.log" $ unlines . map (\t -> show (fst t) ++ "\n" ++ (unlines . map ((:) '\t' . show) . snd $ t)) . Map.toList $ nss

    mapM_ (toCFiles outputPath nss) u

    where

    toNamespace :: Map.Map String Records -> PascalUnit -> Records

    toNamespace nss (System tvs) =

        currentScope $ execState f (emptyState nss)

        where

        f = do

            checkDuplicateFunDecls tvs

            mapM_ (tvar2C True False True False) tvs

    toNamespace nss (Redo tvs) = -- functions that are re-implemented, add prefix to all of them

        currentScope $ execState f (emptyState nss){currentUnit = "fpcrtl_"}

        where

        f = do

            checkDuplicateFunDecls tvs

            mapM_ (tvar2C True False True False) tvs

    toNamespace _ (Program {}) = Map.empty

    toNamespace nss (Unit (Identifier i _) interface _ _ _) =

        currentScope $ execState (interface2C interface True) (emptyState nss){currentUnit = map toLower i ++ "_"}

withState' :: (RenderState -> RenderState) -> State RenderState a -> State RenderState a

withState' f sf = do

    st <- liftM f get

    let (a, s) = runState sf st

    modify(\st -> st{

        lastType = lastType s

        , uniqCounter = uniqCounter s

        , stringConsts = stringConsts s

        })

    return a

withLastIdNamespace f = do

    li <- gets lastIdentifier

    nss <- gets namespaces

    withState' (\st -> st{currentScope = fromMaybe Map.empty $ Map.lookup li (namespaces st)}) f

withRecordNamespace :: String -> [Record] -> State RenderState Doc -> State RenderState Doc

withRecordNamespace _ [] = error "withRecordNamespace: empty record"

withRecordNamespace prefix recs = withState' f

    where

        f st = st{currentScope = Map.unionWith un records (currentScope st), currentUnit = ""}

        records = Map.fromList $ map (\(Record a b d) -> (map toLower a, [Record (prefix ++ a) b d])) recs

        un [a] b = a : b

toCFiles :: String -> Map.Map String Records -> (String, PascalUnit) -> IO ()

toCFiles _ _ (_, System _) = return ()

toCFiles _ _ (_, Redo _) = return ()

toCFiles outputPath ns p@(fn, pu) = do

    hPutStrLn stdout $ "Rendering '" ++ fn ++ "'..."

    toCFiles' p

    where

    toCFiles' (fn, p@(Program {})) = writeFile (outputPath ++ fn ++ ".c") $ "#include \"fpcrtl.h\"\n" ++ (render2C initialState . pascal2C) p

    toCFiles' (fn, (Unit unitId@(Identifier i _) interface implementation _ _)) = do

        let (a, s) = runState (id2C IOInsert (setBaseType BTUnit unitId) >> interface2C interface True) initialState{currentUnit = map toLower i ++ "_"}

            (a', s') = runState (id2C IOInsert (setBaseType BTUnit unitId) >> interface2C interface False) initialState{currentUnit = map toLower i ++ "_"}

        writeFile (outputPath ++ fn ++ ".h") $ "#pragma once\n\n#include \"pas2c.h\"\n\n" ++ (render (a $+$ text ""))

        writeFile (outputPath ++ fn ++ ".c") $ "#include \"fpcrtl.h\"\n\n#include \"" ++ fn ++ ".h\"\n" ++ render (a' $+$ text "") ++ (render2C s . implementation2C) implementation

    initialState = emptyState ns

    render2C :: RenderState -> State RenderState Doc -> String

    render2C a = render . ($+$ empty) . flip evalState a

usesFiles :: PascalUnit -> [String]

usesFiles (Program _ (Implementation uses _) _) = ["pas2cSystem", "pas2cRedo"] ++ uses2List uses

usesFiles (Unit _ (Interface uses1 _) (Implementation uses2 _) _ _) = ["pas2cSystem", "pas2cRedo"] ++ uses2List uses1 ++ uses2List uses2

usesFiles (System {}) = []

usesFiles (Redo {}) = []

pascal2C :: PascalUnit -> State RenderState Doc

pascal2C (Unit _ interface implementation init fin) =

    liftM2 ($+$) (interface2C interface True) (implementation2C implementation)

pascal2C (Program _ implementation mainFunction) = do

    impl <- implementation2C implementation

    [main] <- tvar2C True False True True (FunctionDeclaration (Identifier "main" BTInt) False (SimpleType $ Identifier "int" BTInt) [VarDeclaration False False ([Identifier "argc" BTInt], SimpleType (Identifier "Integer" BTInt)) Nothing, VarDeclaration False False ([Identifier "argv" BTUnknown], SimpleType (Identifier "PPChar" BTUnknown)) Nothing] (Just (TypesAndVars [], mainFunction)))

    return $ impl $+$ main

-- the second bool indicates whether do normal interface translation or generate variable declarations

-- that will be inserted into implementation files

interface2C :: Interface -> Bool -> State RenderState Doc

interface2C (Interface uses tvars) True = do

    u <- uses2C uses

    tv <- typesAndVars2C True True True tvars

    r <- renderStringConsts

    return (u $+$ r $+$ tv)

interface2C (Interface uses tvars) False = do

    u <- uses2C uses

    tv <- typesAndVars2C True False False tvars

    r <- renderStringConsts

    return tv

implementation2C :: Implementation -> State RenderState Doc

implementation2C (Implementation uses tvars) = do

    u <- uses2C uses

    tv <- typesAndVars2C True False True tvars

    r <- renderStringConsts

    return (u $+$ r $+$ tv)

checkDuplicateFunDecls :: [TypeVarDeclaration] -> State RenderState ()

checkDuplicateFunDecls tvs =

    modify $ \s -> s{toMangle = Map.keysSet . Map.filter (> 1) . foldr ins initMap $ tvs}

    where

        initMap = Map.empty

        --initMap = Map.fromList [("reset", 2)]

        ins (FunctionDeclaration (Identifier i _) _ _ _ _) m = Map.insertWith (+) (map toLower i) 1 m

        ins _ m = m

-- the second bool indicates whether declare variable as extern or not

-- the third bool indicates whether include types or not

typesAndVars2C :: Bool -> Bool -> Bool -> TypesAndVars -> State RenderState Doc

typesAndVars2C b externVar includeType(TypesAndVars ts) = do

    checkDuplicateFunDecls ts

    liftM (vcat . map (<> semi) . concat) $ mapM (tvar2C b externVar includeType False) ts

setBaseType :: BaseType -> Identifier -> Identifier

setBaseType bt (Identifier i _) = Identifier i bt

uses2C :: Uses -> State RenderState Doc

uses2C uses@(Uses unitIds) = do

    mapM_ injectNamespace (Identifier "pas2cSystem" undefined : unitIds)

    mapM_ injectNamespace (Identifier "pas2cRedo" undefined : unitIds)

    mapM_ (id2C IOInsert . setBaseType BTUnit) unitIds

    return $ vcat . map (\i -> text $ "#include \"" ++ i ++ ".h\"") $ uses2List uses

    where

    injectNamespace (Identifier i _) = do

        getNS <- gets (flip Map.lookup . namespaces)

        modify (\s -> s{currentScope = Map.unionWith (++) (fromMaybe Map.empty (getNS i)) $ currentScope s})

uses2List :: Uses -> [String]

uses2List (Uses ids) = map (\(Identifier i _) -> i) ids

setLastIdValues vv = (\s -> s{lastType = baseType vv, lastIdentifier = lcaseId vv, lastIdTypeDecl = typeDecl vv})

id2C :: InsertOption -> Identifier -> State RenderState Doc

id2C IOInsert i = id2C (IOInsertWithType empty) i

id2C (IOInsertWithType d) (Identifier i t) = do

    ns <- gets currentScope

    tom <- gets (Set.member n . toMangle)

    cu <- gets currentUnit

    let (i', t') = case (t, tom) of

            (BTFunction _ p _, True) -> (cu ++ i ++ ('_' : show p), t)

            (BTFunction _ _ _, _) -> (cu ++ i, t)

            (BTVarParam t', _) -> ('(' : '*' : i ++ ")" , t')

            _ -> (i, t)

    modify (\s -> s{currentScope = Map.insertWith (++) n [Record i' t' d] (currentScope s), lastIdentifier = n})

    return $ text i'

    where

        n = map toLower i

id2C IOLookup i = id2CLookup head i

id2C IOLookupLast i = id2CLookup last i

id2C (IOLookupFunction params) (Identifier i t) = do

    let i' = map toLower i

    v <- gets $ Map.lookup i' . currentScope

    lt <- gets lastType

    if isNothing v then

        error $ "Not defined: '" ++ i' ++ "'\n" ++ show lt ++ "\nwith num of params = " ++ show params ++ "\n" ++ show v

        else

        let vv = fromMaybe (head $ fromJust v) . find checkParam $ fromJust v in

            modify (setLastIdValues vv) >> (return . text . lcaseId $ vv)

    where

        checkParam (Record _ (BTFunction _ p _) _) = p == params

        checkParam _ = False

id2C IODeferred (Identifier i t) = do

    let i' = map toLower i

    v <- gets $ Map.lookup i' . currentScope

    if (isNothing v) then

        modify (\s -> s{lastType = BTUnknown, lastIdentifier = i}) >> return (text i)

        else

        let vv = head $ fromJust v in modify (setLastIdValues vv) >> (return . text . lcaseId $ vv)

id2CLookup :: ([Record] -> Record) -> Identifier -> State RenderState Doc

id2CLookup f (Identifier i t) = do

    let i' = map toLower i

    v <- gets $ Map.lookup i' . currentScope

    lt <- gets lastType

    if isNothing v then

        error $ "Not defined: '" ++ i' ++ "'\n" ++ show lt

        else

        let vv = f $ fromJust v in modify (setLastIdValues vv) >> (return . text . lcaseId $ vv)

id2CTyped :: TypeDecl -> Identifier -> State RenderState Doc

id2CTyped = id2CTyped2 Nothing

id2CTyped2 :: Maybe Doc -> TypeDecl -> Identifier -> State RenderState Doc

id2CTyped2 md t (Identifier i _) = do

    tb <- resolveType t

    case (t, tb) of

        (_, BTUnknown) -> do

            error $ "id2CTyped: type BTUnknown for " ++ show i ++ "\ntype: " ++ show t

        (SimpleType {}, BTRecord _ r) -> do

            ts <- type2C t

            id2C (IOInsertWithType $ ts empty) (Identifier i (BTRecord (render $ ts empty) r))

        (_, BTRecord _ r) -> do

            ts <- type2C t

            id2C (IOInsertWithType $ ts empty) (Identifier i (BTRecord i r))

        _ -> case md of

                Nothing -> id2C IOInsert (Identifier i tb)

                Just ts -> id2C (IOInsertWithType ts) (Identifier i tb)

resolveType :: TypeDecl -> State RenderState BaseType

resolveType st@(SimpleType (Identifier i _)) = do

    let i' = map toLower i

    v <- gets $ Map.lookup i' . currentScope

    if isJust v then return . baseType . head $ fromJust v else return $ f i'

    where

    f "integer" = BTInt

    f "pointer" = BTPointerTo BTVoid

    f "boolean" = BTBool

    f "float" = BTFloat

    f "char" = BTChar

    f "string" = BTString

    f _ = error $ "Unknown system type: " ++ show st

resolveType (PointerTo (SimpleType (Identifier i _))) = return . BTPointerTo $ BTUnresolved (map toLower i)

resolveType (PointerTo t) = liftM BTPointerTo $ resolveType t

resolveType (RecordType tv mtvs) = do

    tvs <- mapM f (concat $ tv : fromMaybe [] mtvs)

    return . BTRecord "" . concat $ tvs

    where

        f :: TypeVarDeclaration -> State RenderState [(String, BaseType)]

        f (VarDeclaration _ _ (ids, td) _) = mapM (\(Identifier i _) -> liftM ((,) i) $ resolveType td) ids

resolveType (ArrayDecl (Just i) t) = do

    t' <- resolveType t

    return $ BTArray i BTInt t'

resolveType (ArrayDecl Nothing t) = liftM (BTArray RangeInfinite BTInt) $ resolveType t

resolveType (FunctionType t a) = liftM (BTFunction False (length a)) $ resolveType t

resolveType (DeriveType (InitHexNumber _)) = return BTInt

resolveType (DeriveType (InitNumber _)) = return BTInt

resolveType (DeriveType (InitFloat _)) = return BTFloat

resolveType (DeriveType (InitString _)) = return BTString

resolveType (DeriveType (InitBinOp {})) = return BTInt

resolveType (DeriveType (InitPrefixOp _ e)) = initExpr2C e >> gets lastType

resolveType (DeriveType (BuiltInFunction{})) = return BTInt

resolveType (DeriveType (InitReference (Identifier{}))) = return BTBool -- TODO: derive from actual type

resolveType (DeriveType _) = return BTUnknown

resolveType (String _) = return BTString

resolveType VoidType = return BTVoid

resolveType (Sequence ids) = return $ BTEnum $ map (\(Identifier i _) -> map toLower i) ids

resolveType (RangeType _) = return $ BTVoid

resolveType (Set t) = liftM BTSet $ resolveType t

resolveType (VarParamType t) = liftM BTVarParam $ resolveType t

resolve :: String -> BaseType -> State RenderState BaseType

resolve s (BTUnresolved t) = do

    v <- gets $ Map.lookup t . currentScope

    if isJust v then

        resolve s . baseType . head . fromJust $ v

        else

        error $ "Unknown type " ++ show t ++ "\n" ++ s

resolve _ t = return t

fromPointer :: String -> BaseType -> State RenderState BaseType

fromPointer s (BTPointerTo t) = resolve s t

fromPointer s t = do

    error $ "Dereferencing from non-pointer type " ++ show t ++ "\n" ++ s

functionParams2C params = liftM (hcat . punctuate comma . concat) $ mapM (tvar2C False False True True) params

numberOfDeclarations :: [TypeVarDeclaration] -> Int

numberOfDeclarations = sum . map cnt

    where

        cnt (VarDeclaration _ _ (ids, _) _) = length ids

        cnt _ = 1

hasPassByReference :: [TypeVarDeclaration] -> Bool

hasPassByReference = or . map isVar

    where

        isVar (VarDeclaration v _ (_, _) _) = v

        isVar _ = error $ "hasPassByReference called not on function parameters"

toIsVarList :: [TypeVarDeclaration] -> [Bool]

toIsVarList = concatMap isVar

    where

        isVar (VarDeclaration v _ (p, _) _) = replicate (length p) v

        isVar _ = error $ "toIsVarList called not on function parameters"

funWithVarsToDefine :: String -> [TypeVarDeclaration] -> Doc

funWithVarsToDefine n params = text "#define" <+> text n <> parens abc <+> text (n ++ "__vars") <> parens cparams

    where

        abc = hcat . punctuate comma . map (char . fst) $ ps

        cparams = hcat . punctuate comma . map (\(c, v) -> if v then char '&' <> parens (char c) else char c) $ ps

        ps = zip ['a'..] (toIsVarList params)

fun2C :: Bool -> String -> TypeVarDeclaration -> State RenderState [Doc]

fun2C _ _ (FunctionDeclaration name inline returnType params Nothing) = do

    t <- type2C returnType

    t'<- gets lastType

    p <- withState' id $ functionParams2C params

    n <- liftM render . id2C IOInsert $ setBaseType (BTFunction hasVars (numberOfDeclarations params) t') name

    let decor = if inline then text "inline" else empty

    if hasVars then

        return [funWithVarsToDefine n params $+$ decor <+> t empty <+> text (n ++ "__vars") <> parens p]

        else

        return [decor <+> t empty <+> text n <> parens p]

    where

        hasVars = hasPassByReference params

fun2C True rv (FunctionDeclaration name@(Identifier i _) inline returnType params (Just (tvars, phrase))) = do

    let res = docToLower $ text rv <> text "_result"

    t <- type2C returnType

    t'<- gets lastType

    notDeclared <- liftM isNothing . gets $ Map.lookup (map toLower i) . currentScope

    n <- liftM render . id2C IOInsert $ setBaseType (BTFunction hasVars (numberOfDeclarations params) t') name

    let isVoid = case returnType of

            VoidType -> True

            _ -> False

    (p, ph) <- withState' (\st -> st{currentScope = Map.insertWith un (map toLower rv) [Record (render res) t' empty] $ currentScope st

            , currentFunctionResult = if isVoid then [] else render res}) $ do

        p <- functionParams2C params

        ph <- liftM2 ($+$) (typesAndVars2C False False True tvars) (phrase2C' phrase)

        return (p, ph)

    let phrasesBlock = if isVoid then ph else t empty <+> res <> semi $+$ ph $+$ text "return" <+> res <> semi

    let define = if hasVars then text "#ifndef" <+> text n $+$ funWithVarsToDefine n params $+$ text "#endif" else empty

    let decor = if inline then text "inline" else empty

    return [

        define

        $+$

        --(if notDeclared && hasVars then funWithVarsToDefine n params else empty) $+$

        decor <+> t empty <+> text (if hasVars then n ++ "__vars" else n) <> parens p

        $+$

        text "{"

        $+$

        nest 4 phrasesBlock

        $+$

        text "}"]