{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE UndecidableSuperClasses #-}

module Plutarch.Repr.Tag (
  PTag (..),
  DeriveAsTag (..),
  TagLiftHelper (..),
) where

import Data.Proxy (Proxy (Proxy))

-- parts have been commented. They will have to be restored after PLiftalbe rework
-- import Data.Coerce (coerce)
import Data.Kind (Type)

-- import GHC.Exts (Any)
import GHC.Generics qualified as GHC
import GHC.TypeError (ErrorMessage (ShowType, Text, (:$$:), (:<>:)), TypeError)
import GHC.TypeLits (type (+))
import Generics.SOP (
  All,
  Code,
  I,
  K (K),
  NP (Nil, (:*)),
  NS (S, Z),
  SOP (SOP),
 )
import Generics.SOP qualified as SOP
import Plutarch.Builtin.Integer (PInteger)

-- import Plutarch.Builtin.Opaque (POpaque, popaque)
import Plutarch.Internal.Lift (pconstant)
import Plutarch.Internal.PlutusType (
  PContravariant',
  PCovariant',
  PInner,
  PVariant',
  PlutusType,
  pcon',
  pmatch',
 )
import Plutarch.Internal.Term (S, Term)
import Plutarch.Repr.Internal (groupHandlers)
import Plutarch.Repr.Newtype (DeriveAsNewtype (DeriveAsNewtype))
import Plutarch.TermCont (pletC, unTermCont)

-- | @since WIP
newtype PTag (struct :: [S -> Type]) (s :: S) = PTag
  { forall (struct :: [S -> Type]) (s :: S).
PTag struct s -> Term s PInteger
unPTag :: Term s PInteger
  -- ^ @since WIP
  }
  deriving stock ((forall x. PTag struct s -> Rep (PTag struct s) x)
-> (forall x. Rep (PTag struct s) x -> PTag struct s)
-> Generic (PTag struct s)
forall (struct :: [S -> Type]) (s :: S) x.
Rep (PTag struct s) x -> PTag struct s
forall (struct :: [S -> Type]) (s :: S) x.
PTag struct s -> Rep (PTag struct s) x
forall x. Rep (PTag struct s) x -> PTag struct s
forall x. PTag struct s -> Rep (PTag struct s) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall (struct :: [S -> Type]) (s :: S) x.
PTag struct s -> Rep (PTag struct s) x
from :: forall x. PTag struct s -> Rep (PTag struct s) x
$cto :: forall (struct :: [S -> Type]) (s :: S) x.
Rep (PTag struct s) x -> PTag struct s
to :: forall x. Rep (PTag struct s) x -> PTag struct s
GHC.Generic)

-- | @since WIP
instance SOP.Generic (PTag struct s)

-- | @since WIP
deriving via DeriveAsNewtype (PTag struct) instance PlutusType (PTag struct)

-- | @since WIP
newtype DeriveAsTag (a :: S -> Type) s = DeriveAsTag
  { forall (a :: S -> Type) (s :: S). DeriveAsTag a s -> a s
unDeriveAsTag :: a s
  -- ^ @since WIP
  }

{- | This derives tag-only PlutusType automatically. Resulted instances will use `PInteger` as underlying type, making this much more efficient than using regular Data/Scott/SOP based encoding. As name suggests, types with no-argument constructors can use this.

Example:
@@
data PFoo s = A | B | C | D | E
  deriving stock (GHC.Generic, Show)
  deriving anyclass (PEq, PIsData)
  deriving (PlutusType, PLiftable)
    via DeriveAsTag PFoo

instance SOP.Generic (PFoo s)
@@

@since WIP
-}
instance (forall s. TagTypeConstraints s a struct) => PlutusType (DeriveAsTag a) where
  type PInner (DeriveAsTag a) = PInteger
  type PCovariant' (DeriveAsTag a) = (PCovariant' a)
  type PContravariant' (DeriveAsTag a) = (PContravariant' a)
  type PVariant' (DeriveAsTag a) = (PVariant' a)
  pcon' :: forall s. DeriveAsTag a s -> Term s (PInner (DeriveAsTag a))
  pcon' :: forall (s :: S). DeriveAsTag a s -> Term s (PInner (DeriveAsTag a))
