{-# LANGUAGE OverloadedStrings #-}

module Aihc.Tc.Kind
  ( TvKindEnv,
    ParamInfo (..),
    checkSurfaceType,
    checkRuntimeType,
    convertSurfaceType,
    convertSurfaceTypeWithKinds,
    defaultKindMetas,
    freeTypeVars,
    freshKindMeta,
    kindToTcType,
    classPredicateArgKinds,
    makeParamEnv,
    makeParamEnvWith,
    runtimeRepToTcType,
    sigToScheme,
    standaloneKindSigToScheme,
    surfacePredToPred,
    tyConKindFromParams,
    tyConKindFromParamsWith,
    unifyKinds,
    zonkKind,
  )
where

import Aihc.Parser.Syntax
  ( Name (..),
    SourceSpan (..),
    TupleFlavor (..),
    TyVarBinder (..),
    Type (..),
    TypeBuiltinCon (..),
    TypePromotion (..),
    UnqualifiedName (..),
    forallTelescopeBinders,
    instanceHeadName,
    instanceHeadTypes,
    nameText,
    peelTypeHead,
    tyVarBinderKind,
    tyVarBinderName,
    unqualifiedNameText,
  )
import Aihc.Tc.Env (TyConInfo (..), TypeSynonymInfo (..))
import Aihc.Tc.Error (TcErrorKind (..))
import Aihc.Tc.Instantiate (applySubst, instantiate)
import Aihc.Tc.Monad
import Aihc.Tc.Types
import Control.Monad (unless, zipWithM)
import Data.List (nub)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as T

type TvKindEnv = Map Text (TyVarId, Kind)

data ParamInfo = ParamInfo
  { paramName :: !Text,
    paramTyVar :: !TyVarId,
    paramKind :: !Kind
  }
  deriving (Show)

sigToScheme :: Type -> TcM TypeScheme
sigToScheme ty = do
  let (explicitBinders, qualifiedBody) = splitForalls ty
      (context, body) = splitContext qualifiedBody
      freeVars = freeTypeVars ty
  rawTvs <- mapM freshSkolemTv freeVars
  kinds <- mapM (const freshKindMeta) freeVars
  let implicitTvs = zipWith setTyVarKind kinds rawTvs
  let implicitEnv = Map.fromList (zip freeVars (zip implicitTvs kinds))
  explicitParams <- makeParamEnvWith implicitEnv explicitBinders
  let explicitTvs = map paramTyVar explicitParams
      tvEnv =
        implicitEnv
          <> Map.fromList
            [ (paramName param, (paramTyVar param, paramKind param))
            | param <- explicitParams
            ]
  tcTy <- checkRuntimeType tvEnv body
  preds <- mapM (surfacePredToPred tvEnv) context
  pure (ForAll (implicitTvs <> explicitTvs) preds tcTy)

standaloneKindSigToScheme :: Type -> TcM TypeScheme
standaloneKindSigToScheme ty = do
  let (explicitBinders, bodyType) = splitForalls ty
      freeVars = freeTypeVars ty
      representationVars = Set.fromList (runtimeRepVariables ty)
      variableKind name
        | name `Set.member` representationVars = KRuntimeRep
        | otherwise = KType
  rawTyVars <- mapM freshSkolemTv freeVars
  let implicitTyVars = zipWith (setTyVarKind . variableKind) freeVars rawTyVars
      implicitEnv = Map.fromList [(tvName tyVar, (tyVar, tvKind tyVar)) | tyVar <- implicitTyVars]
  explicitParams <- makeParamEnvWith implicitEnv explicitBinders
  let explicitTyVars = map paramTyVar explicitParams
      tyVarEnv =
        implicitEnv
          <> Map.fromList
            [ (paramName param, (paramTyVar param, paramKind param))
            | param <- explicitParams
            ]
  body <- checkSurfaceType tyVarEnv bodyType KType
  pure (ForAll (implicitTyVars <> explicitTyVars) [] body)

runtimeRepVariables :: Type -> [Text]
runtimeRepVariables = nub . go
  where
    go surfaceType =
      case peelTypeHead surfaceType of
        TApp function runtimeRep
          | TCon name Unpromoted <- peelTypeHead function,
            nameText name == "TYPE" ->
              freeTypeVars runtimeRep
        TApp function argument -> go function <> go argument
        TTypeApp function argument -> go function <> go argument
        TInfix left _ _ right -> go left <> go right
        TFun _ argument result -> go argument <> go result
        TTuple _ _ arguments -> concatMap go arguments
        TUnboxedSum arguments -> concatMap go arguments
        TList _ arguments -> concatMap go arguments
        TKindSig inner kind -> go inner <> go kind
        TContext predicates body -> concatMap go predicates <> go body
        TForall telescope body -> concatMap binderKindVariables (forallTelescopeBinders telescope) <> go body
        _ -> []
    binderKindVariables binder = maybe [] go (tyVarBinderKind binder)

convertSurfaceType :: Map Text TyVarId -> Type -> TcM TcType
convertSurfaceType tvMap ty = do
  let tvEnv = Map.map (\tv -> (tv, tvKind tv)) tvMap
  checkRuntimeType tvEnv ty

checkSurfaceType :: TvKindEnv -> Type -> Kind -> TcM TcType
checkSurfaceType tvEnv ty expected = do
  (tcTy, actual) <- convertSurfaceTypeWithKinds tvEnv ty
  unifyKinds expected actual
  pure tcTy

-- | Check that a surface type is a value-bearing type of kind @TYPE rep@.
-- Unconstrained kind metas default to lifted representation; explicitly
-- unlifted types retain their fixed representation.
checkRuntimeType :: TvKindEnv -> Type -> TcM TcType
checkRuntimeType tvEnv ty = do
  (tcTy, actual) <- convertSurfaceTypeWithKinds tvEnv ty
  actual' <- zonkKind actual
  case actual' of
    KTYPE {} -> pure tcTy
    KMeta unique -> bindKindMeta unique KType >> pure tcTy
    _ -> emitError NoSourceSpan (KindMismatch KType actual') >> pure tcTy

convertSurfaceTypeWithKinds :: TvKindEnv -> Type -> TcM (TcType, Kind)
convertSurfaceTypeWithKinds tvEnv ty = do
  expanded <- expandTypeSynonym tvEnv (peelTypeHead ty)
  case expanded of
    Just result -> pure result
    Nothing -> convertNonSynonymTypeWithKinds tvEnv (peelTypeHead ty)

convertNonSynonymTypeWithKinds :: TvKindEnv -> Type -> TcM (TcType, Kind)
convertNonSynonymTypeWithKinds tvEnv ty =
  case ty of
    TAnn _ inner ->
      convertSurfaceTypeWithKinds tvEnv inner
    TVar name ->
      inferTypeVariable tvEnv name
    TCon name promoted ->
      inferTypeConstructor promoted name
    TBuiltinCon builtin ->
      inferBuiltinTypeConstructor builtin
    TStar {} ->
      knownType "GHC.Types" "Type" KType
    -- GHC's MutVar# has an inferred RuntimeRep parameter, so its value
    -- argument may have any TYPE r kind even though r is not a visible type
    -- constructor argument.
    TApp function value
      | (TCon name Unpromoted, [state]) <- typeApplicationSpine function,
        nameText name == "MutVar#" -> do
          (mutVarType, _) <- inferTypeConstructor Unpromoted name
          stateType <- checkSurfaceType tvEnv state KType
          valueType <- checkRuntimeType tvEnv value
          pure
            ( applyType (applyType mutVarType stateType) valueType,
              KTYPE (BoxedRep Unlifted)
            )
    TApp f a -> do
      (fTy, fKind) <- convertSurfaceTypeWithKinds tvEnv f
      (aTy, aKind) <- convertSurfaceTypeWithKinds tvEnv a
      resultKind <- freshKindMeta
      unifyKinds fKind (KFun aKind resultKind)
      resultKind' <- zonkKind resultKind
      pure (applyType fTy aTy, resultKind')
    TTypeApp f a -> do
      (fTy, fKind) <- convertSurfaceTypeWithKinds tvEnv f
      (aTy, aKind) <- convertSurfaceTypeWithKinds tvEnv a
      resultKind <- freshKindMeta
      unifyKinds fKind (KFun aKind resultKind)
      resultKind' <- zonkKind resultKind
      pure (applyType fTy aTy, resultKind')
    TInfix lhs name promoted rhs ->
      convertSurfaceTypeWithKinds tvEnv (TApp (TApp (TCon name promoted) lhs) rhs)
    TFun _ a b -> do
      aTy <- checkRuntimeType tvEnv a
      bTy <- checkRuntimeType tvEnv b
      pure (TcFunTy aTy bTy, KType)
    TTuple flavor _ args -> do
      tys <-
        case flavor of
          Boxed -> mapM (\arg -> checkSurfaceType tvEnv arg KType) args
          Unboxed -> mapM (checkRuntimeType tvEnv) args
      let arity = length tys
          fallbackResultKind =
            case flavor of
              Boxed -> KType
              Unboxed -> KTYPE (TupleRep (map runtimeRepOrLifted tys))
          fallbackKind = foldr (KFun . typeKind) fallbackResultKind tys
          typeName = tupleTyConText flavor arity
      maybeTyCon <- lookupTyCon typeName
      tyCon <-
        case maybeTyCon of
          Just info -> pure (tciTyCon info)
          Nothing -> mkKnownTyCon (tupleTyConModule flavor) typeName arity fallbackKind
      let tupleType = TcTyCon tyCon tys
      pure (tupleType, typeKind tupleType)
    TUnboxedSum args -> do
      tys <- mapM (checkRuntimeType tvEnv) args
      let arity = length tys
          resultKind = KTYPE (SumRep (map runtimeRepOrLifted tys))
          tyConKind' = foldr (KFun . typeKind) resultKind tys
      pure (TcTyCon (mkTyCon ("(#" <> bars (arity - 1) <> "#)") arity tyConKind') tys, resultKind)
    TList Unpromoted [arg] -> do
      argTy <- checkSurfaceType tvEnv arg KType
      listTy <- listType argTy
      pure (listTy, KType)
    TList Promoted args -> do
      elemKind <- freshKindMeta
      args' <- mapM (\arg -> checkSurfaceType tvEnv arg elemKind) args
      pure (TcTyCon (TyCon "'[]" (length args')) args', KType)
    TKindSig inner kindTy -> do
      expected <- kindFromSurfaceType tvEnv kindTy
      checkSurfaceType tvEnv inner expected >>= \innerTy -> pure (innerTy, expected)
    TContext preds inner -> do
      mapM_ (\predTy -> checkSurfaceType tvEnv predTy KConstraint) preds
      convertSurfaceTypeWithKinds tvEnv inner
    TForall telescope inner -> do
      params <- makeParamEnv (forallTelescopeBinders telescope)
      let tvEnv' = tvEnv <> Map.fromList [(paramName p, (paramTyVar p, paramKind p)) | p <- params]
      (innerTy, innerKind) <- convertSurfaceTypeWithKinds tvEnv' inner
      pure (foldr (TcForAllTy . paramTyVar) innerTy params, innerKind)
    _ -> do
      emitError NoSourceSpan (OtherError ("unsupported surface type in kind checker: " <> take 80 (show ty)))
      meta <- freshMetaTv
      pure (meta, KType)

expandTypeSynonym :: TvKindEnv -> Type -> TcM (Maybe (TcType, Kind))
expandTypeSynonym tvEnv ty =
  case typeApplicationSpine ty of
    (TCon name Unpromoted, arguments) -> do
      maybeInfo <- lookupResolvedTyCon name
      case maybeInfo >>= tciTypeSynonym of
        Just synonym
          | Just {} <- tsiBody synonym -> Just <$> instantiateTypeSynonym tvEnv (nameText name) synonym arguments
        _ -> pure Nothing
    _ -> pure Nothing

instantiateTypeSynonym :: TvKindEnv -> Text -> TypeSynonymInfo -> [Type] -> TcM (TcType, Kind)
instantiateTypeSynonym tvEnv synonymName synonym arguments =
  case tsiBody synonym of
    Nothing -> do
      emitError NoSourceSpan (OtherError ("recursive or incomplete type synonym: " <> T.unpack synonymName))
      meta <- freshMetaTv
      pure (meta, KType)
    Just body -> do
      let params = tsiParams synonym
          arity = length params
          (synonymArguments, remainingArguments) = splitAt arity arguments
      if length synonymArguments /= arity
        then do
          emitError NoSourceSpan (OtherError ("type synonym " <> T.unpack synonymName <> " is not fully applied"))
          meta <- freshMetaTv
          pure (meta, KType)
        else do
          checkedArguments <- zipWithM checkArgument params synonymArguments
          let substitution = Map.fromList (zip (map tvUnique params) checkedArguments)
          expandedBody <- expandTcTypeSynonyms Set.empty (applySubst substitution body)
          applyRemainingArguments (expandedBody, typeKind expandedBody) remainingArguments
  where
    checkArgument param argument = checkSurfaceType tvEnv argument (tvKind param)

    applyRemainingArguments result [] = pure result
    applyRemainingArguments (functionType, functionKind) (argument : rest) = do
      (argumentType, argumentKind) <- convertSurfaceTypeWithKinds tvEnv argument
      resultKind <- freshKindMeta
      unifyKinds functionKind (KFun argumentKind resultKind)
      zonkedResultKind <- zonkKind resultKind
      applyRemainingArguments (applyType functionType argumentType, zonkedResultKind) rest

typeApplicationSpine :: Type -> (Type, [Type])
typeApplicationSpine = go []
  where
    go arguments (TAnn _ inner) = go arguments inner
    go arguments (TApp function argument) = go (argument : arguments) function
    go arguments (TTypeApp function argument) = go (argument : arguments) function
    go arguments headType = (headType, arguments)

expandTcTypeSynonyms :: Set Text -> TcType -> TcM TcType
expandTcTypeSynonyms expanding ty =
  case ty of
    TcTyVar {} -> pure ty
    TcMetaTv {} -> pure ty
    TcTyCon tyCon arguments -> do
      expandedArguments <- mapM (expandTcTypeSynonyms expanding) arguments
      maybeInfo <- lookupTyConByIdentity tyCon
      case maybeInfo >>= tciTypeSynonym of
        Just synonym
          | Just body <- tsiBody synonym,
            let params = tsiParams synonym,
            length expandedArguments >= length params ->
              if tyConName tyCon `Set.member` expanding
                then do
                  emitError NoSourceSpan (OtherError ("recursive type synonym: " <> T.unpack (tyConName tyCon)))
                  pure (TcTyCon tyCon expandedArguments)
                else do
                  let (synonymArguments, remainingArguments) = splitAt (length params) expandedArguments
                      substitution = Map.fromList (zip (map tvUnique params) synonymArguments)
                      expandedBody = applySubst substitution body
                  normalizedBody <- expandTcTypeSynonyms (Set.insert (tyConName tyCon) expanding) expandedBody
                  expandTcTypeSynonyms expanding (foldl applyType normalizedBody remainingArguments)
        _ -> pure (TcTyCon tyCon expandedArguments)
    TcFunTy argument result -> TcFunTy <$> expandTcTypeSynonyms expanding argument <*> expandTcTypeSynonyms expanding result
    TcForAllTy tyVar body -> TcForAllTy tyVar <$> expandTcTypeSynonyms expanding body
    TcQualTy predicates body -> TcQualTy <$> mapM expandPredicate predicates <*> expandTcTypeSynonyms expanding body
    TcAppTy function argument -> applyType <$> expandTcTypeSynonyms expanding function <*> expandTcTypeSynonyms expanding argument
    TcBuiltinTyCon name arity arguments ->
      TcBuiltinTyCon name arity <$> mapM (expandTcTypeSynonyms expanding) arguments
  where
    expandPredicate predicate =
      case predicate of
        ClassPred className arguments -> ClassPred className <$> mapM (expandTcTypeSynonyms expanding) arguments
        EqPred left right -> EqPred <$> expandTcTypeSynonyms expanding left <*> expandTcTypeSynonyms expanding right

inferTypeVariable :: TvKindEnv -> UnqualifiedName -> TcM (TcType, Kind)
inferTypeVariable tvEnv name =
  let n = unqualifiedNameText name
   in case Map.lookup n tvEnv of
        Just (tv, kind) -> pure (TcTyVar tv, kind)
        Nothing -> inferOpenTypeConstructor n

inferTypeConstructor :: TypePromotion -> Name -> TcM (TcType, Kind)
inferTypeConstructor promoted name =
  case promoted of
    Promoted -> inferPromotedTypeConstructor (nameText name)
    Unpromoted ->
      case nameText name of
        "Type" -> knownType "GHC.Types" "Type" KType
        "Constraint" -> knownType "GHC.Types" "Constraint" KType
        raw -> do
          mInfo <- lookupResolvedTyCon name
          case mInfo of
            Just info -> do
              kind <- instantiateTyConKind info
              pure (TcTyCon (tciTyCon info) [], kind)
            Nothing -> inferBuiltinOrOpenTypeConstructor raw

instantiateTyConKind :: TyConInfo -> TcM Kind
instantiateTyConKind info = do
  (kindType, _) <- instantiate (tyConKindScheme (tciTyCon info))
  pure (kindFromTypeScheme (ForAll [] [] kindType))

inferBuiltinTypeConstructor :: TypeBuiltinCon -> TcM (TcType, Kind)
inferBuiltinTypeConstructor builtin =
  case builtin of
    TBuiltinList ->
      do
        maybeInfo <- lookupTyCon "[]"
        pure (TcTyCon (maybe (mkTyCon "[]" 1 (KFun KType KType)) tciTyCon maybeInfo) [], KFun KType KType)
    TBuiltinCons ->
      pure (TcTyCon (TyCon ":" 2) [], KFun KType (KFun (listTypeKind KType) (listTypeKind KType)))
    TBuiltinTuple arity ->
      let argKinds = replicate arity KType
          kind = foldr KFun KType argKinds
       in knownTypeWithArity "GHC.Tuple" (tupleTyConText Boxed arity) arity kind
    TBuiltinArrow ->
      pure (TcTyCon (TyCon "(->)" 2) [], KFun KType (KFun KType KType))

inferBuiltinOrOpenTypeConstructor :: Text -> TcM (TcType, Kind)
inferBuiltinOrOpenTypeConstructor name =
  case wiredInTypeKind name of
    Just kind -> knownType (knownTypeModule name) name kind
    Nothing -> inferOpenTypeConstructor name

knownType :: Text -> Text -> Kind -> TcM (TcType, Kind)
knownType moduleName name = knownTypeWithArity moduleName name 0

knownTypeWithArity :: Text -> Text -> Int -> Kind -> TcM (TcType, Kind)
knownTypeWithArity moduleName name arity kind = do
  maybeInfo <- lookupTyCon name
  tyCon <- maybe (mkKnownTyCon moduleName name arity kind) (pure . tciTyCon) maybeInfo
  pure (TcTyCon tyCon [], kind)

knownTypeModule :: Text -> Text
knownTypeModule name
  | "#" `T.isSuffixOf` name = "GHC.Prim"
  | otherwise = "GHC.Types"

inferPromotedTypeConstructor :: Text -> TcM (TcType, Kind)
inferPromotedTypeConstructor name =
  case runtimeRepConstructor name of
    Just _ -> pure (TcTyCon (mkTyCon ("'" <> name) 0 KRuntimeRep) [], KRuntimeRep)
    Nothing -> inferOpenTypeConstructor ("'" <> name)

inferOpenTypeConstructor :: Text -> TcM (TcType, Kind)
inferOpenTypeConstructor name = do
  kind <- freshKindMeta
  pure (TcTyCon (mkTyCon name 0 kind) [], kind)

makeParamEnv :: [TyVarBinder] -> TcM [ParamInfo]
makeParamEnv = makeParamEnvWith Map.empty

makeParamEnvWith :: TvKindEnv -> [TyVarBinder] -> TcM [ParamInfo]
makeParamEnvWith = go
  where
    go _ [] = pure []
    go tvEnv (binder : rest) = do
      rawTv <- freshSkolemTv (tyVarBinderName binder)
      kind <- maybe freshKindMeta (kindFromSurfaceType tvEnv) (tyVarBinderKind binder)
      let tv = setTyVarKind kind rawTv
          param =
            ParamInfo
              { paramName = tyVarBinderName binder,
                paramTyVar = tv,
                paramKind = kind
              }
          tvEnv' = Map.insert (paramName param) (tv, kind) tvEnv
      (param :) <$> go tvEnv' rest

tyConKindFromParams :: [ParamInfo] -> Maybe Type -> TcM Kind
tyConKindFromParams = tyConKindFromParamsWith Map.empty

tyConKindFromParamsWith :: TvKindEnv -> [ParamInfo] -> Maybe Type -> TcM Kind
tyConKindFromParamsWith outerEnv params maybeResultKind = do
  let tvEnv = Map.fromList [(paramName param, (paramTyVar param, paramKind param)) | param <- params] <> outerEnv
  resultKind <- maybe (pure KType) (kindFromSurfaceType tvEnv) maybeResultKind
  pure (foldr (KFun . paramKind) resultKind params)

kindFromSurfaceType :: TvKindEnv -> Type -> TcM Kind
kindFromSurfaceType tvEnv ty =
  case peelTypeHead ty of
    TStar {} -> pure KType
    TCon name Unpromoted
      | nameText name == "Type" -> pure KType
      | nameText name == "Constraint" -> pure KConstraint
      | nameText name == "RuntimeRep" -> pure KRuntimeRep
      | nameText name == "Levity" -> pure KLevity
      | nameText name == "VecCount" -> pure KVecCount
      | nameText name == "VecElem" -> pure KVecElem
    TApp function repTy
      | TCon name Unpromoted <- peelTypeHead function,
        nameText name == "TYPE" ->
          KTYPE <$> runtimeRepFromSurfaceType tvEnv repTy
    TFun _ a b -> KFun <$> kindFromSurfaceType tvEnv a <*> kindFromSurfaceType tvEnv b
    TParen inner -> kindFromSurfaceType tvEnv inner
    TAnn _ inner -> kindFromSurfaceType tvEnv inner
    other -> do
      (_tcTy, kind) <- convertSurfaceTypeWithKinds tvEnv other
      unifyKinds kind KType
      pure KType

unifyKinds :: Kind -> Kind -> TcM ()
unifyKinds expected actual = do
  expected' <- zonkKind expected
  actual' <- zonkKind actual
  case (expected', actual') of
    (KMeta u, kind) -> bindKindMeta u kind
    (kind, KMeta u) -> bindKindMeta u kind
    (KTYPE expectedRep, KTYPE actualRep) -> do
      unified <- unifyRuntimeReps expectedRep actualRep
      unless unified (emitError NoSourceSpan (KindMismatch expected' actual'))
    (KConstraint, KConstraint) -> pure ()
    (KRuntimeRep, KRuntimeRep) -> pure ()
    (KLevity, KLevity) -> pure ()
    (KVecCount, KVecCount) -> pure ()
    (KVecElem, KVecElem) -> pure ()
    (KFun a1 b1, KFun a2 b2) -> unifyKinds a1 a2 >> unifyKinds b1 b2
    _ -> emitError NoSourceSpan (KindMismatch expected' actual')

unifyRuntimeReps :: RuntimeRep -> RuntimeRep -> TcM Bool
unifyRuntimeReps expected actual =
  case (expected, actual) of
    (RuntimeRepMeta left, RuntimeRepMeta right)
      | left == right -> pure True
    (RuntimeRepMeta unique, runtimeRep) ->
      writeMetaTv unique (runtimeRepToTcType runtimeRep) >> pure True
    (runtimeRep, RuntimeRepMeta unique) ->
      writeMetaTv unique (runtimeRepToTcType runtimeRep) >> pure True
    (TupleRep expectedFields, TupleRep actualFields)
      | length expectedFields == length actualFields -> and <$> zipWithM unifyRuntimeReps expectedFields actualFields
    (SumRep expectedFields, SumRep actualFields)
      | length expectedFields == length actualFields -> and <$> zipWithM unifyRuntimeReps expectedFields actualFields
    _ -> pure (expected == actual)

bindKindMeta :: Unique -> Kind -> TcM ()
bindKindMeta u kind
  | kind == KMeta u = pure ()
  | occursInKind u kind = emitError NoSourceSpan (KindMismatch (KMeta u) kind)
  | otherwise = writeKindMeta u kind

zonkKind :: Kind -> TcM Kind
zonkKind kind =
  case kind of
    KMeta u -> do
      mKind <- readKindMeta u
      case mKind of
        Nothing -> pure kind
        Just solved -> zonkKind solved
    KFun a b -> KFun <$> zonkKind a <*> zonkKind b
    KTYPE runtimeRep -> KTYPE <$> zonkRuntimeRep runtimeRep
    KConstraint -> pure KConstraint
    KRuntimeRep -> pure KRuntimeRep
    KLevity -> pure KLevity
    KVecCount -> pure KVecCount
    KVecElem -> pure KVecElem

zonkRuntimeRep :: RuntimeRep -> TcM RuntimeRep
zonkRuntimeRep runtimeRep =
  case runtimeRep of
    RuntimeRepMeta unique -> do
      solution <- readMetaTv unique
      pure (maybe runtimeRep runtimeRepFromType solution)
    TupleRep fields -> TupleRep <$> mapM zonkRuntimeRep fields
    SumRep fields -> SumRep <$> mapM zonkRuntimeRep fields
    _ -> pure runtimeRep

defaultKindMetas :: Kind -> TcM Kind
defaultKindMetas kind =
  case kind of
    KMeta u -> do
      mKind <- readKindMeta u
      case mKind of
        Nothing -> writeKindMeta u KType >> pure KType
        Just solved -> defaultKindMetas solved
    KFun a b -> KFun <$> defaultKindMetas a <*> defaultKindMetas b
    KTYPE runtimeRep -> pure (KTYPE runtimeRep)
    KConstraint -> pure KConstraint
    KRuntimeRep -> pure KRuntimeRep
    KLevity -> pure KLevity
    KVecCount -> pure KVecCount
    KVecElem -> pure KVecElem

freshKindMeta :: TcM Kind
freshKindMeta = KMeta <$> freshUnique

occursInKind :: Unique -> Kind -> Bool
occursInKind needle kind =
  case kind of
    KMeta u -> u == needle
    KFun a b -> occursInKind needle a || occursInKind needle b
    KTYPE runtimeRep -> occursInRuntimeRep needle runtimeRep
    KConstraint -> False
    KRuntimeRep -> False
    KLevity -> False
    KVecCount -> False
    KVecElem -> False

applyType :: TcType -> TcType -> TcType
applyType (TcTyCon tc args) arg = TcTyCon tc (args ++ [arg])
applyType f arg = TcAppTy f arg

listType :: TcType -> TcM TcType
listType ty = do
  maybeInfo <- lookupTyCon "[]"
  pure (TcTyCon (maybe (mkTyCon "[]" 1 (KFun KType KType)) tciTyCon maybeInfo) [ty])

listTypeKind :: Kind -> Kind
listTypeKind kind = KFun kind kind

runtimeRepOrLifted :: TcType -> RuntimeRep
runtimeRepOrLifted ty =
  case runtimeRepOfType ty of
    Right runtimeRep -> runtimeRep
    Left _ -> liftedRuntimeRep

wiredInTypeKind :: Text -> Maybe Kind
wiredInTypeKind name =
  case name of
    "Int" -> Just KType
    "Integer" -> Just KType
    "Double" -> Just KType
    "Float" -> Just KType
    "Char" -> Just KType
    "Bool" -> Just KType
    "Array#" -> Just (KFun KType (KTYPE (BoxedRep Unlifted)))
    "ByteArray#" -> Just (KTYPE (BoxedRep Unlifted))
    "MutableArray#" -> Just (KFun KType (KFun KType (KTYPE (BoxedRep Unlifted))))
    "MutableByteArray#" -> Just (KFun KType (KTYPE (BoxedRep Unlifted)))
    "Int#" -> Just (KTYPE IntRep)
    "Int8#" -> Just (KTYPE Int8Rep)
    "Int16#" -> Just (KTYPE Int16Rep)
    "Int32#" -> Just (KTYPE Int32Rep)
    "Int64#" -> Just (KTYPE Int64Rep)
    "Word#" -> Just (KTYPE WordRep)
    "Word8#" -> Just (KTYPE Word8Rep)
    "Word16#" -> Just (KTYPE Word16Rep)
    "Word32#" -> Just (KTYPE Word32Rep)
    "Word64#" -> Just (KTYPE Word64Rep)
    "Addr#" -> Just (KTYPE AddrRep)
    "Float#" -> Just (KTYPE FloatRep)
    "Double#" -> Just (KTYPE DoubleRep)
    "Char#" -> Just (KTYPE WordRep)
    _ -> Nothing

runtimeRepFromSurfaceType :: TvKindEnv -> Type -> TcM RuntimeRep
runtimeRepFromSurfaceType tvEnv ty =
  case ty of
    TAnn _ inner -> runtimeRepFromSurfaceType tvEnv inner
    TParen inner -> runtimeRepFromSurfaceType tvEnv inner
    TVar name ->
      case Map.lookup (unqualifiedNameText name) tvEnv of
        Just (tyVar, KRuntimeRep) -> pure (RuntimeRepVar (tvUnique tyVar))
        _ -> invalidRuntimeRep
    TCon name _ ->
      maybe invalidRuntimeRep pure (runtimeRepConstructor (nameText name))
    TApp function levityTy
      | TCon name _ <- peelTypeHead function,
        nameText name == "BoxedRep" ->
          BoxedRep <$> levityFromSurfaceType levityTy
    TApp function fieldsTy
      | TCon name _ <- peelTypeHead function,
        T.dropWhile (== '\'') (nameText name) == "TupleRep" ->
          TupleRep <$> runtimeRepListFromSurfaceType tvEnv fieldsTy
    _ -> invalidRuntimeRep
  where
    invalidRuntimeRep = do
      emitError NoSourceSpan (OtherError ("invalid RuntimeRep: " <> take 80 (show ty)))
      pure liftedRuntimeRep

runtimeRepListFromSurfaceType :: TvKindEnv -> Type -> TcM [RuntimeRep]
runtimeRepListFromSurfaceType tvEnv ty =
  case ty of
    TAnn _ inner -> runtimeRepListFromSurfaceType tvEnv inner
    TParen inner -> runtimeRepListFromSurfaceType tvEnv inner
    TList Promoted fields -> mapM (runtimeRepFromSurfaceType tvEnv) fields
    _ -> do
      emitError NoSourceSpan (OtherError ("invalid RuntimeRep list: " <> take 80 (show ty)))
      pure []

runtimeRepConstructor :: Text -> Maybe RuntimeRep
runtimeRepConstructor rawName =
  lookup
    (T.dropWhile (== '\'') rawName)
    [ ("LiftedRep", liftedRuntimeRep),
      ("UnliftedRep", BoxedRep Unlifted),
      ("IntRep", IntRep),
      ("Int8Rep", Int8Rep),
      ("Int16Rep", Int16Rep),
      ("Int32Rep", Int32Rep),
      ("Int64Rep", Int64Rep),
      ("WordRep", WordRep),
      ("Word8Rep", Word8Rep),
      ("Word16Rep", Word16Rep),
      ("Word32Rep", Word32Rep),
      ("Word64Rep", Word64Rep),
      ("AddrRep", AddrRep),
      ("FloatRep", FloatRep),
      ("DoubleRep", DoubleRep)
    ]

levityFromSurfaceType :: Type -> TcM Levity
levityFromSurfaceType ty =
  case peelTypeHead ty of
    TCon name _
      | T.dropWhile (== '\'') (nameText name) == "Lifted" -> pure Lifted
      | T.dropWhile (== '\'') (nameText name) == "Unlifted" -> pure Unlifted
    _ -> emitError NoSourceSpan (OtherError ("invalid Levity: " <> take 80 (show ty))) >> pure Lifted

occursInRuntimeRep :: Unique -> RuntimeRep -> Bool
occursInRuntimeRep needle runtimeRep =
  case runtimeRep of
    VecRep {} -> False
    TupleRep reps -> any (occursInRuntimeRep needle) reps
    SumRep reps -> any (occursInRuntimeRep needle) reps
    RuntimeRepVar unique -> unique == needle
    RuntimeRepMeta unique -> unique == needle
    _ -> False

freeTypeVars :: Type -> [Text]
freeTypeVars = nub . go
  where
    go (TVar name) = [unqualifiedNameText name]
    go (TApp f a) = go f ++ go a
    go (TTypeApp f a) = go f ++ go a
    go (TInfix lhs _ _ rhs) = go lhs ++ go rhs
    go (TFun _ a b) = go a ++ go b
    go (TTuple _ _ args) = concatMap go args
    go (TUnboxedSum args) = concatMap go args
    go (TList _ args) = concatMap go args
    go (TParen inner) = go inner
    go (TAnn _ inner) = go inner
    go (TKindSig inner kindTy) = go inner ++ go kindTy
    go (TContext preds inner) = concatMap go preds ++ go inner
    go (TForall telescope inner) =
      filter
        (`Set.notMember` boundNames)
        (concatMap binderKindVars binders ++ go inner)
      where
        binders = forallTelescopeBinders telescope
        boundNames = Set.fromList (map tyVarBinderName binders)
    go _ = []
    binderKindVars binder = maybe [] go (tyVarBinderKind binder)

splitContext :: Type -> ([Type], Type)
splitContext (TAnn _ inner) = splitContext inner
splitContext (TContext preds inner) = (preds, inner)
splitContext ty = ([], ty)

splitForalls :: Type -> ([TyVarBinder], Type)
splitForalls ty =
  case ty of
    TAnn _ inner -> splitForalls inner
    TParen inner -> splitForalls inner
    TForall telescope inner ->
      let (binders, body) = splitForalls inner
       in (forallTelescopeBinders telescope <> binders, body)
    _ -> ([], ty)

surfacePredToPred :: TvKindEnv -> Type -> TcM Pred
surfacePredToPred tvEnv ty =
  case instanceHeadName ty of
    Just className -> do
      let classNameText = nameText className
          headArgs = instanceHeadTypes ty
      maybeClassInfo <- lookupTyCon classNameText
      case maybeClassInfo of
        Just classInfo -> do
          argKinds <- takeClassArgKinds (length headArgs) <$> defaultKindMetas (tyConKind (tciTyCon classInfo))
          args <- zipWithM (checkSurfaceType tvEnv) headArgs argKinds
          pure (ClassPred (tciTyCon classInfo) args)
        Nothing -> do
          emitError NoSourceSpan (OtherError ("unknown class predicate: " <> T.unpack classNameText))
          abortTc ("missing checked type constructor for class predicate " <> T.unpack classNameText)
    Nothing -> do
      emitError NoSourceSpan (OtherError ("invalid class predicate: " <> show ty))
      abortTc "invalid checked class predicate"

classPredicateArgKinds :: Text -> Int -> TcM [Kind]
classPredicateArgKinds className argCount = do
  mInfo <- lookupTyCon className
  case mInfo of
    Just info -> takeClassArgKinds argCount <$> defaultKindMetas (tyConKind (tciTyCon info))
    Nothing -> mapM (const freshKindMeta) [1 .. argCount]

takeClassArgKinds :: Int -> Kind -> [Kind]
takeClassArgKinds n kind
  | n <= 0 = []
  | otherwise =
      case kind of
        KFun arg rest -> arg : takeClassArgKinds (n - 1) rest
        _ -> replicate n KType

tupleTyConText :: TupleFlavor -> Int -> Text
tupleTyConText flavor arity =
  case flavor of
    Boxed -> boxedTupleTyConName arity
    Unboxed -> unboxedTupleTyConName arity

tupleTyConModule :: TupleFlavor -> Text
tupleTyConModule flavor =
  case flavor of
    Boxed -> "GHC.Tuple"
    Unboxed -> "GHC.Types"

bars :: Int -> Text
bars n
  | n <= 0 = ""
  | otherwise = mconcat (replicate n "|")
