{-# LANGUAGE ScopedTypeVariables #-}

module Pas2C 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 Data.List (find)

import Numeric

import PascalParser

import PascalUnitSyntaxTree

data InsertOption = 

    IOInsert

    | IOLookup

    | IODeferred

type Record = (String, (String, BaseType))

data RenderState = RenderState 

    {

        currentScope :: [Record],

        lastIdentifier :: String,

        lastType :: BaseType,

        stringConsts :: [(String, String)],

        uniqCounter :: Int,

        namespaces :: Map.Map String [Record]

    }

emptyState = RenderState [] "" BTUnknown [] 0

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 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 "const string255" <+> (text a) <+> text "=" <+> strInit b <> semi)) 

    $ gets stringConsts

docToLower :: Doc -> Doc

docToLower = text . map toLower . render

pas2C :: String -> IO ()

pas2C fn = do

    setCurrentDirectory "../hedgewars/"

    s <- flip execStateT initState $ f fn

    renderCFiles s

    where

    printLn = liftIO . hPutStrLn stderr

    print = liftIO . hPutStr stderr

    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) 

                $ preprocess (fileName ++ ".pas")

            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 "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 -> IO ()

renderCFiles units = do

    let u = Map.toList units

    let nss = Map.map (toNamespace nss) units

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

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

    mapM_ (toCFiles nss) u

    where

    toNamespace :: Map.Map String [Record] -> PascalUnit -> [Record]

    toNamespace nss (System tvs) = 

        currentScope $ execState (mapM_ (tvar2C True) tvs) (emptyState nss)

    toNamespace _ (Program {}) = []

    toNamespace nss (Unit _ interface _ _ _) = 

        currentScope $ execState (interface2C interface) (emptyState nss)

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 :: State RenderState Doc -> State RenderState Doc

withLastIdNamespace f = do

    li <- gets lastIdentifier

    nss <- gets namespaces

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

withRecordNamespace :: String -> [(String, BaseType)] -> State RenderState Doc -> State RenderState Doc

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

withRecordNamespace prefix recs = withState' f

    where

        f st = st{currentScope = records ++ currentScope st}

        records = map (\(a, b) -> (map toLower a, (prefix ++ a, b))) recs

toCFiles :: Map.Map String [Record] -> (String, PascalUnit) -> IO ()

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

