{-# LANGUAGE TemplateHaskell #-}

-- |
-- Module: Options.Help
-- License: MIT
module Options.Help
	( addHelpFlags
	, checkHelpFlag
	, helpFor
	, HelpFlag(..)
	) where

import           Control.Monad.Writer
import           Data.List (intercalate, partition, stripPrefix)
import           Data.Maybe (isNothing, listToMaybe, maybeToList)
import qualified Data.Set as Set
import qualified Data.Map as Map

import           Language.Haskell.TH (location, loc_package, loc_module)

import           Options.Types

data HelpFlag = HelpSummary | HelpAll | HelpGroup String
	deriving (Eq, Show)

addHelpFlags :: OptionDefinitions a -> OptionDefinitions a
addHelpFlags (OptionDefinitions opts subcmds) = OptionDefinitions withHelp subcmdsWithHelp where
	shortFlags = Set.fromList $ do
		opt <- opts
		optionInfoShortFlags opt
	longFlags = Set.fromList $ do
		opt <- opts
		optionInfoLongFlags opt
	
	withHelp = optHelpSummary ++ optsGroupHelp ++ opts
	
	groupHelp = GroupInfo
		{ groupInfoName = "all"
		, groupInfoTitle = "Help Options"
		, groupInfoDescription = "Show all help options."
		}
	
	optSummary = OptionInfo
		{ optionInfoKey = keyFor "optHelpSummary"
		, optionInfoShortFlags = []
		, optionInfoLongFlags = []
		, optionInfoDefault = "false"
		, optionInfoUnary = True
		, optionInfoDescription = "Show option summary."
		, optionInfoGroup = Just groupHelp
		}
	
	optGroupHelp group flag = OptionInfo
		{ optionInfoKey = keyFor "optHelpGroup" ++ ":" ++ groupInfoName group
		, optionInfoShortFlags = []
		, optionInfoLongFlags = [flag]
		, optionInfoDefault = "false"
		, optionInfoUnary = True
		, optionInfoDescription = groupInfoDescription group
		, optionInfoGroup = Just groupHelp
		}
	
	optHelpSummary = if Set.member 'h' shortFlags
		then if Set.member "help" longFlags
			then []
			else [optSummary
				{ optionInfoLongFlags = ["help"]
				}]
		else if Set.member "help" longFlags
			then [optSummary
				{ optionInfoShortFlags = ['h']
				}]
			else [optSummary
				{ optionInfoShortFlags = ['h']
				, optionInfoLongFlags = ["help"]
				}]
	
	optsGroupHelp = do
		let (groupsAndOpts, _) = uniqueGroupInfos opts
		let groups = [g | (g, _) <- groupsAndOpts]
		group <- (groupHelp : groups)
		let flag = "help-" ++ groupInfoName group
		if Set.member flag longFlags
			then []
			else [optGroupHelp group flag]
	
	subcmdsWithHelp = do
		(subcmdName, subcmdOpts) <- subcmds
		let subcmdLongFlags = Set.fromList $ do
			opt <- subcmdOpts ++ optsGroupHelp
			optionInfoLongFlags opt
		
		let (groupsAndOpts, _) = uniqueGroupInfos subcmdOpts
		let groups = [g | (g, _) <- groupsAndOpts]
		let newOpts = do
			group <- groups
			let flag = "help-" ++ groupInfoName group
			if Set.member flag (Set.union longFlags subcmdLongFlags)
				then []
				else [optGroupHelp group flag]
		return (subcmdName, newOpts ++ subcmdOpts)

checkHelpFlag :: TokensFor a -> Maybe HelpFlag
checkHelpFlag (TokensFor tokens _) = flag where
	flag = listToMaybe helpKeys
	helpKeys = do
		(k, _) <- tokens
		if k == keySummary
			then return HelpSummary
			else if k == keyAll
				then return HelpAll
				else do
					groupName <- maybeToList (stripPrefix keyGroupPrefix k)
					return (HelpGroup groupName)
	keySummary = keyFor "optHelpSummary"
	keyAll = keyFor "optHelpGroup:all"
	keyGroupPrefix = keyFor "optHelpGroup:"

helpFor :: HelpFlag -> OptionDefinitions a -> Maybe String -> String
helpFor flag defs subcmd = case flag of
	HelpSummary -> execWriter (showHelpSummary defs subcmd)
	HelpAll -> execWriter (showHelpAll defs subcmd)
	HelpGroup groupName -> execWriter (showHelpOneGroup defs groupName subcmd)

showOptionHelp :: OptionInfo -> Writer String ()
showOptionHelp info = do
	let safeHead xs = case xs of
		[] -> []
		(x:_) -> [x]
	let shorts = optionInfoShortFlags info
	let longs = optionInfoLongFlags info
	let optStrings = map (\x -> ['-', x]) (safeHead shorts) ++ map (\x -> "--" ++ x) (safeHead longs)
	unless (null optStrings) $ do
		let optStringCsv = intercalate ", " optStrings
		tell "  "
		tell optStringCsv
		
		let desc = optionInfoDescription info
		unless (null desc) $ do
			if length optStringCsv > 27
				then do
					tell "\n"
					tell "    "
					tell (optionInfoDescription info)
				else do
					tell (replicate (28 - length optStringCsv) ' ')
					tell (optionInfoDescription info)
		
		tell "\n"

showHelpSummary :: OptionDefinitions a -> Maybe String -> Writer String ()
showHelpSummary (OptionDefinitions mainOpts subcmds) subcmd = do
	let subcmdOptions = do
		subcmdName <- subcmd
		opts <- lookup subcmdName subcmds
		return (subcmdName, opts)
	
	let (groupInfos, ungroupedMainOptions) = uniqueGroupInfos mainOpts
	
	-- Always print --help group
	let hasHelp = filter (\(g,_) -> groupInfoName g == "all") groupInfos
	forM_ hasHelp showHelpGroup
	
	tell "Application Options:\n"
	forM_ ungroupedMainOptions showOptionHelp
	unless (null subcmds) (tell "\n")
	
	case subcmdOptions of
		Nothing -> unless (null subcmds) $ do
			tell "Subcommands:\n"
			forM_ subcmds $ \(subcmdName, _) -> do
				tell "  "
				tell subcmdName
				-- TODO: subcommand help description
				tell "\n"
			tell "\n"
		Just (n, subOpts) -> do
			-- TODO: subcommand description
			-- TODO: handle grouped options in subcommands?
			tell ("Options for subcommand " ++ show n ++ ":\n")
			forM_ subOpts showOptionHelp
			tell "\n"

showHelpAll :: OptionDefinitions a -> Maybe String -> Writer String ()
showHelpAll (OptionDefinitions mainOpts subcmds) subcmd = do
	let subcmdOptions = do
		subcmdName <- subcmd
		opts <- lookup subcmdName subcmds
		return (subcmdName, opts)
	
	let (groupInfos, ungroupedMainOptions) = uniqueGroupInfos mainOpts
	
	-- Always print --help group first, if present
	let (hasHelp, noHelp) = partition (\(g,_) -> groupInfoName g == "all") groupInfos
	forM_ hasHelp showHelpGroup
	forM_ noHelp showHelpGroup
	
	tell "Application Options:\n"
	forM_ ungroupedMainOptions showOptionHelp
	unless (null subcmds) (tell "\n")
	
	case subcmdOptions of
		Nothing -> forM_ subcmds $ \(subcmdName, subcmdOpts) -> do
			-- no subcommand description
			tell ("Options for subcommand " ++ show subcmdName ++ ":\n")
			forM_ subcmdOpts showOptionHelp
			tell "\n"
		Just (n, subOpts) -> do
			-- TODO: subcommand description
			-- TODO: handle grouped options in subcommands?
			tell ("Options for subcommand " ++ show n ++ ":\n")
			forM_ subOpts showOptionHelp
			tell "\n"

showHelpGroup :: (GroupInfo, [OptionInfo]) -> Writer String ()
showHelpGroup (groupInfo, opts) = do
	tell (groupInfoTitle groupInfo ++ ":\n")
	forM_ opts showOptionHelp
	tell "\n"

showHelpOneGroup :: OptionDefinitions a -> String -> Maybe String -> Writer String ()
showHelpOneGroup (OptionDefinitions mainOpts subcmds) groupName subcmd = do
	let opts = case subcmd of
		Nothing -> mainOpts
		Just n -> case lookup n subcmds of
			Just infos -> mainOpts ++ infos -- both
			Nothing -> mainOpts
	let (groupInfos, _) = uniqueGroupInfos opts
	
	-- Always print --help group
	let group = filter (\(g,_) -> groupInfoName g == groupName) groupInfos
	forM_ group showHelpGroup

keyFor :: String -> String
keyFor fieldName = this_pkg ++ ":" ++ this_mod ++ ":" ++ fieldName where
	(this_pkg, this_mod) = $(do
		loc <- location
		let pkg = loc_package loc
		let mod' = loc_module loc
		[| (pkg, mod') |])

uniqueGroupInfos :: [OptionInfo] -> ([(GroupInfo, [OptionInfo])], [OptionInfo])
uniqueGroupInfos allOptions = (Map.elems infoMap, ungroupedOptions) where
	infoMap = Map.fromListWith merge $ do
		opt <- allOptions
		case optionInfoGroup opt of
			Nothing -> []
			Just g -> [(groupInfoName g, (g, [opt]))]
	merge (g, opts1) (_, opts2) = (g, opts2 ++ opts1)
	ungroupedOptions = [o | o <- allOptions, isNothing (optionInfoGroup o)]
