{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PolyKinds #-}
module Data.GADT.Compare.TH
    ( DeriveGEQ(..)
    , DeriveGCompare(..)
    , module Data.GADT.Compare.Monad
    ) where

import Control.Monad
import Control.Monad.Writer
import Data.GADT.TH.Internal
import Data.Functor.Identity
import Data.GADT.Compare
import Data.GADT.Compare.Monad
import Data.Type.Equality ((:~:) (..))
import qualified Data.Set as Set
import Data.Set (Set)
import qualified Data.Map as Map
import qualified Data.Map.Merge.Lazy as Map
import Data.Map (Map)
import Language.Haskell.TH
import Language.Haskell.TH.Datatype

-- A type class purely for overloading purposes
class DeriveGEQ t where
    deriveGEq :: t -> Q [Dec]

instance DeriveGEQ Name where
  deriveGEq typeName = do
    typeInfo <- reifyDatatype typeName
    let instTypes = datatypeInstTypes typeInfo
        paramVars = Set.unions [freeTypeVariables t | t <- instTypes]
        instTypes' = case reverse instTypes of
          [] -> fail "deriveGEq: Not enough type parameters"
          (_:xs) -> reverse xs
        instanceHead = AppT (ConT ''GEq) (foldl AppT (ConT typeName) instTypes')
    (clauses, cxt) <- runWriterT (mapM (geqClause paramVars) (datatypeCons typeInfo))

    return [InstanceD Nothing cxt instanceHead [geqFunction clauses]]

instance DeriveGEQ Dec where
    deriveGEq = deriveForDec ''GEq $ \typeInfo -> do
      let
        instTypes = datatypeInstTypes typeInfo
        paramVars = Set.unions [freeTypeVariables t | t <- instTypes]
      clauses <- mapM (geqClause paramVars) (datatypeCons typeInfo)
      return $ geqFunction clauses

instance DeriveGEQ t => DeriveGEQ [t] where
  deriveGEq [it] = deriveGEq it
  deriveGEq _ = fail "deriveGEq: [] instance only applies to single-element lists"

instance DeriveGEQ t => DeriveGEQ (Q t) where
  deriveGEq = (>>= deriveGEq)

geqFunction :: [Clause] -> Dec
geqFunction clauses = FunD 'geq $ clauses ++ [ Clause [WildP, WildP] (NormalB (ConE 'Nothing)) [] ]
 -- TODO: only include last clause if there's more than one constructor?

