{-# LANGUAGE LambdaCase, RecordWildCards, OverloadedStrings #-}
module Gen where

import Control.Monad.State
import Control.Monad.Except
import Data.Maybe
import Text.Printf
import Parser
import Syntax
import Lambda.Syntax (External(..), SimpleType(..), packName, Name)
import qualified Lambda.Syntax as L
import qualified Lambda.Pretty as L
import Text.PrettyPrint.ANSI.Leijen hiding ((<$>), cat)
--import Numeric

import Data.Char
import Data.Map (Map)
import qualified Data.Map as Map

readInfo :: IO Info
readInfo = do
  s <- readFile "primops.txt"
  case parse s of
     Left err -> error ("parse error at " ++ (show err))
     Right p_o_specs@(Info _ _)
        -> seq (sanityTop p_o_specs) (pure p_o_specs)

deriveNewName :: Name -> G Name
deriveNewName name = do
  i <- Map.findWithDefault 0 name <$> gets envNameCounter
  modify $ \env@Env{..} -> env {envNameCounter = Map.insert name (succ i) envNameCounter}
  {-
  let code    = ['0'..'9'] ++ ['a'..'z'] ++ ['A'..'Z']
      baseX a = showIntAtBase (length code) (code !!) a ""
  pure $ packName $ printf "%s.%s" name (baseX i)
  -}
  pure $ packName $ printf "%s_%d" name i

genName :: Name -> G Name
genName name = do
  m <- gets envNameMap
  case Map.lookup name m of
    Just n  -> pure n
    Nothing -> do
      n <- deriveNewName name
      modify $ \env@Env{..} -> env {envNameMap = Map.insert name n envNameMap}
      pure n

clearTyVarMap :: G ()
clearTyVarMap = modify $ \env -> env {envNameMap = mempty}

cvtExternal :: String -> Ty -> Bool -> ExceptT String G External
cvtExternal n t isEffectful = do
  (retTy, argsTy) <- cvtFunTy $ removeStateFromRetType t
  pure External
    { eName       = packName n
    , eRetType    = retTy
    , eArgsType   = argsTy
    , eEffectful  = isEffectful
    , eKind       = L.PrimOp
    }

mkUnboxedTuple :: [L.Ty] -> G L.Ty
mkUnboxedTuple args = do
  ut <- deriveNewName "t"
  pure $ case length args of
    0 -> L.TyCon ut (packName "ghc-prim_GHC.Prim.(##)") []
    1 -> L.TyCon ut (packName "ghc-prim_GHC.Prim.Solo#") args
    n -> L.TyCon ut (packName $ "ghc-prim_GHC.Prim.(#" ++ replicate (max 0 $ n-1) ',' ++ "#)") args

isStateTy :: Ty -> Bool
isStateTy = \case
  TyApp (TyCon "State#") [_] -> True
  _ -> False

isTyF :: Ty -> Bool
isTyF = \case
  TyF{} -> True
  _     -> False

removeStateFromRetType :: Ty -> Ty
removeStateFromRetType = \case
  TyF a b
    | isTyF b -> TyF (removeStateFromRetType a) $ removeStateFromRetType b
    -- last ty item aka return type
    | otherwise -> TyF (removeStateFromRetType a) $ case b of
        -- remove returning state
        TyUTup args -> TyUTup [removeStateFromRetType t | t <- args, not $ isStateTy t]
        t | isStateTy t -> TyUTup []
          | otherwise   -> t

  t -> t

-- generate unique names for type nodes
cvtFunTy :: Ty -> ExceptT String G (L.Ty, [L.Ty])
cvtFunTy a = do
  lift clearTyVarMap
  ty <- cvtTy a
  case flattenFunTy ty of
    L.TyFun _ r args -> pure (r, args)
    t -> throwError $ "ivalid function type: " ++ show t

flattenFunTy :: L.Ty -> L.Ty
flattenFunTy = \case
  L.TyFun n1 (L.TyFun _n2 r args2) args1 -> flattenFunTy (L.TyFun n1 r (map flattenFunTy $ args1 ++ args2))
  L.TyCon a b args -> L.TyCon a b (map flattenFunTy args)
  t -> t

cvtTy :: Ty -> ExceptT String G L.Ty
cvtTy ty = case ty of
  TyVar n -> L.TyVar <$> lift (genName (packName n))
  TyApp (TyCon n) []
    | Just t <- cvtType n
    -> L.TySimple <$> lift (deriveNewName "t") <*> pure t
  TyApp (TyCon n) args -> do
    xs <- mapM cvtTy args
    tn <- lift (deriveNewName "t")
    pure $ L.TyCon tn (packName n) xs
  TyUTup args -> do
    xs <- mapM cvtTy args
    lift $ mkUnboxedTuple xs
  TyF a b -> do
    a2 <- cvtTy a
    b2 <- cvtTy b
    tn <- lift (deriveNewName "tf")
    pure $ L.TyFun tn b2 [a2]
  t -> throwError $ "unsupported type: " ++ show t

cvtType :: String -> Maybe SimpleType
cvtType = \case
  "Char#"   -> Just T_Char
  "Int#"    -> Just T_Int64
  "Word#"   -> Just T_Word64
  "Double#" -> Just T_Double
  "Float#"  -> Just T_Float
  "Addr#"   -> Just T_Addr
  _ -> Nothing

data Env
  = Env
  { envSections         :: [(String, [External], [(String, String)])] -- section, externals, unsupported ops
  , envDefaults         :: [Option]
  , envSectionTitle     :: String
  , envExternals        :: [External]
  , envUnsupported      :: [(String, String)] -- NOTE: primop name, message
  , envNameMap          :: Map Name Name
  , envNameCounter      :: Map Name Int
  }

emptyEnv :: [Option] -> Env
emptyEnv opts = Env
  { envSections         = []
  , envDefaults         = opts
  , envSectionTitle     = ""
  , envExternals        = []
  , envUnsupported      = []
  , envNameMap          = mempty
  , envNameCounter      = mempty
  }

type G = State Env

attr :: [Option] -> (Option -> Maybe a) -> G a
attr opts f = do
  defaults <- gets envDefaults
  pure $ head $ catMaybes $ map f (opts ++ defaults)

attrBool :: String -> [Option] -> G Bool
attrBool name opts = attr opts $ \case
  OptionFalse n | n == name -> Just False
  OptionTrue  n | n == name -> Just True
  _ -> Nothing

unsupportedPrimop :: String -> String -> G ()
unsupportedPrimop name msg = modify $ \s@Env{..} -> s {envUnsupported = envUnsupported ++ [(name, msg)]}

addExternal :: External -> G ()
addExternal e = modify $ \s@Env{..} -> s {envExternals = envExternals ++ [e]}

flushSection :: G ()
flushSection = do
  exts <- gets envExternals
  errs <- gets envUnsupported
  if null exts && null errs
    then modify $ \s@Env{..} -> s {envExternals = [], envUnsupported = []}
    else modify $ \s@Env{..} -> s {envExternals = [], envUnsupported = [], envSections = envSections ++ [(envSectionTitle, exts, errs)]}

setSection :: String -> G ()
setSection title = modify $ \s@Env{..} -> s {envSectionTitle = title}

visitEntry :: Entry -> G ()
visitEntry = \case
  PrimVecOpSpec{..} -> unsupportedPrimop name "SIMD is not supported yet"
  PseudoOpSpec{..}  -> unsupportedPrimop name "pseudo ops are not supported"

  PrimOpSpec{..} -> do
    has_side_effects <- attrBool "has_side_effects" opts
    err <- runExceptT $ cvtExternal name ty has_side_effects
    case err of
      Right e@External{..}
        -- | not (hasUnknownTyVar eArgsType eRetType) || Set.member name whiteList
        -> addExternal e
        -- | otherwise
        -- -> unsupportedPrimop name "unknown type parameters in the result type"
      Left msg
        -> unsupportedPrimop name msg

  Section{..} -> do
    flushSection
    setSection title

  PrimTypeSpec{}    -> pure ()
  PrimVecTypeSpec{} -> pure ()

processEntries :: [Entry] -> G ()
processEntries ents = do
  mapM_ visitEntry ents
  flushSection

showWidth :: Int -> Doc -> String
showWidth w x = displayS (renderPretty 0.4 w x) ""

genGHCPrimOps :: IO ()
genGHCPrimOps = do
  i <- readInfo
  let Info opts ents = i
      Env{..} = execState (processEntries ents) (emptyEnv opts)
  let header =
        [ "{-# LANGUAGE OverloadedStrings, QuasiQuotes #-}"
        , "module Lambda.Stg.GHCPrimOps where"
        , ""
        , "import qualified Data.Set as Set"
        , "import Lambda.Syntax"
        , "import Lambda.TH"
        , ""
        ]

      tab l
        | all isSpace l = ""
        | otherwise     = "  " ++ l

      comment t =
        [ "{-"
        , "  " ++ t
        , "-}"
        ]

      primPrelude =
        [ "primPrelude :: Program"
        , "primPrelude = [progConst|"
        ] ++ map tab (concat [comment title ++ (lines $ showWidth 800 $ plain $ L.prettyExternals exts) ++ [""] | (title, exts, _) <- envSections]) ++
        ["  |]\n"]

      unsupported =
        [ "unsupported :: Set.Set String"
        , "unsupported = Set.fromList"
        ] ++ ppUnsupported True envSections ++
        ["  ]"]

      ppUnsupported _ [] = []
      ppUnsupported isFirst ((_s,_,[]):xs) = ppUnsupported isFirst xs
      ppUnsupported isFirst ((s,_,l):xs) = ["\n  -- " ++ s] ++ goSection isFirst l ++ ppUnsupported False xs

      goSection _ [] = []
      goSection True ((n,msg):xs)   = ["  [ " ++ (take 40 $ show n ++ repeat ' ') ++ " -- " ++ msg] ++ goSection False xs
      goSection False ((n,msg):xs)  = ["  , " ++ (take 40 $ show n ++ repeat ' ') ++ " -- " ++ msg] ++ goSection False xs

  writeFile "GHCPrimOps.hs" $ unlines $ header ++ primPrelude ++ unsupported

{-
  generate GHC primop prelude module for Lambda
    done - handle attributes
    done - collect:
        done - unsupported primops
        done - supported primops group by sections
    done - generate header
    done - generate ghcUnsupportedPrimopSet
    done - generate ghcPrimopPrelude
    no - transform primop names to   GHC:op_name
    no - transform primtype names to GHC:type_name
    done - filter out void representations from the result ; State#
    done - map unboxed tuples properly
-}
{-
  STG state convention:
    - passes as arg
    - does not return ; filtered out from the result
-}