toCFiles ns p@(fn, pu) = do

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

    toCFiles' p

    where

    toCFiles' (fn, p@(Program {})) = writeFile (fn ++ ".c") $ (render2C initialState . pascal2C) p

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

        let (a, s) = runState (id2C IOInsert (setBaseType BTUnit unitId) >> interface2C interface) initialState

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

        writeFile (fn ++ ".c") $ "#include \"" ++ fn ++ ".h\"\n" ++ (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" : uses2List uses

usesFiles (Unit _ (Interface uses1 _) (Implementation uses2 _) _ _) = "pas2cSystem" : uses2List uses1 ++ uses2List uses2

usesFiles (System {}) = []

pascal2C :: PascalUnit -> State RenderState Doc

pascal2C (Unit _ interface implementation init fin) =

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

pascal2C (Program _ implementation mainFunction) = do

    impl <- implementation2C implementation

    [main] <- tvar2C True 

        (FunctionDeclaration (Identifier "main" BTInt) (SimpleType $ Identifier "int" BTInt) [] (Just (TypesAndVars [], mainFunction)))

    return $ impl $+$ main

interface2C :: Interface -> State RenderState Doc

interface2C (Interface uses tvars) = do

    u <- uses2C uses

    tv <- typesAndVars2C True tvars

    r <- renderStringConsts

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

implementation2C :: Implementation -> State RenderState Doc

implementation2C (Implementation uses tvars) = do

    u <- uses2C uses

    tv <- typesAndVars2C True tvars

    r <- renderStringConsts

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

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

typesAndVars2C b (TypesAndVars ts) = liftM (vcat . map (<> semi) . concat) $ mapM (tvar2C b) 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_ (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)

        let f = flip (foldl (flip (:))) (fromMaybe [] (getNS i))

        modify (\s -> s{currentScope = f $ currentScope s})

uses2List :: Uses -> [String]

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

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

id2C IOInsert (Identifier i t) = do

    ns <- gets currentScope

{--    case t of 

        BTUnknown -> do

            ns <- gets currentScope

            error $ "id2C IOInsert: type BTUnknown for " ++ show i ++ "\nnamespace: " ++ show (take 100 ns)

        _ -> do --}

    modify (\s -> s{currentScope = (n, (i, t)) : currentScope s, lastIdentifier = n})

    return $ text i

    where

        n = map toLower i

id2C IOLookup (Identifier i t) = do

    let i' = map toLower i

    v <- gets $ find (\(a, _) -> a == i') . currentScope

    ns <- gets currentScope

    lt <- gets lastType

    if isNothing v then 

        error $ "Not defined: '" ++ i' ++ "'\n" ++ show lt ++ "\n" ++ show (take 100 ns)

        else 

        let vv = snd $ fromJust v in modify (\s -> s{lastType = snd vv, lastIdentifier = fst vv}) >> (return . text . fst $ vv)

id2C IODeferred (Identifier i t) = do

    let i' = map toLower i

    v <- gets $ find (\(a, _) -> a == i') . currentScope

    if (isNothing v) then

        return $ text i

        else

        return . text . fst . snd . fromJust $ v

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

id2CTyped t (Identifier i _) = do

    tb <- resolveType t

    ns <- gets currentScope

    case tb of 

        BTUnknown -> do

            ns <- gets currentScope

            error $ "id2CTyped: type BTUnknown for " ++ show i ++ "\ntype: " ++ show t ++ "\nnamespace: " ++ show (take 100 ns)

        _ -> return ()

    id2C IOInsert (Identifier i tb)

resolveType :: TypeDecl -> State RenderState BaseType

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

    let i' = map toLower i

    v <- gets $ find (\(a, _) -> a == i') . currentScope

    if isJust v then return . snd . snd $ 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 _) = liftM BTFunction $ 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 {})) = return BTInt

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

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

resolve s (BTUnresolved t) = do

    v <- gets $ find (\(a, _) -> a == t) . currentScope

    if isJust v then

        resolve s . snd . snd . 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 (BTFunctionReturn _ (BTPointerTo t)) = resolve s t

fromPointer s t = do

    ns <- gets currentScope

    error $ "Dereferencing from non-pointer type " ++ show t ++ "\n" ++ s ++ "\n\n" ++ show (take 100 ns)

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

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

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

    t <- type2C returnType 

    t'<- gets lastType

    p <- withState' id $ functionParams2C params

    n <- id2C IOInsert $ setBaseType (BTFunction t') name

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

fun2C True rv (FunctionDeclaration name returnType params (Just (tvars, phrase))) = do

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

    t <- type2C returnType

    t'<- gets lastType

    n <- id2C IOInsert $ setBaseType (BTFunction t') name

    (p, ph) <- withState' (\st -> st{currentScope = (map toLower rv, (render res, BTFunctionReturn (render n) t')) : currentScope st}) $ do

        p <- functionParams2C params

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

        return (p, ph)

    let phrasesBlock = case returnType of

            VoidType -> ph

            _ -> t empty <+> res <> semi $+$ ph $+$ text "return" <+> res <> semi

    return [ 

        t empty <+> n <> parens p

        $+$

        text "{" 

        $+$ 

        nest 4 phrasesBlock

        $+$

        text "}"]

    where

    phrase2C' (Phrases p) = liftM vcat $ mapM phrase2C p

    phrase2C' p = phrase2C p

fun2C False _ (FunctionDeclaration (Identifier name _) _ _ _) = error $ "nested functions not allowed: " ++ name

fun2C _ tv _ = error $ "fun2C: I don't render " ++ show tv

tvar2C :: Bool -> TypeVarDeclaration -> State RenderState [Doc]

tvar2C b f@(FunctionDeclaration (Identifier name _) _ _ _) =

    fun2C b name f

