{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Symantic.Typer.Unify where

import Control.Monad (Monad (..))
import Control.Monad.Classes as MC (MonadExcept, throw)
import Control.Monad.Classes.Run as MC ()
import Data.Bool (otherwise)
import Data.Eq (Eq (..))
import Data.Function (id, ($))
import Data.Functor ((<$>))
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (Maybe (..))
import Data.Monoid (Monoid (..))
import Data.Semigroup (Semigroup (..))
import Text.Show (Show)
import Type.Reflection ((:~~:) (HRefl))
import Unsafe.Coerce (unsafeCoerce)
import Prelude qualified

import Symantic.Parser.Error (ErrorInj (..))
import Symantic.Typer.Eq (EqTy (eqTy))
import Symantic.Typer.Type (
  IndexVar,
  KindK,
  KindOf (kindOf),
  Ty (TyApp, TyConst, TyVar, tyConst, tyVar),
  TyT (..),
  TyVT (..),
  Tys (..),
  Var,
  VarOccursIn (..),
  indexVar,
  proj_ConstKi,
  type (#>),
 )

-- import Symantic.Typer.Show ()

-- * Type 'Subst'

-- | /Type variable substitution/.
--
-- WARNING: a 'Subst' MUST be without loops, and fully expanded.
newtype Subst meta vs
  = Subst (Map IndexVar (VT meta vs))

deriving instance (Show meta, Monoid meta) => Show (Subst meta vs)
instance Monoid meta => Semigroup (Subst meta vs) where
  (<>) = unionSubst
instance Monoid meta => Monoid (Subst meta vs) where
  mempty = Subst Map.empty
  mappend = (<>)

-- | Unify two 'Subst's.
--
-- NOTE: the union is left-biased: in case of duplicate 'Var's,
-- it keeps the one from the first 'Subst' given.
--
-- NOTE: the first 'Subst' given is applied to the second (with 'subst'),
-- this way each 'Var' directly maps to an expanded 'Type',
-- so that, when using the resulting 'Subst',
-- there is no need to apply it multiple times
-- until there is no more substitution to be done.
unionSubst :: Subst meta vs -> Subst meta vs -> Subst meta vs
unionSubst sx@(Subst x) (Subst y) = Subst $ x `Map.union` ((\(VT v r) -> VT v $ subst sx r) Data.Functor.<$> y)

-- * Type 'VT'

-- | A 'Var' and a 'Type' existentialized over their type index.
data VT meta vs = forall t. VT (Var meta vs t) (Ty meta vs t)

deriving instance (Show meta, Monoid meta) => Show (VT meta vs)

insertSubst :: Monoid meta => Var meta vs v -> Ty meta vs v -> Subst meta vs -> Subst meta vs
insertSubst v t (Subst s) = Subst $ Map.insert (indexVar v) (VT v t) s

lookupSubst :: Monoid meta => Var meta vs v -> Subst meta vs -> Maybe (Ty meta vs v)
lookupSubst v (Subst s)
  | Just (VT v' t) <- Map.lookup (indexVar v) s
  , Just HRefl <- v `eqTy` v' =
      Just t
lookupSubst _v _m = Nothing

-- * Class 'Substable'
class Substable t where
  -- | Like 'substVar', but without the /occurence check/.
  substVarUnsafe ::
    Var meta vs v -> Ty meta vs v -> t meta vs a -> t meta vs a

  -- | Substitute all the 'Var's which have a match in given 'Subst'.
  subst ::
    Subst meta vs -> t meta vs a -> t meta vs a
instance Substable Ty where
  substVarUnsafe _v _r t@TyConst{} = t
  substVarUnsafe v r (TyApp f a) =
    TyApp
      (substVarUnsafe v r f)
      (substVarUnsafe v r a)
  substVarUnsafe v r t@TyVar{tyVar = tv} =
    case v `eqTy` tv of
      Just HRefl -> r
      Nothing -> t

  -- substVarUnsafe v r (TyFam meta len fam as) =
  --  TyFam meta len fam $ substVarUnsafe v r as

  subst _s t@TyConst{} = t
  subst s (TyApp f a) = TyApp (subst s f) (subst s a)
  subst (Subst s) t@TyVar{tyVar = tv} =
    case Map.lookup (indexVar tv) s of
      Nothing -> t
      Just (VT vr r) ->
        case tv `eqTy` vr of
          Nothing -> Prelude.error "[BUG] subst: kind mismatch"
          Just HRefl -> r

-- subst s (TyFam meta len fam as) = TyFam meta len fam $ subst s as
instance Substable (Tys as) where
  substVarUnsafe _v _r TysZ = TysZ
  substVarUnsafe v r (TysS t ts) =
    substVarUnsafe v r t
      `TysS` substVarUnsafe v r ts
  subst _s TysZ = TysZ
  subst s (TysS t ts) = subst s t `TysS` subst s ts

-- | Substitute the given 'Var' by the given 'Type',
-- returning 'Nothing' if this 'Type' contains
-- the 'Var' (occurence check).
substVar ::
  VarOccursIn t =>
  Substable t =>
  Var meta vs v ->
  Ty meta vs v ->
  t meta vs a ->
  Maybe (t meta vs a)
substVar v r t =
  if v `varOccursIn` r
    then Nothing -- NOTE: occurence check
    else Just $ substVarUnsafe v r t

-- ** Type 'ErrorUnify'

-- | Reasons why two 'Type's cannot be unified.
data ErrorUnify meta
  = -- | /occurence check/: a 'Var' is unified with a 'Type'
    --   which contains this same 'Var'.
    ErrorUnifyVarLoop IndexVar (TyVT meta)
  | -- | Two 'TyConst's should be the same, but are different.
    ErrorUnifyConst (TyVT meta) (TyVT meta)
  | -- | Two 'Kind's should be the same, but are different.
    ErrorUnifyKind (KindK meta) (KindK meta)
  | -- | Cannot unify those two 'Type's.
    ErrorUnifyType (TyVT meta) (TyVT meta)

deriving instance Monoid meta => Eq (ErrorUnify meta)
deriving instance (Show meta, Monoid meta) => Show (ErrorUnify meta)

instance ErrorInj (ErrorUnify meta) (ErrorUnify meta) where
  errorInj = id

data SourceTypeArg meta
  = SourceTypeArg (TyVT meta)

-- | Return the left spine of a 'Ty':
-- the root 'Ty' and its 'Ty' parameters,
-- from the left to the right.
spineTy ::
  forall meta vs t.
  -- MC.MonadWriter (SourceTypeArg meta) (Sourced meta) =>
  Ty meta vs t ->
  (TyT meta vs, [TyT meta vs])
spineTy = go []
  where
    go :: forall kx (x :: kx). [TyT meta vs] -> Ty meta vs x -> (TyT meta vs, [TyT meta vs])
    -- Skip the constraint @q@.
    go ctx (TyApp (TyApp (TyConst{tyConst = c}) _q) t)
      | Just HRefl <- proj_ConstKi @_ @(#>) c =
          go ctx t
    -- Add a type to the spine.
    go ctx (TyApp f a) = go (TyT a {-(SourceTypeArg (TyVT ty)-} : ctx) f
    -- Add the root.
    go ctx t = (TyT t {-SourceTypeArg (TyVT ty)-}, ctx)

{-
spineTy
 :: Ty meta ctx ss cs (t::k)
 -> (forall kx (x::kx) xs. Ty meta ctx ss cs x -> Tys meta ctx ss cs xs -> ret)
 -> ret
spineTy = go TysZ
  where
  go :: Tys meta ctx ss cs hs
     -> Ty meta ctx ss cs (t::k)
     -> (forall kx (x::kx) xs. Ty meta ctx ss cs x -> Tys meta ctx ss cs xs -> ret)
     -> ret
  go ctx (TyApp  _src f a) k = go (a `TysS` ctx) f k
  go ctx (Term x _te)    k = go ctx x k
  go ctx (TyAny x)       k = go ctx x k
  go ctx t k = k t ctx
-}

-- | Return the /most general unification/ of two 'Type's, when it exists.
unifyType ::
  forall xyK (x :: xyK) (y :: xyK) vs meta m.
  -- ErrorInj (Con_Kind meta) (ErrorUnify meta) =>
  Monoid meta =>
  MC.MonadExcept (ErrorUnify meta) m =>
  Subst meta vs ->
  Ty meta vs (x :: xyK) ->
  Ty meta vs (y :: xyK) ->
  m (Subst meta vs)
unifyType vs x y =
  case (spineTy x, spineTy y) of
    ((TyT xRootTy, px), (TyT yRootTy, py)) ->
      case (xRootTy, yRootTy) of
        (TyVar{tyVar = xV}, _) | Just HRefl <- xK `eqTy` kindOf xV -> goVar vs xV y
        (_, TyVar{tyVar = yV}) | Just HRefl <- xK `eqTy` kindOf yV -> goVar vs yV x
        (TyConst{tyConst = cx}, TyConst{tyConst = cy})
          | Just HRefl <- cx `eqTy` cy -> goZip vs px py
          | otherwise -> MC.throw $ ErrorUnifyConst (TyVT xRootTy) (TyVT yRootTy)
        _ ->
          case (x, y) of
            (TyApp fx ax, TyApp fy ay) -> goZip vs [TyT fx, TyT ax] [TyT fy, TyT ay]
            _ -> MC.throw $ ErrorUnifyType (TyVT x) (TyVT y)
  where
    xK = kindOf x
    goVar ::
      forall k (a :: k) (b :: k) ws.
      Subst meta ws ->
      Var meta ws a ->
      Ty meta ws b ->
      m (Subst meta ws)
    goVar ws va b =
      case va `lookupSubst` ws of
        Just a -> unifyType ws b a
        Nothing ->
          case ws `subst` b of
            TyVar{tyVar = vb} | Just HRefl <- va `eqTy` vb -> return ws
            b'
              | va `varOccursIn` b' -> MC.throw $ ErrorUnifyVarLoop (indexVar va) (TyVT b')
              | HRefl :: a :~~: b <- unsafeCoerce HRefl ->
                  return $ insertSubst va b' mempty <> ws
    goZip ::
      forall ws.
      Subst meta ws ->
      [TyT meta ws] ->
      [TyT meta ws] ->
      m (Subst meta ws)
    goZip ws [] [] = return ws
    goZip ws (TyT a : as) (TyT b : bs) =
      case aK `eqTy` bK of
        Nothing -> MC.throw $ ErrorUnifyKind (TyT aK) (TyT bK)
        Just HRefl -> unifyType ws a b >>= \ws' -> goZip ws' as bs
      where
        aK = kindOf a
        bK = kindOf b
    goZip _vs _a _b = Prelude.error "[BUG] unifyType: kinds mismatch"