geqClause :: Set Name -> ConstructorInfo -> WriterT Cxt Q Clause
geqClause paramVars con = do
  let conName = constructorName con
      argTypes = constructorFields con
      conTyVars = Set.fromList (map tvName (constructorVars con))
      needsGEq argType = not . Set.null $
        Set.intersection (freeTypeVariables argType) (Set.union paramVars conTyVars)
  lArgNames <- forM argTypes $ \_ -> lift $ newName "x"
  rArgNames <- forM argTypes $ \_ -> lift $ newName "y"

  stmts <- forM (zip3 lArgNames rArgNames argTypes) $ \(l, r, t) -> do
    case t of
      AppT tyFun tyArg | needsGEq t -> do
        u <- lift $ reifyInstancesWithRigids paramVars ''GEq [tyFun]
        case u of
          [] -> tell [AppT (ConT ''GEq) tyFun]
          [(InstanceD _ cxt _ _)] -> tell cxt
          _ -> fail $ "More than one instance found for GEq (" ++ show (ppr tyFun) ++ "), and unsure what to do. Please report this."
        lift $ bindS (conP 'Refl []) [| geq $(varE l) $(varE r) |]
      _ -> lift $ noBindS [| guard ($(varE l) == $(varE r)) |]
  ret <- lift $ noBindS [| return Refl |]

  pats <- lift $ sequence
    [ conP conName (map varP lArgNames)
    , conP conName (map varP rArgNames)
    ]
  pure $ Clause pats
    (NormalB (doUnqualifiedE (stmts ++ [ret])))
    []

class DeriveGCompare t where
    deriveGCompare :: t -> Q [Dec]

instance DeriveGCompare Name where
    deriveGCompare typeName = do
      typeInfo <- reifyDatatype typeName
      let instTypes = datatypeInstTypes typeInfo
          paramVars = Set.unions [freeTypeVariables t | t <- instTypes]
          instTypes' = case reverse instTypes of
            [] -> fail "deriveGCompare: Not enough type parameters"
            (_:xs) -> reverse xs
          instanceHead = AppT (ConT ''GCompare) (foldl AppT (ConT typeName) instTypes')
      (clauses, cxt) <- runWriterT (fmap concat $ mapM (gcompareClauses paramVars) (datatypeCons typeInfo))
      dec <- gcompareFunction clauses
      return [InstanceD Nothing cxt instanceHead [dec]]

instance DeriveGCompare Dec where
    deriveGCompare = deriveForDec ''GCompare $ \typeInfo -> do
      let
        instTypes = datatypeInstTypes typeInfo
        paramVars = Set.unions [freeTypeVariables t | t <- instTypes]
      clauses <- mapM (gcompareClauses paramVars) (datatypeCons typeInfo)
      lift $ gcompareFunction (concat clauses)

instance DeriveGCompare t => DeriveGCompare [t] where
    deriveGCompare [it] = deriveGCompare it
    deriveGCompare _ = fail "deriveGCompare: [] instance only applies to single-element lists"

instance DeriveGCompare t => DeriveGCompare (Q t) where
    deriveGCompare = (>>= deriveGCompare)

gcompareFunction :: [Clause] -> Q Dec
gcompareFunction [] = funD 'gcompare [clause [] (normalB [| \x y -> seq x (seq y undefined) |]) []]
gcompareFunction clauses = return $ FunD 'gcompare clauses

gcompareClauses :: Set Name -> ConstructorInfo -> WriterT Cxt Q [Clause]
gcompareClauses paramVars con = do
  let conName = constructorName con
      argTypes = constructorFields con
      conTyVars = Set.fromList (map tvName (constructorVars con))
      needsGCompare argType = not . Set.null $
        Set.intersection (freeTypeVariables argType) (Set.union paramVars conTyVars)

  lArgNames <- forM argTypes $ \_ -> lift $ newName "x"
  rArgNames <- forM argTypes $ \_ -> lift $ newName "y"

  stmts <- forM (zip3 lArgNames rArgNames argTypes) $ \(lArg, rArg, argType) ->
    case argType of
      AppT tyFun tyArg | needsGCompare argType -> do
        u <- lift $ reifyInstancesWithRigids paramVars ''GCompare [tyFun]
        case u of
          [] -> tell [AppT (ConT ''GCompare) tyFun]
          [(InstanceD _ cxt _ _)] -> tell cxt -- this might not be enough, may want to do full instance resolution.
          _ -> fail $ "More than one instance of GCompare (" ++ show (ppr tyFun) ++ ") found, and unsure what to do. Please report this."
        lift $ bindS (conP 'Refl []) [| geq' $(varE lArg) $(varE rArg) |]
      _ -> lift $ noBindS [| compare' $(varE lArg) $(varE rArg) |]

  ret <- lift $ noBindS [| return GEQ |]


  pats <- lift $ sequence
        [ conP conName (map varP lArgNames)
        , conP conName (map varP rArgNames)
        ]
  let main = Clause pats
        (NormalB (AppE (VarE 'runGComparing) (doUnqualifiedE (stmts ++ [ret]))))
        []
      lt = Clause [RecP conName [], WildP] (NormalB (ConE 'GLT)) []
      gt = Clause [WildP, RecP conName []] (NormalB (ConE 'GGT)) []
  return [main, lt, gt]

#if MIN_VERSION_template_haskell(2,17,0)
doUnqualifiedE = DoE Nothing
#else
doUnqualifiedE = DoE
#endif