module Tessera.Env where

import Data.List (foldl')
import qualified Data.Map.Strict as Map

import Tessera.Syntax

data TypeDef = TypeDef
  { typeDefinitionParams :: [String]
  , typeDefinitionConstructors :: [Constructor]
  } deriving (Eq, Show)

data Env = Env
  { envTypes :: Map.Map String TypeDef
  , envConstructorTypes :: Map.Map String (String, [Type])
  } deriving (Eq, Show)

builtinTypes :: Map.Map String TypeDef
builtinTypes = Map.fromList
  [ ("Int", TypeDef [] [])
  , ("Bool", TypeDef [] [Constructor "True" [], Constructor "False" []])
  ]

builtinConstructors :: Map.Map String (String, [Type])
builtinConstructors = Map.fromList
  [ ("True", ("Bool", []))
  , ("False", ("Bool", []))
  ]

buildEnv :: Program -> Either Diagnostic Env
buildEnv program = do
  types <- foldl' addType (Right builtinTypes) [d | d <- program, isData d]
  let constructors = foldl' addConstructors builtinConstructors [d | d <- program, isData d]
  return (Env types constructors)
  where
    isData (DataDeclaration _ _ _) = True
    isData _ = False

    addType acc (DataDeclaration name params constructors)
      | name `elem` ["Int", "Bool"] = Left (Diagnostic (Span 0 0) ("cannot redeclare builtin type " ++ name))
      | otherwise = do
          table <- acc
          case Map.lookup name table of
            Just _ -> Left (Diagnostic (Span 0 0) ("duplicate type declaration " ++ name))
            Nothing -> Right (Map.insert name (TypeDef params constructors) table)
    addType acc _ = acc

    addConstructors table (DataDeclaration name _ constructors) =
      foldl' (\m c -> Map.insert (constructorName c) (name, constructorArgs c) m) table constructors
    addConstructors table _ = table

lookupConstructor :: Env -> String -> Maybe (String, [Type])
lookupConstructor env name = Map.lookup name (envConstructorTypes env)

lookupType :: Env -> String -> Maybe TypeDef
lookupType env name = Map.lookup name (envTypes env)

constructorArity :: Env -> String -> Maybe Int
constructorArity env name = length . snd <$> lookupConstructor env name

substitute :: Map.Map String Type -> Type -> Type
substitute mapping ty = case ty of
  TVar name -> Map.findWithDefault (TVar name) name mapping
  TName name arguments -> TName name (map (substitute mapping) arguments)
  TTuple elements -> TTuple (map (substitute mapping) elements)

instantiate :: [String] -> [Type] -> [Type] -> [Type]
instantiate params arguments = map (substitute mapping)
  where mapping = Map.fromList (zip params arguments)

typeEquals :: Type -> Type -> Bool
typeEquals left right = case (left, right) of
  (TName a as, TName b bs) -> a == b && length as == length bs && and (zipWith typeEquals as bs)
  (TVar a, TVar b) -> a == b
  (TTuple as, TTuple bs) -> length as == length bs && and (zipWith typeEquals as bs)
  _ -> False