tvar2C _ td@(TypeDeclaration i' t) = do

    i <- id2CTyped t i'

    tp <- case t of

        FunctionType {} -> type2C (PointerTo t)

        _ -> type2C t

    return [text "typedef" <+> tp i]

tvar2C _ (VarDeclaration isConst (ids, t) mInitExpr) = do

    t' <- liftM (((if isConst then text "const" else empty) <+>) . ) $ type2C t

    ie <- initExpr mInitExpr

    lt <- gets lastType

    case (isConst, lt, ids, mInitExpr) of

         (True, BTInt, [i], Just _) -> do

             i' <- id2CTyped t i

             return [text "enum" <> braces (i' <+> ie)]

         _ -> liftM (map(\i -> t' i <+> ie)) $ mapM (id2CTyped t) ids

    where

    initExpr Nothing = return $ empty

    initExpr (Just e) = liftM (text "=" <+>) (initExpr2C e)

tvar2C f (OperatorDeclaration op (Identifier i _) ret params body) = do

    r <- op2CTyped op (extractTypes params)

    fun2C f i (FunctionDeclaration r ret params body)

op2CTyped :: String -> [TypeDecl] -> State RenderState Identifier

op2CTyped op t = do

    t' <- liftM (render . hcat . punctuate (char '_') . map (\t -> t empty)) $ mapM type2C t

    bt <- gets lastType

    return $ case bt of

         BTRecord {} -> Identifier (t' ++ "_op_" ++ opStr) bt

         _ -> Identifier t' bt

    where 

    opStr = case op of

                    "+" -> "add"

                    "-" -> "sub"

                    "*" -> "mul"

                    "/" -> "div"

                    "=" -> "eq"

                    "<" -> "lt"

                    ">" -> "gt"

                    _ -> error $ "op2CTyped: unknown op '" ++ op ++ "'"

extractTypes :: [TypeVarDeclaration] -> [TypeDecl]

extractTypes = concatMap f

    where

        f (VarDeclaration _ (ids, t) _) = replicate (length ids) t

        f a = error $ "extractTypes: can't extract from " ++ show a

initExpr2C :: InitExpression -> State RenderState Doc

initExpr2C InitNull = return $ text "NULL"

initExpr2C (InitAddress expr) = liftM ((<>) (text "&")) (initExpr2C expr)

initExpr2C (InitPrefixOp op expr) = liftM (text (op2C op) <>) (initExpr2C expr)

initExpr2C (InitBinOp op expr1 expr2) = do

    e1 <- initExpr2C expr1

    e2 <- initExpr2C expr2

    return $ parens $ e1 <+> text (op2C op) <+> e2

initExpr2C (InitNumber s) = return $ text s

initExpr2C (InitFloat s) = return $ text s

initExpr2C (InitHexNumber s) = return $ text "0x" <> (text . map toLower $ s)

initExpr2C (InitString [a]) = return . quotes $ text [a]

initExpr2C (InitString s) = return $ strInit s

initExpr2C (InitChar a) = return $ quotes $ text "\\x" <> text (showHex (read a) "")

initExpr2C (InitReference i) = id2C IOLookup i

initExpr2C (InitRecord fields) = do

    (fs :: [Doc]) <- mapM (\(Identifier a _, b) -> liftM (text "." <> text a <+> equals <+>) $ initExpr2C b) fields

    return $ lbrace $+$ (nest 4 . vcat . punctuate comma $ fs) $+$ rbrace

initExpr2C (InitArray [value]) = initExpr2C value

initExpr2C (InitArray values) = liftM (braces . vcat . punctuate comma) $ mapM initExpr2C values

initExpr2C r@(InitRange (Range i@(Identifier i' _))) = do

    id2C IOLookup i

    t <- gets lastType

    case t of

         BTEnum s -> return . int $ length s

         BTInt -> case i' of

                       "byte" -> return $ int 256

                       _ -> error $ "InitRange identifier: " ++ i'

         _ -> error $ "InitRange: " ++ show r

initExpr2C (InitRange (RangeFromTo (InitNumber "0") r)) = initExpr2C $ BuiltInFunction "succ" [r]

initExpr2C (InitRange (RangeFromTo (InitChar "0") (InitChar r))) = initExpr2C $ BuiltInFunction "succ" [InitNumber r]

initExpr2C (InitRange a) = error $ show a --return $ text "<<range>>"

initExpr2C (InitSet []) = return $ text "0"

initExpr2C (InitSet a) = return $ text "<<set>>"

initExpr2C (BuiltInFunction "low" [InitReference e]) = return $ 

    case e of

         (Identifier "LongInt" _) -> int (-2^31)

         (Identifier "SmallInt" _) -> int (-2^15)

         _ -> error $ "BuiltInFunction 'low': " ++ show e

initExpr2C (BuiltInFunction "high" [e]) = do

    initExpr2C e

    t <- gets lastType

    case t of

         (BTArray i _ _) -> initExpr2C $ BuiltInFunction "pred" [InitRange i]

         a -> error $ "BuiltInFunction 'high': " ++ show a

initExpr2C (BuiltInFunction "succ" [BuiltInFunction "pred" [e]]) = initExpr2C e

initExpr2C (BuiltInFunction "pred" [BuiltInFunction "succ" [e]]) = initExpr2C e

initExpr2C (BuiltInFunction "succ" [e]) = liftM (<> text " + 1") $ initExpr2C e

initExpr2C (BuiltInFunction "pred" [e]) = liftM (<> text " - 1") $ initExpr2C e

initExpr2C b@(BuiltInFunction _ _) = error $ show b    

initExpr2C a = error $ "initExpr2C: don't know how to render " ++ show a

range2C :: InitExpression -> State RenderState [Doc]

range2C (InitString [a]) = return [quotes $ text [a]]

range2C (InitRange (Range i)) = liftM (flip (:) []) $ id2C IOLookup i

range2C (InitRange (RangeFromTo (InitString [a]) (InitString [b]))) = return $ map (\i -> quotes $ text [i]) [a..b]

range2C a = liftM (flip (:) []) $ initExpr2C a

baseType2C :: String -> BaseType -> Doc

baseType2C _ BTFloat = text "float"

baseType2C _ BTBool = text "bool"

baseType2C _ BTString = text "string255"

baseType2C s a = error $ "baseType2C: " ++ show a ++ "\n" ++ s

type2C :: TypeDecl -> State RenderState (Doc -> Doc)

type2C (SimpleType i) = liftM (\i a -> i <+> a) $ id2C IOLookup i

type2C t = do

    r <- type2C' t