{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
module Hdoc.DTC.Analyze.Collect where

import Control.Applicative (Applicative(..), liftA2)
import Control.Monad
import Data.Bool
import Data.Default.Class (Default(..))
import Data.Foldable (Foldable(..), concat)
import Data.Function (($), (.))
import Data.Functor ((<$>))
import Data.Map.Strict (Map)
import Data.Maybe (Maybe(..), fromMaybe, maybe)
import Data.Monoid (Monoid(..))
import Data.Semigroup (Semigroup(..))
import Data.Sequence (Seq)
import Data.TreeSeq.Strict (Tree(..))
import Text.Show (Show(..))
import qualified Control.Monad.Trans.Reader as R
import qualified Data.HashMap.Strict as HM
import qualified Data.Map.Strict as Map
import qualified Data.Sequence as Seq
import qualified Data.Text.Lazy as TL
import qualified Data.TreeSeq.Strict as TS
import qualified Data.TreeMap.Strict as TM

import qualified Hdoc.TCT.Cell as TCT
import Hdoc.DTC.Document as DTC
import qualified Hdoc.XML as XML

-- * Type 'Reader'
newtype Reader = Reader
 { reader_section :: Section
 }
instance Default Reader where
	def = Reader
	 { reader_section = def
	 }

-- * Type 'All'
data All = All
 { all_figure     :: HM.HashMap TL.Text (Map XML.Pos (Maybe Title))
 , all_grades     :: HM.HashMap Ident (Seq [Grade])
 , all_irefs      :: TM.TreeMap Word (Seq Section)
 , all_judges     :: HM.HashMap Ident (Seq Judges)
 , all_judgesRef  :: HM.HashMap Ident (Seq ((TCT.Location, XML.Pos), Section))
 , all_notes      :: TS.Trees (Seq [Para])
 , all_reference  :: HM.HashMap Ident (Seq Reference)
 , all_ref        :: HM.HashMap Ident (Seq ((TCT.Location, XML.Pos), Section))
 , all_pageRef    :: HM.HashMap PathPage (Seq ((TCT.Location, XML.Pos), Section))
 , all_section    :: HM.HashMap Ident (Seq Section)
 , all_at         :: HM.HashMap Ident (Seq ((TCT.Location, XML.Pos), Section))
 , all_tag        :: HM.HashMap Ident (Seq ((TCT.Location, XML.Pos), Section))
 } deriving (Show)
instance Default All where
	def = All
	 { all_figure    = def
	 , all_grades    = def
	 , all_irefs     = def
	 , all_judges    = def
	 , all_judgesRef = def
	 , all_notes     = def
	 , all_reference = def
	 , all_ref       = def
	 , all_pageRef   = def
	 , all_section   = def
	 , all_at        = def
	 , all_tag       = def
	 }
instance Semigroup All where
	x<>y = All
	 { all_figure    = hm_union all_figure
	 , all_grades    = hm_union all_grades
	 , all_irefs     = tm_union all_irefs
	 , all_judges    = hm_union all_judges
	 , all_judgesRef = hm_union all_judgesRef
	 , all_notes     = ts_union (all_notes x) (all_notes y)
	 , all_reference = hm_union all_reference
	 , all_ref       = hm_union all_ref
	 , all_pageRef   = hm_union all_pageRef
	 , all_section   = hm_union all_section
	 , all_at        = hm_union all_at
	 , all_tag       = hm_union all_tag
	 } where
		hm_union get = HM.unionWith (<>) (get x) (get y)
		tm_union get = TM.union     (<>) (get x) (get y)
		ts_union :: TS.Trees (Seq [Para]) -> TS.Trees (Seq [Para]) -> TS.Trees (Seq [Para])
		ts_union sx sy = lx <> Seq.singleton union <> ry
			where
			filter   = null . TS.subTrees
			union    = TS.tree0 $ TS.unTree =<< (rx <> ly)
			(rx, lx) = Seq.spanr filter sx
			(ly, ry) = Seq.spanl filter sy
instance Monoid All where
	mempty  = def
	mappend = (<>)
instance Semigroup (R.Reader Reader All) where
	(<>) = liftA2 (<>)
instance Monoid (R.Reader Reader All) where
	mempty  = pure def
	mappend = (<>)

-- * Class 'Collect'
class Collect a where
	collect :: a -> R.Reader Reader All
instance Collect Document where
	collect Document{..} =
		foldMap collect document_head <>
		collect document_body
		-- R.local (\ro -> ro{reader_section = head_section head}) $
			{- FIXME: see whether to handle that or not
			return def
			 { all_section =
				HM.fromListWith (<>) $
				(\t -> (t, pure $ Left head))
				 <$> about_titles (head_about head)
			 } <>
			-}
instance Collect Head where
	collect Head{..} =
		collect $ Seq.singleton $
			TS.Tree (BodySection head_section) head_body
instance Collect (TS.Trees BodyNode) where
	collect = foldMap $ \(TS.Tree b bs) ->
		case b of
		 BodyBlock blk ->
			collect blk
		 BodySection section@Section{..} ->
			R.local (\ro -> ro{reader_section = section}) $
				((\all -> all{all_notes = pure $ TS.Tree Seq.empty $ all_notes all}) <$>) $
				return def
				 { all_section =
					let titles = (<$> about_titles section_about) $ \section_title ->
						Alias
						 { alias_attrs = section_attrs
						 , alias_title = section_title
						 } in
					HM.fromListWith (<>) $ concat $
					(\Alias{..} -> maybe [] (pure . (, pure section)) $ attrs_id alias_attrs)
					 <$> (titles <> about_aliases section_about)
				 } <>
				foldMap collect (about_titles section_about) <>
				collect bs
instance Collect Block where
	collect = \case
	 BlockPara p     -> collect p
	 BlockBreak{}    -> pure def
	 BlockToC{}      -> pure def
	 BlockToF{}      -> pure def
	 BlockAside{..}  -> foldMap collect blocks
	 BlockIndex{..}  -> pure def
	 BlockFigure{..} ->
		return def
		 { all_figure = HM.singleton type_ (Map.singleton posXML mayTitle)
		 } <> foldMap collect paras
	 BlockReferences{..} ->
		return def
		 { all_reference=
			HM.fromListWith (<>) $
			(<$> refs) $
			\ref@DTC.Reference{..} -> (reference_id, pure ref)
		 } <> foldMap collect refs
	 BlockGrades{attrs, ..} ->
		return def{all_grades = HM.singleton (fromMaybe "" $ attrs_id attrs) $ pure scale}
	 BlockJudges judges@Judges{judges_attrs, ..} ->
		return def{all_judges = HM.singleton (fromMaybe "" $ attrs_id judges_attrs) $ pure judges}
instance Collect Para where
	collect = \case
	 ParaItem item -> collect item
	 ParaItems{..} -> foldMap collect items
instance Collect ParaItem where
	collect = \case
	 ParaPlain p    -> collect p
	 ParaArtwork{}  -> return def
	 ParaQuote{..}  -> foldMap collect paras
	 ParaComment{}  -> return def
	 ParaOL items   -> foldMap collect items
	 ParaUL items   -> foldMap (foldMap collect) items
	 ParaJudgment j -> collect j
instance Collect Judgment where
	collect Judgment{..} = do
		Reader{..} <- R.ask
		return def{all_judgesRef = HM.singleton judgment_judgesId $ pure ((judgment_locTCT, judgment_posXML), reader_section) }
instance Collect ListItem where
	collect ListItem{..} = foldMap collect paras
instance Collect Title where
	collect (Title t) = collect t
instance Collect Plain where
	collect = foldMap collect
instance Collect (Tree PlainNode) where
	collect (Tree n ts) =
		case n of
		 PlainBreak    -> return def
		 PlainText{}   -> return def
		 PlainGroup    -> collect ts
		 PlainB        -> collect ts
		 PlainCode     -> collect ts
		 PlainDel      -> collect ts
		 PlainI        -> collect ts
		 PlainSpan{}   -> collect ts
		 PlainSub      -> collect ts
		 PlainSup      -> collect ts
		 PlainSC       -> collect ts
		 PlainU        -> collect ts
		 PlainNote{..} ->
			return def
			 { all_notes = pure $ TS.tree0 $ pure note_paras
			 } <> foldMap collect note_paras
		 PlainQ        -> collect ts
		 PlainEref{}   -> collect ts
		 PlainIref{..} -> do
			Reader{..} <- R.ask
			case pathFromWords iref_term of
			 Nothing -> collect ts
			 Just path ->
				return def
				 { all_irefs = TM.singleton path $ pure reader_section
				 } <> collect ts
		 PlainAt{..}
		  | at_back -> return def
		  | otherwise -> do
			Reader{..} <- R.ask
			return def
			 { all_at = HM.singleton at_ident $ pure ((at_locTCT, at_posXML), reader_section) }
		 PlainTag{..}
		  | tag_back -> return def
		  | otherwise -> do
			Reader{..} <- R.ask
			return def
			 { all_tag = HM.singleton tag_ident $ pure ((tag_locTCT, tag_posXML), reader_section) }
		 PlainRef{..} -> do
			Reader{..} <- R.ask
			return def
			 { all_ref = HM.singleton ref_ident $ pure ((ref_locTCT, ref_posXML), reader_section)
			 } <> collect ts
		 PlainPageRef{..} -> do
			Reader{..} <- R.ask
			return def
			 { all_pageRef = HM.singleton pageRef_path $ pure ((pageRef_locTCT, pageRef_posXML), reader_section)
			 } <> collect ts
instance Collect Reference where
	collect Reference{..} =
		collect reference_about
instance Collect About where
	collect About{..} =
		foldMap collect about_titles <>
		foldMap collect about_description