{-# 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 qualified Data.Set as Set

import Data.List (find)

import Numeric

import PascalParser

import PascalUnitSyntaxTree

data InsertOption = 

    IOInsert

    | IOLookup

    | IOLookupLast

    | IOLookupFunction Int

    | IODeferred

type Record = (String, BaseType)

type Records = Map.Map String [Record]

data RenderState = RenderState 

    {

        currentScope :: Records,

        lastIdentifier :: String,

        lastType :: BaseType,

        stringConsts :: [(String, String)],

        uniqCounter :: Int,

        toMangle :: Set.Set String,

        currentUnit :: String,

        namespaces :: Map.Map String Records

    }

emptyState = RenderState Map.empty "" BTUnknown [] 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 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 . Map.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 Records -> PascalUnit -> Records

    toNamespace nss (System tvs) = 

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

    toNamespace _ (Program {}) = Map.empty

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

        currentScope $ execState (interface2C interface) (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 :: State RenderState Doc -> State RenderState Doc

withLastIdNamespace f = do

    li <- gets lastIdentifier

    nss <- gets namespaces

    withState' (\st -> st{currentScope = fromMaybe Map.empty $ 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 = Map.unionWith un records (currentScope st), currentUnit = ""}

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

        un [a] b = a : b

toCFiles :: Map.Map String Records -> (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@(Identifier i _) interface implementation _ _)) = do

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

        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)

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

checkDuplicateFunDecls tvs =

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

    where

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

        ins _ m = m

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

typesAndVars2C b (TypesAndVars ts) = do

    checkDuplicateFunDecls 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)

        modify (\s -> s{currentScope = Map.unionWith (++) (fromMaybe Map.empty (getNS i)) $ 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

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

    cu <- gets currentUnit

    let i' = case (t, tom) of

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

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

            _ -> i

    modify (\s -> s{currentScope = Map.insertWith (++) n [(i', t)] (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 (\s -> s{lastType = snd vv, lastIdentifier = fst vv}) >> (return . text . fst $ vv)

    where

        checkParam (_, 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 (\s -> s{lastType = snd vv, lastIdentifier = fst vv}) >> (return . text . fst $ vv)

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

id2CLookup f (Identifier i _) = 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 (\s -> s{lastType = snd vv, lastIdentifier = fst vv}) >> (return . text . fst $ vv)

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

id2CTyped 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 IOInsert (Identifier i (BTRecord (render $ ts empty) r))

        (_, BTRecord _ r) -> do

            ts <- type2C t

            id2C IOInsert (Identifier i (BTRecord i r))

        _ -> id2C IOInsert (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 . snd . 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 (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 {})) = 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 $ Map.lookup t . currentScope

    if isJust v then

        resolve s . snd . 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) params

numberOfDeclarations :: [TypeVarDeclaration] -> Int

numberOfDeclarations = sum . map cnt

    where

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

        cnt _ = 1

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 (numberOfDeclarations params) 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 (numberOfDeclarations params) t') name

    (p, ph) <- withState' (\st -> st{currentScope = Map.insertWith un (map toLower rv) [(render res, 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

    un [a] b = a : b

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 <- 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)]

         (True, BTFloat, [i], Just e) -> do

             i' <- id2CTyped t i

             ie <- initExpr2C e

             return [text "#define" <+> i' <+> parens ie <> text "\n"]

         (_, BTFunction{}, _, Nothing) -> liftM (map(\i -> t' $ text "*" <+> i)) $ mapM (id2CTyped t) ids

         _ -> 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 $ Identifier (t' ++ "_op_" ++ opStr) bt

    where 

    opStr = case op of

                    "+" -> "add"

                    "-" -> "sub"

                    "*" -> "mul"

                    "/" -> "div"

                    "=" -> "eq"

                    "<" -> "lt"

                    ">" -> "gt"

                    "<>" -> "neq"

                    _ -> 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, initExpr2C' :: InitExpression -> State RenderState Doc

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

initExpr2C a = initExpr2C' a

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' 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