{-
	Copyright (c) 2013, Alex Cole
	
	This Source Code Form is subject to the terms of the Mozilla Public License,
	v. 2.0. If a copy of the MPL was not distributed with this file, You can
	obtain one at http://mozilla.org/MPL/2.0/.
-}

{-# LANGUAGE TemplateHaskell        #-}

module Data.Types.Reorder.TH (
		reorderable,
		reorderableGroup
	) where

import Data.Types.Reorder.Base
import Data.Types.Reorder.Quoter

import Language.Haskell.TH
import Language.Haskell.TH.Syntax
import Language.Haskell.TH.Quote

import Data.Maybe (mapMaybe)
import Data.List (nub)

import Language.Haskell.Exts.Parser
import Language.Haskell.Exts.Extension
import qualified Language.Haskell.Exts.Syntax as Hs

import Language.Haskell.Meta.Parse
import Language.Haskell.Meta.Syntax.Translate
-- import qualified Language.Haskell.Exts.Annotated.Fixity as Fix

-- Type class on which to store types that are declared using "reorderable"
-- instead of "reorderableGroup", i.e. types that need interaction code
-- generating for new reorderable types.
class ReorderableGlobalGroup a

--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

-- | Takes a type name, i.e. @reorderable ''IP@.
reorderable :: Name -> Q [Dec]
reorderable new
	| new == ''ReorderableEnd = return [] -- Deny this type.
	| otherwise = generateCoreCode new `cat` (
		getAllGlobals >>= (\ old -> concatMapM (generatePairCode new) old `cat` (
			if new `elem` old
				then return []
				else return [InstanceD [] (AppT (ConT ''ReorderableGlobalGroup) (ConT new)) []])))
		where
			getAllGlobals = reify ''ReorderableGlobalGroup >>= (\ (ClassI _ instances) -> return $
				map (\ (InstanceD _ (AppT _ (ConT t)) _) -> t) instances)

-- | Takes a list of names, and generates pair instances for them all.
--  
-- > reorderableGroup [''IP, ''TIC, ''RT]
--  
reorderableGroup new' = concatMapM generateCoreCode new `cat` genPairs new
	where
		-- Filter all types that are NOT "ReorderableEnd", and remove any
		-- duplicates.  The code above will not generate any code in the case of
		-- an empty list, so don't check for that case explicitly as it is by
		-- far an uncommon case, thus a waste of time in the common case.
		new = nub $ filter (not . (==) ''ReorderableEnd) new'
		-- For each element in the array, generate it paired with all LATER
		-- items.  We know that the array has all duplicates removed, so as long
		-- as we ensure that when generating a pair we generated in both
		-- directions (i.e. "TypeOrder a b" and "TypeOrder b a"), this will
		-- generate all combinations from the array.
		genPairs []       = return []
		genPairs (x : []) = return []
		genPairs (x : xs) = concatMapM (generatePairCode x) xs `cat` genPairs xs

--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

-- | This code generates the code required to compare two reorderable types.  It
-- generates two instances for every valid combination of reorderable types, it
-- then stores all the generated combinations in case one comes up again.  E.g:
--  
-- > reorderableGroup [''Type1, ''Type2, ''Type3]
-- > reorderableGroup [''Type1, ''Type2, ''Type4]
--  
-- This will not generate a pair for @(''Type3, ''Type4)@, and will also
-- remember that it has already generated a pair for @(''Type1, ''Type2)@ when
-- the second group comes to be compiled.
generatePairCode :: Name -> Name -> Q [Dec]
generatePairCode new old = getAllPairs >>=
		(\ p -> return $ generatePairCode' new old p ++ generatePairCode' old new p)
	where
		-- Find all existing type pairs.
		getAllPairs = reify ''TypeOrder >>= (\ (FamilyI _ instances) -> return $
			mapMaybe exTypes instances)
		-- Extract the two types.
		exTypes (TySynInstD _ [ConT a, ConT b] _) = Just (a, b)
		-- exTypes (TySynInstD _ [] _) = error "0"
		-- exTypes (TySynInstD _ [_] _) = error "1"
		-- exTypes (TySynInstD a xs b) = error $ show a ++ ", " ++ show xs ++ ", " ++ show b
		exTypes _ = Nothing
		-- Generate one pair one way around.
		generatePairCode' new old pairs =
			if ((new, old) `elem` pairs)
				then []
				else if (nameBase new < nameBase old)
					then [_TypeOrder [ConT old, ConT new] _FirstHigher]
					else [_TypeOrder [ConT old, ConT new] _FirstLower ]
		-- Useful type constructors.
		_FirstLower  = ConT ''TypeOrder_Lower
		_FirstHigher = ConT ''TypeOrder_Higher
		_TypeOrder   = TySynInstD ''TypeOrder

-- | This code generates all the instances of classes required by EVERY
-- reorderable type, irrespective of what other types there are.
generateCoreCode :: Name -> Q [Dec]
generateCoreCode new = getAllGenerators >>= concatMapM (compile new' . replaceName) . groupCode
	where
		-- The name of the thing being compiled.
		name = nameBase new
		new' = ConT new
		
		-- Replace "???" with the actual type name being currently compiled.
		replaceName' [] = []
		replaceName' ('?' : '?' : '?' : rest) = name ++ replaceName' rest
		replaceName' (r : rest) = r : replaceName' rest
		replaceName (a, b) = (a, replaceName' (unlines b))
		
		-- Find the code for all the available generators in the current code.
		getAllGenerators = reify ''ReorderableInstance >>= (\ (ClassI _ instances) -> return $ map
				(\ (InstanceD _ (AppT (AppT (AppT _ (ConT t)) (LitT (NumTyLit n))) (LitT (StrTyLit s))) _) ->
					-- "t" is the parent type's name.
					-- "n" is the allocated number.
					-- "s" is the code.
					(t, n, s)
				) instances
			)
		
		-- The base compiler
		compile new (t, s) = isInstance t [new] >>=
			-- Check if there is already an instance of this code for this type.
			(\ b -> return $ if b
				then [] -- Already done.
				else case compileDecls s of -- Compile the code.
					Left  err -> error err -- Return the error.
					-- Compiled OK, add on a new instance declaration too.
					Right dec -> InstanceD [] (AppT (ConT t) new) [] : dec)
		
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
-- This is the compiler code.  A lot of this is ripped from haskell-src-meta,
-- other bits are written by me.  I did submit a pull request to get many of
-- these updates integrated in to the original library, but they haven't been
-- integrated, so I'm putting them here instead for compatibility.  In short,
-- haskell-src-meta doesn't parse "type instance", but "haskell-src-exts", on
-- which it is based, does.  As we need those, we intercept the translation,
-- look for any type family instances, and handle those here.  Any other
-- declarations we pass on to be handled as normal.  Unfortunately, this needs a
-- lot of code as we need to intercept ALL intermediate stages to get our new
-- code in to the right spot and pass the correct extensions.
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
		
		-- Note: Despite the confusingly indented comment block above, this code
		-- is still within the "where" clause of "generateCoreCode".
		compileDecls :: String -> Either String [Dec]
		compileDecls = either Left (Right . toDecs2 . moduleDecls)
		  . parseResultToEither . parseModuleWithMode parseMode
		
		parseMode = ParseMode {
			parseFilename = []            ,
			baseLanguage = Haskell2010,
			extensions = map EnableExtension parseExts,
			ignoreLinePragmas = False     ,
			ignoreLanguagePragmas = False ,
			fixities = Nothing    }
		
		parseExts = [
			EmptyDataDecls           ,
			TypeFamilies             ,
			PostfixOperators         ,
			UnicodeSyntax            ,
			FlexibleContexts         ,
			PatternSignatures        ,
			RankNTypes               ,
			MultiParamTypeClasses    ]

-----------------------------------------------------------------------------

class ToDecs2 a where toDecs2 :: a -> [Dec]

-- * ToDecs InstDecl
instance ToDecs2 Hs.InstDecl where
  toDecs2 (Hs.InsDecl decl) = toDecs2 decl
  toDecs2 d              = todo "toDec2" d

-- * ToDecs HsDecl HsBinds

instance ToDecs2 Hs.Decl where
	toDecs2 a@(Hs.TypeSig _ ns t) = xs
		where
			xs = fmap (flip SigD (fixForall $ toType t) . toName) ns
	
	toDecs2 (Hs.InfixDecl _ assoc fixity ops) =
			map (\op -> InfixD (Fixity fixity dir) (toName op)) ops
		where
			dir = case assoc of
				Hs.AssocNone -> InfixN
				Hs.AssocLeft -> InfixL
				Hs.AssocRight -> InfixR
	
	toDecs2 a = [toDec2 a]

instance ToDecs2 a => ToDecs2 [a] where
	toDecs2 a = concatMap toDecs2 a

instance ToDecs2 Hs.Binds where
	toDecs2 (Hs.BDecls ds) = toDecs2 ds

hsTypeToNameAndTypes2 :: Hs.Type -> (Name, [Type])
hsTypeToNameAndTypes2 = getTypes [] . toType
	where
		getTypes ts (AppT a b)      = getTypes (b : ts) a
		getTypes ts (ForallT _ _ a) = getTypes ts a
		getTypes ts (ConT con)      = (con, ts)
		getTypes ts (VarT con)      = (con, ts)
		getTypes ts (PromotedT con) = (con, ts)
		getTypes _ a = nonsense "hsTypeToNameAndTypes2" "Unexpected type" a

-- Had family declarations, but no instances.
toDec2 (Hs.TypeInsDecl _ t0 t1) =
		TySynInstD n ts (toType t1)
	where
		(n, ts) = hsTypeToNameAndTypes2 t0
	
toDec2 (Hs.DataInsDecl _ Hs.DataType t qcds qns) =
		DataInstD [] n ts (fmap qualConDeclToCon qcds) (fmap (toName . fst) qns)
	where
		(n, ts) = hsTypeToNameAndTypes2 t
	
toDec2 (Hs.DataInsDecl _ Hs.NewType t (q:_) qns) =
		NewtypeInstD [] n ts (qualConDeclToCon q) (fmap (toName . fst) qns)
	where
		(n, ts) = hsTypeToNameAndTypes2 t
	
-- Default to the standard translation.
toDec2 a = toDec a

--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

-- | Do `concatMap' with a function returning a monad.
concatMapM :: (Monad m) => (a -> m [b]) -> [a] -> m [b]
concatMapM f a = mapM f a >>= return . concat

-- | Concatenate two arrays in a monad.
cat :: (Monad m) => m [a] -> m [a] -> m [a]
cat a0 a1 = do
	a0' <- a0
	a1' <- a1
	return $ a0' ++ a1'

mapMaybes :: (a -> Maybe b) -> [a] -> ([b], [a])
mapMaybes f [] = ([], [])
mapMaybes f (x : xs) = case f x of
		Nothing -> (l, x : r)
		Just x' -> (x' : l, r)
	where
		(l, r) = mapMaybes f xs

-- Intra-Splice order is unimportant.
groupCode [] = []
groupCode ((t, _, b) : xs) = (t, b : ms) : groupCode ns
	where
		sameGroup (t', _, b')
			| t == t'   = Just b'
			| otherwise = Nothing
		(ms, ns) = mapMaybes sameGroup xs