pcon' (DeriveAsTag a s
x) =
    forall (a :: S -> Type) (s :: S).
PLiftable a =>
AsHaskell a -> Term s a
pconstant @PInteger (AsHaskell PInteger -> Term s PInteger)
-> AsHaskell PInteger -> Term s PInteger
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ SOP @Type I (Code (a s)) -> Int
forall k l (h :: (k -> Type) -> l -> Type) (f :: k -> Type)
       (xs :: l).
HIndex @k @l h =>
h f xs -> Int
forall (f :: Type -> Type) (xs :: [[Type]]). SOP @Type f xs -> Int
SOP.hindex (SOP @Type I (Code (a s)) -> Int)
-> SOP @Type I (Code (a s)) -> Int
forall a b. (a -> b) -> a -> b
$ a s -> SOP @Type I (Code (a s))
forall a. Generic a => a -> Rep a
SOP.from a s
x

  pmatch' :: forall s b. Term s (PInner (DeriveAsTag a)) -> (DeriveAsTag a s -> Term s b) -> Term s b
  pmatch' :: forall (s :: S) (b :: S -> Type).
Term s (PInner (DeriveAsTag a))
-> (DeriveAsTag a s -> Term s b) -> Term s b
pmatch' Term s (PInner (DeriveAsTag a))
tag DeriveAsTag a s -> Term s b
f = TermCont @b s (Term s b) -> Term s b
forall (a :: S -> Type) (s :: S).
TermCont @a s (Term s a) -> Term s a
unTermCont (TermCont @b s (Term s b) -> Term s b)
-> TermCont @b s (Term s b) -> Term s b
forall a b. (a -> b) -> a -> b
$ do
    -- plet here because tag might be a big computation
    Term s PInteger
tag' <- Term s (PInner (DeriveAsTag a))
-> TermCont @b s (Term s (PInner (DeriveAsTag a)))
forall {r :: S -> Type} (s :: S) (a :: S -> Type).
Term s a -> TermCont @r s (Term s a)
pletC Term s (PInner (DeriveAsTag a))
tag
    let
      g :: SOP I (Code (a s)) -> Term s b
      g :: SOP @Type I (Code (a s)) -> Term s b
