{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE UndecidableInstances #-}

-- | A more strongly typed alternative to 'Nix.Match'
module Nix.Match.Typed
  ( matchNix,
    matchNixLoc,
    TypedMatcher (..),
    TypedMatch (..),
    get,
    getOptional,
    matchTyped,
    findMatchesTyped,
  )
where

import Control.Category ((>>>))
import Data.Coerce (coerce)
import Data.Data
import Data.Fix
import Data.Generics.Aliases
import Data.Kind (Constraint)
import Data.Maybe
import qualified Data.Text as T
import Data.Type.Equality (type (==))
import GHC.TypeLits
  ( ErrorMessage (..),
    KnownSymbol,
    Symbol,
    TypeError,
    symbolVal,
  )
import Language.Haskell.TH
  ( Exp (AppE, VarE),
    ExpQ,
    Pat (..),
    PatQ,
    Q,
    TyLit (StrTyLit),
    Type (..),
    appTypeE,
    litT,
    mkName,
    newName,
    strTyLit,
    tupE,
    tupP,
    varE,
    varP,
  )
import Language.Haskell.TH.Lib
  ( appE,
    conE,
  )
import Language.Haskell.TH.Quote (QuasiQuoter (..))
import Language.Haskell.TH.Syntax
  ( dataToExpQ,
    liftString,
  )
import Nix
import Nix.Match
import Nix.TH

----------------------------------------------------------------
-- Typed matching
----------------------------------------------------------------

-- | A QuasiQuoter for safely generating 'TypedMatcher's from nix source
--
-- The expression has the type @'TypedMatcher' opts reqs 'NExprF'@ where @opts@
-- and @reqs@ are the optional and required holes from the source expression.
--
-- The pattern, if matched, will bring into scope variables named according to
-- the holes present in the expression. These will have type 'NExpr' if they
-- are required, and @Maybe 'NExpr'@ if they are optional.
--
-- This requires ViewPatterns, TypeApplications and DataKinds
--
-- >>> case [nix|{a="hello";}|] of [matchNix|{a=^a;}|] -> a
-- Fix (NStr (DoubleQuoted [Plain "hello"]))
--
-- >>> :t [matchNix|{a = ^a; b = {c = ^c; _d = ^d;};}|]
-- [matchNix|{a = ^a; b = {c = ^c; _d = ^d;};}|] :: TypedMatcher '["d"] '["a", "c"] NExprF
--
-- >>> [matchNix|let a = ^a; _b = ^b; in x|] = undefined
-- >>> :t (a, b)
-- (a, b) :: (Fix NExprF, Maybe (Fix NExprF))
matchNix :: QuasiQuoter
matchNix =
  QuasiQuoter
    { quoteExp = typedMatcherExp,
      quotePat = typedMatcherPat,
      quoteDec = error "No dec quoter for typedMatcher",
      quoteType = error "No type quoter for typedMatcher"
    }

-- | A QuasiQuoter for safely generating 'TypedMatcher's from nix source along
-- with source location annotations
--
-- The expression has the type @'TypedMatcher' opts reqs 'NExprLocF'@ where
-- @opts@ and @reqs@ are the optional and required holes from the source
-- expression.
--
-- This requires ViewPatterns, TypeApplications and DataKinds
--
-- The pattern, if matched, will bring into scope variables named according to
-- the holes present in the expression. These will have type 'NExprLoc' if they
-- are required, and @Maybe 'NExprLoc'@ if they are optional.
matchNixLoc :: QuasiQuoter
matchNixLoc =
  QuasiQuoter
    { quoteExp = typedMatcherLocExp,
      quotePat = typedMatcherLocPat,
      quoteDec = error "No dec quoter for typedMatcherLoc",
      quoteType = error "No type quoter for typedMatcherLoc"
    }

-- | A matcher with the names of the required and optional holes encoded at the
-- type level.
newtype TypedMatcher (opts :: [Symbol]) (reqs :: [Symbol]) t = TypedMatcher {unTypedMatcher :: WithHoles t VarName}

-- | The results of matching with a 'TypedMatcher'. The values in the required
-- list are guaranteed to be present. The values in the optional list may be
-- present. Use 'get' and 'getOptional' to extract them safely.
newtype TypedMatch (opts :: [Symbol]) (reqs :: [Symbol]) a
  = TypedMatch [(T.Text, a)]

-- | Extract a required key from a match
get ::
  forall x opts reqs a.
  (Elem "Required" x reqs, KnownSymbol x) =>
  TypedMatch opts reqs a ->
  a
get (TypedMatch ms) =
  fromMaybe (error "Required key not present in TypedMatch") $
    lookup (T.pack (symbolVal (Proxy @x))) ms

-- | Maybe extract an optional key from a match
getOptional ::
  forall x opts reqs a.
  (Elem "Optional" x opts, KnownSymbol x) =>
  TypedMatch opts reqs a ->
  Maybe a
getOptional (TypedMatch ms) = lookup (T.pack (symbolVal (Proxy @x))) ms

-- | A typed version of 'match'
matchTyped ::
  Matchable t =>
  TypedMatcher opts reqs t ->
  Fix t ->
  Maybe (TypedMatch opts reqs (Fix t))
matchTyped = coerce match

-- | A typed version of 'findMatches'
findMatchesTyped ::
  Matchable t =>
  TypedMatcher opts reqs t ->
  Fix t ->
  [(Fix t, TypedMatch opts reqs (Fix t))]
findMatchesTyped = coerce findMatches

typedMatcherExp :: String -> ExpQ
typedMatcherExp =
  fmap snd . typedMatcherGen parseNixText collectHoles addHoles id

typedMatcherLocExp :: String -> ExpQ
typedMatcherLocExp =
  fmap snd
    . typedMatcherGen
      parseNixTextLoc
      collectHolesLoc
      addHolesLoc
      stripAnnotation

typedMatcherPat :: String -> PatQ
typedMatcherPat = typedMatcherPatGen parseNixText collectHoles addHoles id

typedMatcherLocPat :: String -> PatQ
typedMatcherLocPat =
  typedMatcherPatGen parseNixTextLoc collectHolesLoc addHolesLoc stripAnnotation

typedMatcherPatGen ::
  Data a =>
  (T.Text -> Result t) ->
  (t -> ([VarName], [VarName])) ->
  (t -> a) ->
  (t -> NExpr) ->
  String ->
  Q Pat
typedMatcherPatGen parseNix collect add strip s = do
  ((opt, req), matcher) <- typedMatcherGen parseNix collect add strip s
  -- e' <- [|fmap (\x -> $()) . matchTyped $(pure matcher)|]
  x <- newName "x"
  let pat = tupP (varP . mkName . T.unpack . unVarName <$> (req <> opt))
      textSymbol = litT . strTyLit . T.unpack . unVarName
      getters =
        tupE
          ( ((\r -> [|get @($r) $(varE x)|]) . textSymbol <$> req)
              <> ((\o -> [|getOptional @($o) $(varE x)|]) . textSymbol <$> opt)
          )
  [p|(fmap (\ $(varP x) -> $getters) . matchTyped $(pure matcher) -> Just $pat)|]

unVarName :: VarName -> T.Text
unVarName (VarName x) = x

typedMatcherGen ::
  Data a =>
  (T.Text -> Result t) ->
  (t -> ([VarName], [VarName])) ->
  (t -> a) ->
  (t -> NExpr) ->
  String ->
  Q (([VarName], [VarName]), Exp)
typedMatcherGen parseNix collect add strip s = do
  expr <- case parseNix (T.pack s) of
    Left err -> fail $ show err
    Right e -> pure e
  let (opt, req) = collect expr
      optT = symbolList opt
      reqT = symbolList req
      holed = add expr
      exprExp =
        dataToExpQ
          ( const Nothing
              `extQ` metaExp (getFreeVars (strip expr))
              `extQ` (Just . liftText)
          )
          holed
  e <-
    conE 'TypedMatcher `appTypeE` pure optT `appTypeE` pure reqT `appE` exprExp
  pure ((opt, req), e)

liftText :: T.Text -> Q Exp
liftText txt = AppE (VarE 'T.pack) <$> liftString (T.unpack txt)

-- | Make a list of promoted strings
symbolList :: [VarName] -> Type
symbolList =
  foldr
    (\(VarName n) -> (PromotedConsT `AppT` LitT (StrTyLit (T.unpack n)) `AppT`))
    PromotedNilT

-- | Collect optional and required holes
collectHoles :: NExpr -> ([VarName], [VarName])
collectHoles =
  unFix >>> \case
    NSynHole n -> ([], [n])
    NSet _ bs -> foldMap (bindingHoles collectHoles) bs
    NLet bs e -> collectHoles e <> foldMap (bindingHoles collectHoles) bs
    e -> foldMap collectHoles e

-- | Collect optional and required holes
collectHolesLoc :: NExprLoc -> ([VarName], [VarName])
collectHolesLoc =
  unFix >>> \case
    Compose (AnnUnit _ (NSynHole n)) -> ([], [n])
    Compose (AnnUnit _ (NSet _ bs)) -> foldMap (bindingHoles collectHolesLoc) bs
    Compose (AnnUnit _ (NLet bs e)) ->
      collectHolesLoc e <> foldMap (bindingHoles collectHolesLoc) bs
    e -> foldMap collectHolesLoc e

-- | Find the optional and required holees in a binding
bindingHoles :: (r -> ([a], [a])) -> Binding r -> ([a], [a])
bindingHoles f = \case
  b@(NamedVar p _ _)
    | isJust (isOptionalPath p) ->
        let (opt, req) = foldMap f b in (opt <> req, [])
  b -> foldMap f b

----------------------------------------------------------------
-- Helpers
----------------------------------------------------------------

type family Bool' (f :: k) (t :: k) (x :: Bool) :: k where
  Bool' f _ 'False = f
  Bool' _ t 'True = t

type family Elem n x ys :: Constraint where
  Elem n x '[] = TypeError ('Text n ':<>: 'Text " key \"" ':<>: 'Text x ':<>: 'Text "\" not found in TypedMatch")
  Elem n x (y : ys) = Bool' (Elem n x ys) (() :: Constraint) (x == y)