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

module Symantic.Typer.Unify where

import Control.Monad (Monad (..))
import Data.Bool (otherwise)
import Data.Either (Either (..))
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

-- * Type 'Subst'

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

deriving instance (Show prov) => Show (Subst prov vs)
instance Semigroup (Subst prov vs) where
  (<>) = unionSubst
instance Monoid (Subst prov 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 prov vs -> Subst prov vs -> Subst prov 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 prov vs = forall t. VT (Var prov vs t) (Ty prov vs t)

deriving instance (Show prov) => Show (VT prov vs)

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

lookupSubst :: Var prov vs v -> Subst prov vs -> Maybe (Ty prov 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 prov vs v -> Ty prov vs v -> t prov vs a -> t prov vs a

  -- | Substitute all the 'Var's which have a match in given 'Subst'.
  subst ::
    Subst prov vs -> t prov vs a -> t prov vs a
instance Substable Ty where
  substVarUnsafe _v _r t@TyConst{} = t
  substVarUnsafe v r ty@TyProv{tyUnProv} =
    ty{tyUnProv = substVarUnsafe v r tyUnProv}
  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 prov len fam as) =
  --  TyFam prov len fam $ substVarUnsafe v r as

  subst _s t@TyConst{} = t
  subst s ty@TyProv{tyUnProv} = ty{tyUnProv = subst s tyUnProv}
  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 prov len fam as) = TyFam prov 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 prov vs v ->
  Ty prov vs v ->
  t prov vs a ->
  Maybe (t prov 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 prov
  = -- | /occurence check/: a 'Var' is unified with a 'Type'
    --   which contains this same 'Var'.
    ErrorUnifyVarLoop IndexVar (TyVT prov)
  | -- | Two 'TyConst's should be the same, but are different.
    ErrorUnifyConst (TyVT prov) (TyVT prov)
  | -- | Two 'Kind's should be the same, but are different.
    ErrorUnifyKind (KindK prov) (KindK prov) (TyVT prov) (TyVT prov)
  | -- | Cannot unify those two 'Type's.
    ErrorUnifyType (TyVT prov) (TyVT prov)

deriving instance Eq (ErrorUnify prov)
deriving instance (Show prov) => Show (ErrorUnify prov)

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

data SourceTypeArg prov
  = SourceTypeArg (TyVT prov)

-- | Return the left spine of a 'Ty':
-- the root 'Ty' and its 'Ty' parameters,
-- from the left to the right.
spineTy ::
  forall prov vs t.
  -- MC.MonadWriter (SourceTypeArg prov) (Sourced prov) =>
  Ty prov vs t ->
  (TyT prov vs, [TyT prov vs])
spineTy = go []
  where
    go :: forall kx (x :: kx). [TyT prov vs] -> Ty prov vs x -> (TyT prov vs, [TyT prov 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)

-- | Return the /most general unification/ of two 'Type's, when it exists.
unifyType ::
  forall xyK (x :: xyK) (y :: xyK) vs prov.
  (forall k. ProvenanceKindOf (Ty @k) prov) =>
  (forall k. ProvenanceKindOf (Var @k) prov) =>
  Show prov =>
  Subst prov vs ->
  Ty prov vs (x :: xyK) ->
  Ty prov vs (y :: xyK) ->
  Either (ErrorUnify prov) (Subst prov 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 -> Left $ ErrorUnifyConst (TyVT xRootTy) (TyVT yRootTy)
        _ -> Left $ ErrorUnifyType (TyVT x) (TyVT y)
  where
    xK = kindOf x
    goVar ::
      forall k (a :: k) (b :: k) ws.
      (forall k'. ProvenanceKindOf (Ty @k') prov) =>
      (forall k'. ProvenanceKindOf (Var @k') prov) =>
      Show prov =>
      Subst prov ws ->
      Var prov ws a ->
      Ty prov ws b ->
      Either (ErrorUnify prov) (Subst prov ws)
    goVar sub va b =
      case va `lookupSubst` sub of
        Just a -> unifyType sub b a
        Nothing ->
          case sub `subst` b of
            TyVar{tyVar = vb} | Just HRefl <- va `eqTy` vb -> Right sub
            b'
              | va `varOccursIn` b' -> Left $ ErrorUnifyVarLoop (indexVar va) (TyVT b')
              | HRefl :: a :~~: b <- unsafeCoerce HRefl ->
                  Right $ insertSubst va b' mempty <> sub
    goZip ::
      forall ws.
      Subst prov ws ->
      [TyT prov ws] ->
      [TyT prov ws] ->
      Either (ErrorUnify prov) (Subst prov ws)
    goZip sub [] [] = Right sub
    goZip sub (TyT a : as) (TyT b : bs) =
      case aK `eqTy` bK of
        Nothing -> Left $ ErrorUnifyKind (TyT aK) (TyT bK) (TyVT a) (TyVT b)
        Just HRefl -> unifyType sub a b >>= \sub' -> goZip sub' as bs
      where
        aK = kindOf a
        bK = kindOf b
    goZip _vs _a _b = Prelude.error "[BUG] unifyType: kinds mismatch"