g = DeriveAsTag a s -> Term s b
f (DeriveAsTag a s -> Term s b)
-> (SOP @Type I (Code (a s)) -> DeriveAsTag a s)
-> SOP @Type I (Code (a s))
-> Term s b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a s -> DeriveAsTag a s
forall (a :: S -> Type) (s :: S). a s -> DeriveAsTag a s
DeriveAsTag (a s -> DeriveAsTag a s)
-> (SOP @Type I (Code (a s)) -> a s)
-> SOP @Type I (Code (a s))
-> DeriveAsTag a s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOP @Type I (Code (a s)) -> a s
forall a. Generic a => Rep a -> a
SOP.to

      go :: forall x xs. IsEmpty x => TagMatchHandler s b xs -> TagMatchHandler s b (x ': xs)
      go :: forall (x :: [Type]) (xs :: [[Type]]).
IsEmpty @Type x =>
TagMatchHandler s b xs -> TagMatchHandler s b ((':) @[Type] x xs)
go (TagMatchHandler Integer
-> (SOP @Type I xs -> Term s b)
-> NP @[Type] (K @[Type] (Integer, Term s b)) xs
rest) =
        (Integer
 -> (SOP @Type I ((':) @[Type] x xs) -> Term s b)
 -> NP @[Type] (K @[Type] (Integer, Term s b)) ((':) @[Type] x xs))
-> TagMatchHandler s b ((':) @[Type] x xs)
forall (s :: S) (b :: S -> Type) (struct :: [[Type]]).
(Integer
 -> (SOP @Type I struct -> Term s b)
 -> NP @[Type] (K @[Type] (Integer, Term s b)) struct)
-> TagMatchHandler s b struct
TagMatchHandler ((Integer
  -> (SOP @Type I ((':) @[Type] x xs) -> Term s b)
  -> NP @[Type] (K @[Type] (Integer, Term s b)) ((':) @[Type] x xs))
 -> TagMatchHandler s b ((':) @[Type] x xs))
-> (Integer
    -> (SOP @Type I ((':) @[Type] x xs) -> Term s b)
    -> NP @[Type] (K @[Type] (Integer, Term s b)) ((':) @[Type] x xs))
-> TagMatchHandler s b ((':) @[Type] x xs)
forall a b. (a -> b) -> a -> b
$ \Integer
idx SOP @Type I ((':) @[Type] x xs) -> Term s b
handle ->
          (Integer, Term s b) -> K @[Type] (Integer, Term s b) x
forall k a (b :: k). a -> K @k a b
K (Integer
idx, SOP @Type I ((':) @[Type] x xs) -> Term s b
handle (SOP @Type I ((':) @[Type] x xs) -> Term s b)
-> SOP @Type I ((':) @[Type] x xs) -> Term s b
forall a b. (a -> b) -> a -> b
$ NS @[Type] (NP @Type I) ((':) @[Type] x xs)
-> SOP @Type I ((':) @[Type] x xs)
forall k (f :: k -> Type) (xss :: [[k]]).
NS @[k] (NP @k f) xss -> SOP @k f xss
SOP (NS @[Type] (NP @Type I) ((':) @[Type] x xs)
 -> SOP @Type I ((':) @[Type] x xs))
-> NS @[Type] (NP @Type I) ((':) @[Type] x xs)
-> SOP @Type I ((':) @[Type] x xs)
forall a b. (a -> b) -> a -> b
$ NP @Type I ('[] @Type)
-> NS @[Type] (NP @Type I) ((':) @[Type] ('[] @Type) xs)
forall {k} (a :: k -> Type) (x :: k) (xs :: [k]).
a x -> NS @k a ((':) @k x xs)
Z NP @Type I ('[] @Type)
forall {k} (a :: k -> Type). NP @k a ('[] @k)
Nil) K @[Type] (Integer, Term s b) x
-> NP @[Type] (K @[Type] (Integer, Term s b)) xs
-> NP @[Type] (K @[Type] (Integer, Term s b)) ((':) @[Type] x xs)
forall {k} (a :: k -> Type) (x :: k) (xs :: [k]).
a x -> NP @k a xs -> NP @k a ((':) @k x xs)
:* Integer
-> (SOP @Type I xs -> Term s b)
-> NP @[Type] (K @[Type] (Integer, Term s b)) xs
rest (Integer
idx Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) (\(SOP NS @[Type] (NP @Type I) xs
x) -> SOP @Type I ((':) @[Type] x xs) -> Term s b
handle (SOP @Type I ((':) @[Type] x xs) -> Term s b)
-> SOP @Type I ((':) @[Type] x xs) -> Term s b
forall a b. (a -> b) -> a -> b
$ NS @[Type] (NP @Type I) ((':) @[Type] x xs)
-> SOP @Type I ((':) @[Type] x xs)
forall k (f :: k -> Type) (xss :: [[k]]).
NS @[k] (NP @k f) xss -> SOP @k f xss
SOP (NS @[Type] (NP @Type I) ((':) @[Type] x xs)
 -> SOP @Type I ((':) @[Type] x xs))
-> NS @[Type] (NP @Type I) ((':) @[Type] x xs)
-> SOP @Type I ((':) @[Type] x xs)
forall a b. (a -> b) -> a -> b
$ NS @[Type] (NP @Type I) xs
-> NS @[Type] (NP @Type I) ((':) @[Type] x xs)
forall {k} (a :: k -> Type) (xs :: [k]) (x :: k).
NS @k a xs -> NS @k a ((':) @k x xs)
S NS @[Type] (NP @Type I) xs
x)

      handlers' :: TagMatchHandler s b struct
      handlers' :: TagMatchHandler s b struct
handlers' = Proxy @([Type] -> Constraint) (IsEmpty @Type)
-> TagMatchHandler s b ('[] @[Type])
-> (forall (y :: [Type]) (ys :: [[Type]]).
    (IsEmpty @Type y, All @[Type] (IsEmpty @Type) ys) =>
    TagMatchHandler s b ys -> TagMatchHandler s b ((':) @[Type] y ys))
-> TagMatchHandler s b struct
forall k (c :: k -> Constraint) (xs :: [k])
       (proxy :: (k -> Constraint) -> Type) (r :: [k] -> Type).
All @k c xs =>
proxy c
-> r ('[] @k)
-> (forall (y :: k) (ys :: [k]).
    (c y, All @k c ys) =>
    r ys -> r ((':) @k y ys))
-> r xs
forall (proxy :: ([Type] -> Constraint) -> Type)
       (r :: [[Type]] -> Type).
proxy (IsEmpty @Type)
-> r ('[] @[Type])
-> (forall (y :: [Type]) (ys :: [[Type]]).
    (IsEmpty @Type y, All @[Type] (IsEmpty @Type) ys) =>
    r ys -> r ((':) @[Type] y ys))
-> r struct
SOP.cpara_SList (forall {k} (t :: k). Proxy @k t
forall (t :: [Type] -> Constraint). Proxy @([Type] -> Constraint) t
Proxy @IsEmpty) ((Integer
 -> (SOP @Type I ('[] @[Type]) -> Term s b)
 -> NP @[Type] (K @[Type] (Integer, Term s b)) ('[] @[Type]))
-> TagMatchHandler s b ('[] @[Type])
forall (s :: S) (b :: S -> Type) (struct :: [[Type]]).
(Integer
 -> (SOP @Type I struct -> Term s b)
 -> NP @[Type] (K @[Type] (Integer, Term s b)) struct)
-> TagMatchHandler s b struct
TagMatchHandler ((Integer
  -> (SOP @Type I ('[] @[Type]) -> Term s b)
  -> NP @[Type] (K @[Type] (Integer, Term s b)) ('[] @[Type]))
 -> TagMatchHandler s b ('[] @[Type]))
-> (Integer
    -> (SOP @Type I ('[] @[Type]) -> Term s b)
    -> NP @[Type] (K @[Type] (Integer, Term s b)) ('[] @[Type]))
-> TagMatchHandler s b ('[] @[Type])
forall a b. (a -> b) -> a -> b
$ \Integer
_ SOP @Type I ('[] @[Type]) -> Term s b
_ -> NP @[Type] (K @[Type] (Integer, Term s b)) ('[] @[Type])
forall {k} (a :: k -> Type). NP @k a ('[] @k)
Nil) TagMatchHandler s b ys -> TagMatchHandler s b ((':) @[Type] y ys)
forall (y :: [Type]) (ys :: [[Type]]).
(IsEmpty @Type y, All @[Type] (IsEmpty @Type) ys) =>
TagMatchHandler s b ys -> TagMatchHandler s b ((':) @[Type] y ys)
forall (x :: [Type]) (xs :: [[Type]]).
IsEmpty @Type x =>
TagMatchHandler s b xs -> TagMatchHandler s b ((':) @[Type] x xs)
go

      handlers :: [(Integer, Term s b)]
      handlers :: [(Integer, Term s b)]
handlers = NP @[Type] (K @[Type] (Integer, Term s b)) struct
-> CollapseTo @[Type] @[[Type]] (NP @[Type]) (Integer, Term s b)
forall (xs :: [[Type]]) a.
SListIN @[Type] @[[Type]] (NP @[Type]) xs =>
NP @[Type] (K @[Type] a) xs
-> CollapseTo @[Type] @[[Type]] (NP @[Type]) a
forall k l (h :: (k -> Type) -> l -> Type) (xs :: l) a.
(HCollapse @k @l h, SListIN @k @l h xs) =>
h (K @k a) xs -> CollapseTo @k @l h a
SOP.hcollapse (NP @[Type] (K @[Type] (Integer, Term s b)) struct
 -> CollapseTo @[Type] @[[Type]] (NP @[Type]) (Integer, Term s b))
-> NP @[Type] (K @[Type] (Integer, Term s b)) struct
-> CollapseTo @[Type] @[[Type]] (NP @[Type]) (Integer, Term s b)
forall a b. (a -> b) -> a -> b
$ TagMatchHandler s b struct
-> Integer
-> (SOP @Type I struct -> Term s b)
-> NP @[Type] (K @[Type] (Integer, Term s b)) struct
forall (s :: S) (b :: S -> Type) (struct :: [[Type]]).
TagMatchHandler s b struct
-> Integer
-> (SOP @Type I struct -> Term s b)
-> NP @[Type] (K @[Type] (Integer, Term s b)) struct
unTagMatchHandler TagMatchHandler s b struct
handlers' Integer
0 SOP @Type I struct -> Term s b
SOP @Type I (Code (a s)) -> Term s b
g

    Term s b -> TermCont @b s (Term s b)
forall a. a -> TermCont @b s a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Term s b -> TermCont @b s (Term s b))
-> Term s b -> TermCont @b s (Term s b)
forall a b. (a -> b) -> a -> b
$ [(Integer, Term s b)] -> Term s PInteger -> Term s b
forall (s :: S) (r :: S -> Type).
[(Integer, Term s r)] -> Term s PInteger -> Term s r
groupHandlers [(Integer, Term s b)]
handlers Term s PInteger
tag'

-- | @since WIP
newtype TagLiftHelper r struct = TagLiftHelper
  { forall r (struct :: [[Type]]).
TagLiftHelper r struct -> Integer -> (SOP @Type I struct -> r) -> r
unTagLiftHelper :: Integer -> (SOP I struct -> r) -> r
  -- ^ @since WIP
  }

-- Helpers

class x ~ '[] => IsEmpty (x :: [k])

instance x ~ '[] => IsEmpty '[]

type family TagTypePrettyError' n (xs :: [[Type]]) :: Bool where
  TagTypePrettyError' n ('[] ': rest) = TagTypePrettyError' (n + 1) rest
  TagTypePrettyError' n (invalid ': _) =
    TypeError
      ( 'Text "DeriveAsTag only supports constructors without arguments. However, at constructor #"
          ':<>: 'ShowType n
          ':<>: 'Text ", I got:"
          ':$$: 'ShowType invalid
      )
  TagTypePrettyError' _ '[] = 'True

type TagTypePrettyError struct = TagTypePrettyError' 1 struct ~ 'True

class (SOP.Generic (a s), TagTypePrettyError (Code (a s)), Code (a s) ~ struct, All IsEmpty struct) => TagTypeConstraints s a struct | s a -> struct
instance (SOP.Generic (a s), TagTypePrettyError (Code (a s)), Code (a s) ~ struct, All IsEmpty struct) => TagTypeConstraints s a struct

newtype TagMatchHandler s b struct = TagMatchHandler
  { forall (s :: S) (b :: S -> Type) (struct :: [[Type]]).
TagMatchHandler s b struct
-> Integer
-> (SOP @Type I struct -> Term s b)
-> NP @[Type] (K @[Type] (Integer, Term s b)) struct
unTagMatchHandler :: Integer -> (SOP I struct -> Term s b) -> NP (K (Integer, Term s b)) struct
  }

-- instance
--   ( PlutusType (DeriveAsTag a)
--   , SOP.Generic (a Any)
--   , TagTypeConstraints Any a struct
--   ) =>
--   PLiftable (DeriveAsTag a)
--   where
--   type AsHaskell (DeriveAsTag a) = a Any
--   type PlutusRepr (DeriveAsTag a) = Integer
--   toPlutarchRepr = toInteger . hindex . from
--   toPlutarch = PLifted . popaque . pconstant @PInteger . toPlutarchRepr @(DeriveAsTag a)
--   fromPlutarchRepr idx =
--     let
--       go :: IsEmpty x => TagLiftHelper r xs -> TagLiftHelper r (x ': xs)
--       go (TagLiftHelper rest) =
--         TagLiftHelper $ \n f ->
--           if idx == n
--             then f $ SOP $ Z Nil
--             else rest (n + 1) \(SOP s) -> f $ SOP $ S s

--       helper :: TagLiftHelper (Maybe (a Any)) (Code (a Any))
--       helper = cpara_SList (Proxy @IsEmpty) (TagLiftHelper \_ _ -> Nothing) go
--      in
--       unTagLiftHelper helper 0 (Just <$> to)
--   fromPlutarch t = do
--     idx <- fromPlutarch @PInteger $ coerce t
--     case fromPlutarchRepr @(DeriveAsTag a) idx of
--       Nothing -> Left TagTypeInvalidIndex
--       Just x -> pure x