module Core.State
  ( DsM
  , DsEnum
  , DsCtor
  , DsState
  , initialState
  , runDesugar
  , var
  , contVar
  , addEnum
  , findEnumForCtor
  , arity
  , enumCtors
  , ctorName
  , importDsState
  ) where

import qualified Core.Absyn as CA

import Absyn.Typed
import Lib.Registry
import Typing.Types

import Control.Monad.State (State, runState, gets, modify)
import Data.Tuple (swap)

data DsEnum = DsEnum { _enumName :: String, enumCtors :: [DsCtor] }
data DsCtor = DsCtor { ctorName :: String, arity :: Int }
data DsState = DsState { contCount :: Int, varCount :: Int, enums :: [DsEnum] }

runDesugar :: (([CA.Var] -> DsM CA.Term) -> DsM CA.Term) -> DsState -> (DsState, CA.Term)
runDesugar d =
  let k z = return (CA.AppCont (CA.ContVar "halt") z)
   in swap . runState (d k)

initialState :: DsState
initialState = DsState
  { contCount = 0
  , varCount = 0
  , enums
  }
    where
      enums = map (uncurry DsEnum) $ foldl aux [] ctors

      ctors = filter isCtor registry

      aux enums ctor =
        let (name, ty) = decl ctor
            (enum, ctor') = mkDsCtor name ty
         in addTo enum ctor' enums

      mkDsCtor name (Fun _ args ty) =
        (getEnumForTy ty, DsCtor name (length args))
      mkDsCtor name ty = (getEnumForTy ty, DsCtor name 0)

      getEnumForTy (Con enum) = enum
      getEnumForTy (Forall _ ty) = getEnumForTy ty
      getEnumForTy (TyApp ty _) = getEnumForTy ty
      getEnumForTy _ = undefined

      addTo enum ctor [] = [(enum, [ctor])]
      addTo enum ctor ((enum', ctors):rest) | enum == enum' =
        (enum, ctor : ctors) : rest
      addTo enum ctor (x:xs) = x : addTo enum ctor xs



type DsM = State DsState

contVar :: DsM CA.ContVar
contVar = do
  count <- gets contCount
  modify $ \s -> s { contCount = count + 1 }
  return $ CA.ContVar ("#k" ++ show count)

var :: DsM CA.Var
var = do
  count <- gets varCount
  modify $ \s -> s { varCount = count + 1 }
  return $ CA.Var ("#x" ++ show count)

addEnum :: ASTNode BaseDecl Id c -> DsM ()
addEnum (Enum name _ ctors) =
  modify $ \s ->
    s { enums = DsEnum name ctors' : (enums s) }
  where
    ctors' = map mkCtor ctors
    mkCtor (name, args) = DsCtor { ctorName = name, arity = maybe 0 length args }

addEnum _ = error "addEnum called with a Decl that is not an enum"

findEnumForCtor :: String -> DsM DsEnum
findEnumForCtor c = do
  es <- gets enums
  return $ f es
    where
      f :: [DsEnum] -> DsEnum
      f [] = undefined
      f (e:es) =
        if g (enumCtors e) then e else f es
      g :: [DsCtor] -> Bool
      g [] = False
      g (c' : cs)
        | ctorName c' == c = True
        | otherwise = g cs

importDsState :: [String] -> DsState -> DsState -> DsState
importDsState imports targetState importState =
  importState { enums = mergedEnums }
    where
      mergedEnums = importedEnums ++ enums targetState
      importedEnums = filter pred (enums importState)
      pred (DsEnum name _) = name `elem` imports
