{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
-- editorconfig-checker-disable-file
-- GHC doesn't like the definition of 'TrySpecializeAsVar'.
{-# OPTIONS_GHC -fno-warn-redundant-constraints #-}

module PlutusCore.Builtin.Elaborate
  ( ElaborateFromTo
  ) where

import PlutusCore.Builtin.KnownTypeAst
import PlutusCore.Builtin.Polymorphism
import PlutusCore.Core.Type

import Data.Kind qualified as GHC
import Data.Type.Bool
import Data.Type.Equality
import GHC.TypeLits

{-| Check if two values of different kinds are in fact the same value (with the same kind).
A heterogeneous version of @Type.Equality.(==)@. -}
type (===) :: forall a b. a -> b -> Bool
type family x === y where
  x === x = 'True
  x === y = 'False

-- Explained in detail in https://github.com/effectfully-ou/sketches/tree/cbf3ee9d11e0e3d4fc397ce7bf419418224122e2/poly-type-of-saga/part1-try-unify
{-| Unify two values if possible, otherwise leave them alone. Useful for instantiating type
variables. -}
type TryUnify :: forall a b. Bool -> a -> b -> GHC.Constraint
class same ~ (x === y) => TryUnify same x y

instance (x === y) ~ 'False => TryUnify 'False x y
instance {-# INCOHERENT #-} (x ~~ y, same ~ 'True) => TryUnify same x y

{-| Unify two values unless they're obviously distinct (in which case do nothing).
Allows us to detect and specialize type variables, since a type variable is not obviously
distinct from anything else and so the INCOHERENT instance of 'TryUnify' gets triggered and the
variable gets unified with whatever we want it to. -}
type (~?~) :: forall a b. a -> b -> GHC.Constraint
type x ~?~ y = TryUnify (x === y) x y

-- | Get the element at an @i@th position in a list.
type Lookup :: forall a. Nat -> [a] -> a
type family Lookup n xs where
  Lookup _ '[] = TypeError ('Text "Not enough elements")
  Lookup 0 (x ': xs) = x
  Lookup n (_ ': xs) = Lookup (n - 1) xs

{-| Get the name at the @i@th position in the list of default names. We could use @a_0@, @a_1@,
@a_2@ etc instead, but @a@, @b@, @c@ etc are nicer. -}
type GetName :: GHC.Type -> Nat -> Symbol
type family GetName k i where
  GetName GHC.Type i = Lookup i '["a", "b", "c", "d", "e", "i", "j", "k", "l"]
  GetName _ i = Lookup i '["f", "g", "h", "m", "n"] -- For higher-kinded types.

-- | Apply the function stored in the provided 'Maybe' if there's one.
type MaybeApply :: forall k. Maybe (k -> k) -> k -> k
type family MaybeApply mayVal x where
  MaybeApply 'Nothing a = a
  MaybeApply ('Just f) a = f a

{-| Try to specialize @a@ as a type representing a PLC type variable.
@i@ is a fresh id and @j@ is a final one (either @i + 1@ or @i@ depending on whether
specialization attempt is successful or not).
@mw@ is for wrapping 'TyVarRep', if there's a wrapper inside (see 'HandleHole' for how it's
used). -}
type TrySpecializeAsVar :: forall k. Nat -> Nat -> Maybe (k -> k) -> k -> GHC.Constraint
class TrySpecializeAsVar i j mw a | i mw a -> j

instance
  ( var ~ MaybeApply mw (TyVarRep @k ('TyNameRep (GetName k i) i))
  , -- Try to unify @a@ with a freshly created @var@.
    a ~?~ var
  , -- If @a@ is equal to @var@ then unification was successful and we just used the fresh id and
    -- so we need to bump it up. Otherwise @var@ was discarded and so the fresh id is still fresh.
    -- Replacing @(===)@ with @(==)@ causes errors at use site, for whatever reason.
    j ~ If (a === var) (i + 1) i
  )
  => TrySpecializeAsVar i j mw (a :: k)

type NoAppliedVarsHeader =
  'Text "A built-in function is not allowed to have applied type variables in its type"

-- See Note [Rep vs Type context].
-- | Throw an error telling the user not to apply type variables to anything.
type ThrowNoAppliedVars :: (GHC.Type -> GHC.Type) -> GHC.Constraint
type family ThrowNoAppliedVars hole where
  -- In the Rep context higher-kinded type variables are allowed, but need to be applied via
  -- 'TyAppRep', hence the error message.
  ThrowNoAppliedVars RepHole =
    TypeError
      ( NoAppliedVarsHeader
          ':$$: 'Text "To fix this error apply type variables via explicit ‘TyAppRep’"
      )
  -- In the Type context no higher-kinded type variables are allowed.
  ThrowNoAppliedVars TypeHole =
    TypeError
      ( NoAppliedVarsHeader
          ':$$: 'Text "To fix this error specialize all higher-kinded type variables"
      )
  -- In case we add more contexts.
  ThrowNoAppliedVars _ =
    TypeError
      ( NoAppliedVarsHeader
          ':$$: 'Text "Internal error: the context is not recognized. Please report"
      )

-- | Check that the higher-kinded type does not represent a PLC type variable and if it does.
type CheckNotAppliedVar :: forall k. (GHC.Type -> GHC.Type) -> k -> GHC.Constraint
type family CheckNotAppliedVar hole a where
  CheckNotAppliedVar hole (TyVarRep _) = ThrowNoAppliedVars hole
  CheckNotAppliedVar _ _ = ()

{-| Try to specialize the head of a (possibly nullary) type application to a type representing a
PLC variable and throw if that succeeds or if it already was one. -}
type TrySpecializeHeadAsVar
  :: forall a b. Nat -> Nat -> (GHC.Type -> GHC.Type) -> (a -> b) -> GHC.Constraint
class TrySpecializeHeadAsVar i j hole f | i f -> j

-- | Recurse to reach the head.
instance {-# OVERLAPPABLE #-} TrySpecializeHeadAsVar i j hole f => TrySpecializeHeadAsVar i j hole (f x)

{-| Reached the head, it's a 'TyVarRep', throwing. Mostly to ensure that a 'TyVarRep' doesn't slip
in via user input, but also out of pure paranoia (GHC seems to occasionally solve constraints
for the same type multiple times by racing through different routes in the presence of
@INCOHERENT@ pragmas, but this is not confirmed. It wouldn't be unreasonable, though). -}
instance {-# OVERLAPPING #-} (ThrowNoAppliedVars hole, i ~ j) => TrySpecializeHeadAsVar i j hole (TyVarRep name)

-- | Reached the head, try to specialize it as a variable and throw if that succeeds.
instance
  {-# INCOHERENT #-}
  ( TrySpecializeAsVar i j 'Nothing f
  , CheckNotAppliedVar hole f
  )
  => TrySpecializeHeadAsVar i j hole f

{-| Try to specialize @a@ as a type representing a PLC type variable.
Same as 'TrySpecializeAsVar' (in particular, the parameters are the same), except this one also
checks if the given type is a type application, in which case it tries to specialize the head of the
application to a type representing a PLC type variable and throws if that succeeds.

We need this because blindly specializing an @f A@ (where @f@ is a type variable and @A@ is an
arbitrary type) in the Type context would give us @Opaque val A@ (by @f ~ Opaque val@) instead of
@Opaque val (f' A)@, which would be very confusing to the user: the only reasonable way to elaborate
an application of a type variable in Haskell is to make an application of a type variable in PLC,
not to randomly drop the type variable.

There's no way we could elaborate an applied type variable without surprising the user, because no
instantiation of @f@ could turn @f A@ into @Opaque val (f' A)@, hence we simply forbid applied type
variables and throw upon encountering one. It would be a pain to handle an iterated application of a
type representing a PLC type variable anyway.

It's not possible to determine if @Opaque val A@ is a specialization of @f a@ or @a@ (for a type
variable @a@), hence we have to make this whole check during elaboration and not any later.

Regardless of whether the given argument is a type application or not, it's attempted to be
specialized as a type representing a PLC type variable, because such a type is a type application
itself. Unifying an already known type variable with a fresh one is useful when the name part
of the former is unknown. -}
type TrySpecializeAsUnappliedVar
  :: forall k. Nat -> Nat -> (GHC.Type -> GHC.Type) -> Maybe (k -> k) -> k -> GHC.Constraint
class TrySpecializeAsUnappliedVar i j hole mw a | i mw a -> j

instance
  ( TrySpecializeHeadAsVar i j hole f
  , TrySpecializeAsVar j k mw (f x)
  )
  => TrySpecializeAsUnappliedVar i k hole mw (f x)
instance {-# INCOHERENT #-} TrySpecializeAsVar i j mw a => TrySpecializeAsUnappliedVar i j hole mw a

-- See Note [Rep vs Type context]
{-| First try to specialize the hole using 'TrySpecializeAsVar' and then recurse on the result of
that using 'HandleHoles'.
@i@ is a fresh id and @j@ is a final one as in 'TrySpecializeAsVar', but since 'HandleHole' can
specialize multiple variables, @j@ can be equal to @i + n@ for any @n@ (including @0@). -}
type HandleHole :: (GHC.Type -> GHC.Type) -> Nat -> Nat -> GHC.Type -> Hole -> GHC.Constraint
class HandleHole uni i j val hole | uni i val hole -> j

-- In the Rep context @x@ is attempted to be specialized as a 'TyVarRep'.
instance
  ( TrySpecializeAsUnappliedVar i j RepHole 'Nothing x
  , HandleHoles uni j k val RepHole x
  )
  => HandleHole uni i k val (RepHole x)

-- In the Type context @a@ is attempted to be specialized as a 'TyVarRep' wrapped in @Opaque val@.
instance
  ( TrySpecializeAsUnappliedVar i j TypeHole ('Just (Opaque val)) a
  , HandleHoles uni j k val TypeHole a
  )
  => HandleHole uni i k val (TypeHole a)

{-| Call 'HandleHole' over each hole from the list, threading the state (the fresh unique) through
the calls. -}
type HandleHolesGo :: (GHC.Type -> GHC.Type) -> Nat -> Nat -> GHC.Type -> [Hole] -> GHC.Constraint
class HandleHolesGo uni i j val holes | uni i val holes -> j

instance i ~ j => HandleHolesGo uni i j val '[]
instance
  ( HandleHole uni i j val hole
  , HandleHolesGo uni j k val holes
  )
  => HandleHolesGo uni i k val (hole ': holes)

{-| If the outermost constructor of the second argument is known and happens to be one of the
constructors of the list data type, then the second argument is returned back. Otherwise the
first one is returned, which is meant to be a custom type error. -}
type ThrowOnStuckList :: forall a. [a] -> [a] -> [a]
type family ThrowOnStuckList err xs where
  ThrowOnStuckList _ '[] = '[]
  ThrowOnStuckList _ (x ': xs) = x ': xs
  ThrowOnStuckList err _ = err

type UnknownTypeError :: forall a any. GHC.Type -> a -> any
type family UnknownTypeError val x where
  UnknownTypeError val x =
    TypeError
      ( 'Text "No instance for ‘KnownTypeAst "
          ':<>: 'ShowType (UniOf val)
          ':<>: 'Text " ("
          ':<>: 'ShowType x
          ':<>: 'Text ")’"
          ':$$: 'Text "Which means"
          ':$$: 'Text "  ‘" ':<>: 'ShowType x ':<>: 'Text "’"
          ':$$: 'Text "is neither a built-in type, nor one of the control types."
          ':$$: 'Text "If it can be represented in terms of one of the built-in types"
          ':$$: 'Text "  then go add the instance (you may need a few others too)"
          ':$$: 'Text "  alongside the instance for the built-in type."
          ':$$: 'Text "Otherwise you may need to add a new built-in type"
          ':$$: 'Text "  (provided you're doing something that can be supported in principle)"
      )

-- | Get the holes of @x@ and recurse into them.
type HandleHoles
  :: forall a. (GHC.Type -> GHC.Type) -> Nat -> Nat -> GHC.Type -> (GHC.Type -> GHC.Type) -> a -> GHC.Constraint
type HandleHoles uni i j val hole x =
  -- Here we detect a stuck application of 'ToHoles' and throw 'UnknownTypeError' on it.
  -- See https://blog.csongor.co.uk/report-stuck-families for a detailed description of how
  -- detection of stuck type families works.
  HandleHolesGo uni i j val (ThrowOnStuckList (UnknownTypeError val x) (ToHoles uni hole x))

-- Check out the following for a detailed explanation of the idea (after learning about 'TryUnify'):
-- https://github.com/effectfully-ou/sketches/tree/cbf3ee9d11e0e3d4fc397ce7bf419418224122e2/poly-type-of-saga/part2-enumerate-type-vars
{-| Specialize each Haskell type variable in @a@ as a type representing a PLC type variable.
@i@ is a fresh id and @j@ is a final one as in 'TrySpecializeAsVar', but since 'HandleHole' can
specialize multiple variables, @j@ can be equal to @i + n@ for any @n@ (including @0@). -}
type ElaborateFromTo
  :: (GHC.Type -> GHC.Type) -> Nat -> Nat -> GHC.Type -> GHC.Type -> GHC.Constraint
type ElaborateFromTo uni i j val a = HandleHole uni i j val (TypeHole a)