module Language.BASIC.Translate(translateBASIC) where
import Control.Monad
import Data.List
import qualified Data.Map as M
import Data.Map((!), fromList)
import Data.Word

import LLVM.Core
import LLVM.Util.File

import Language.BASIC.Types

renumber :: [Expr a] -> [Expr a]
renumber cs =
    let m = M.fromList $ zip (map cmdLabel cs) [10, 20 ..]
        ren (Cmd l c es) = Cmd (m M.! l) c (map ren es)
	ren (Label l) = Label (m M.! l)
	ren e = e
    in  map ren cs

-- This assumes some sanity in loop nesting.
removeFor :: [Expr a] -> [Expr a]
removeFor [] = []
removeFor (Cmd l For [v, lo, hi] : cs) =
    let cs' = removeFor cs
        (n, cs'') = removeNext cs'
	removeNext [] = error $ "No NEXT for line " ++ show (l, v)
	removeNext (Cmd ln Next [v'] : bs) | v == v' = (ln+2,
	    [Cmd ln Let [v, Binop v "+" (Dbl 1)], Cmd (ln+1) Goto [Label (l+1)], Cmd (ln+2) Rem []] ++ bs)
	removeNext (c:bs) = (ln, c:bs') where (ln, bs') = removeNext bs
	loopStart = [Cmd l Let [v, lo], Cmd (l+1) If [Binop v ">" hi, Label n]]
    in  loopStart ++ cs''
removeFor (c:cs) = c : removeFor cs

translateBASIC :: [Expr ()] -> IO (IO ())
translateBASIC cmds = do
    let cmds' = removeFor $ renumber cmds
--    putStrLn $ unlines $ map show cmds'

    let mfunc = trans cmds'
    writeCodeGenModule "run.bc" mfunc
    func <- optimizeFunctionCG mfunc
--    writeCodeGenModule "runo.bc" mfunc'
--    func <- simpleFunction mfunc
    return func


trans :: [Expr ()] -> CodeGenModule (Function (IO ()))
trans acmds = do
    atof     <- newNamedFunction ExternalLinkage "atof"     :: TFunction (Ptr Word8 -> IO Double)
    gets     <- newNamedFunction ExternalLinkage "gets"     :: TFunction (Ptr Word8 -> IO (Ptr Word8))
    power    <- newNamedFunction ExternalLinkage "power"    :: TFunction (Double -> Double -> IO Double)
    printfv  <- newNamedFunction ExternalLinkage "printf"   :: TFunction (Ptr Word8 -> VarArgs Word32)
    rand     <- newNamedFunction ExternalLinkage "rand"     :: TFunction (IO Word32)
    sranddev <- newNamedFunction ExternalLinkage "sranddev" :: TFunction (IO ())
    let printfd :: Function (Ptr Word8 -> Double -> IO Word32)
        printfd = castVarArgs printfv
        printfs :: Function (Ptr Word8 -> Ptr Word8 -> IO Word32)
        printfs = castVarArgs printfv
        printfn :: Function (Ptr Word8 -> IO Word32)
        printfn = castVarArgs printfv

    fmtg <- createStringNul "%g"
    fmts <- createStringNul "%s"
    fmtn <- createStringNul "\n"

    let cmds = acmds ++ [Cmd 99999 End []]
        nextmap = fromList $ zip (map cmdLabel cmds) (map cmdLabel (tail cmds))
        strings = nub $ concatMap getCmdStrings cmds
	getCmdStrings (Cmd _ _ es) = concatMap getExprStrings es
	getCmdStrings _ = error "getCmdStrings"
	getExprStrings (Str s) = [s]
	getExprStrings (Binop e1 _ e2) = getExprStrings e1 ++ getExprStrings e2
	getExprStrings _ = []
    strmap <- liftM (fromList . zip strings) $ mapM createStringNul strings

    let mkGlobal x = do
            v <- createNamedGlobal False InternalLinkage (show x) (constOf (0 :: Double))
	    return (x, v)
    globmap <- liftM fromList $ mapM mkGlobal [I,S,X,Y,Z]

    createFunction ExternalLinkage $ do
        let mkBlk c = do b <- newBasicBlock; return (cmdLabel c, b)
        blkmap <- liftM fromList $ mapM mkBlk cmds
	let block = (blkmap !)
	    next = block . (nextmap !)
	    gen (Cmd l kw es) = do
                defineBasicBlock (block l)
		case (kw, es) of
		    (End, _) -> ret ()
		    (Goto, [Label d]) -> br (block d)
		    (Print, as) -> do mapM_ pr as; newline; br (next l)
		    (Let, [v, e]) -> do
                        d <- genExpr e
                        store d (globmap ! v)
		        br (next l)
		    (If, [b, Label d]) -> do
                        v <- genBool b
                        condBr v (block d) (next l)
                    (Input, [v]) -> do
                        buff <- arrayMalloc (100 :: Word32)
			call gets buff
			d <- call atof buff
			store d (globmap ! v)
			free buff
			br (next l)
		    (Rem, _) -> br (next l)
		    x -> error $ "Unimplemented construct " ++ show x
            gen _ = error "gen"

            newline = do
                tmp <- getElementPtr fmtn (0::Word32, (0::Word32, ()))
                call printfn tmp

            pr (Str s) = do
	        tmp <- getElementPtr fmts (0::Word32, (0::Word32, ()))
	        tmpa <- getElementPtr (strmap ! s) (0::Word32, (0::Word32, ()))
                call printfs tmp tmpa
            pr e = do
                d <- genExpr e
                tmp <- getElementPtr fmtg (0::Word32, (0::Word32, ()))
                call printfd tmp d

--	    genExpr e | trace (show e) False = undefined
            genExpr (Dbl d) = return $ value $ constOf d
	    genExpr (Binop e1 "+" e2) = binop add e1 e2
	    genExpr (Binop e1 "-" e2) = binop sub e1 e2
	    genExpr (Binop e1 "*" e2) = binop mul e1 e2
	    genExpr (Binop e1 "/" e2) = binop fdiv e1 e2
	    genExpr (Binop e1 "^" e2) = binop (call power) e1 e2
	    genExpr (INT e) = do
	        v <- genExpr e
--		r <- frem v (1 :: Double)
--		sub v r
                i <- fptoui v
                uitofp (i :: Value Word64)
            genExpr (RND _) = do
                r <- call rand
                d <- uitofp r
		fdiv (d :: Value Double) (0x7fffffff :: Double)
            genExpr (SGN e) = do
                d <- genExpr e
                n <- fcmp FPOLT d (0 :: Double)
                p <- fcmp FPOGT d (0 :: Double)
		nd <- uitofp n
		pd <- uitofp p
		sub (pd :: Value Double) (nd :: Value Double)
	    genExpr e | e > Var = load (globmap ! e)
	    genExpr e = error $ "genExpr: " ++ show e

            genBool (Binop e1 "<>" e2) = binop (fcmp FPONE) e1 e2
	    genBool (Binop e1 "==" e2) = binop (fcmp FPOEQ) e1 e2
	    genBool (Binop e1 "<"  e2) = binop (fcmp FPOLT) e1 e2
	    genBool (Binop e1 "<=" e2) = binop (fcmp FPOLE) e1 e2
	    genBool (Binop e1 ">"  e2) = binop (fcmp FPOGT) e1 e2
	    genBool (Binop e1 ">=" e2) = binop (fcmp FPOGE) e1 e2
	    genBool e = error $ "Unknown bool op " ++ show e

	    binop :: (Value Double -> Value Double -> CodeGenFunction r (Value c)) ->
	             Expr () -> Expr () -> CodeGenFunction r (Value c)
            binop op e1 e2 = do
                d1 <- genExpr e1
                d2 <- genExpr e2
                op d1 d2

        call sranddev
	br (block $ cmdLabel $ head cmds)
        mapM_ gen cmds