import Text.Pandoc
import Text.Pandoc.JSON
import Text.Pandoc.Builder
import Control.Monad.State
import qualified Data.Map as M

main :: IO ()
main = toJSONFilter go
 where
  go (Just (Format "latex")) p                  = bottomUp runLatex p
  go _                       (Pandoc meta body) = Pandoc meta $ runOther body
  runLatex s@(Span (_id', cls, _props) val:xs)
    | "index" `elem` cls = concat
      [ [RawInline (Format "latex") "\\index{"]
      , val
      , [RawInline (Format "latex") "}"]
      , xs
      ]
    | otherwise = s
  runLatex x = x
  runOther :: [Block] -> [Block]
  runOther body = addIndex $ evalState (bottomUpM makeIds body) (1 :: Integer)
  addIndex :: [Block] -> [Block]
  addIndex bs = bs ++ toList idx
   where
    idx = header 1 (str "Index") <> definitionList (cat $ queryWith getIdx bs)
  makeIds s@(Span (_id', cls, props) val)
    | "index" `elem` cls = do
      curIdx <- get
      put $ curIdx + 1
      return $ Span ("index-id-no-" ++ show curIdx, cls, props) val
    | otherwise = return s
  makeIds x = return x
  getIdx :: Inline -> [(Inlines, Inlines)]
  getIdx (Span (id', cls, _props) val)
    | "index" `elem` cls
    = [(strong (fromList val), link ('#' : id') [] (str "idx"))]
    | otherwise
    = []
  getIdx _ = []
  cat = M.toList . M.map (return . plain) . M.fromListWith
    (\x acc -> acc <> str "," <> space <> x)