content: {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE DeriveDataTypeable #-}
  {-# LANGUAGE DeriveGeneric #-}
  {- |
     Module      : Text.Pandoc.Chunks
     Copyright   : Copyright (C) 2022-2024 John MacFarlane
     License     : GNU GPL, version 2 or above
  
     Maintainer  : John MacFarlane <jgm@berkeley.edu>
     Stability   : alpha
     Portability : portable
  
  Functions and types for splitting a Pandoc into subdocuments,
  e.g. for conversion into a set of HTML pages.
  -}
  module Text.Pandoc.Chunks
    ( Chunk(..)
    , ChunkedDoc(..)
    , PathTemplate(..)
    , splitIntoChunks
    , toTOCTree
    , tocToList
    , SecInfo(..)
    ) where
  
  import Text.Pandoc.Definition
  import Text.Pandoc.Shared (makeSections, stringify, inlineListToIdentifier,
                             tshow, uniqueIdent)
  import Text.Pandoc.Walk (Walkable(..), query)
  import Data.Aeson (FromJSON, ToJSON)
  import Data.Text (Text)
  import Text.Printf (printf)
  import Data.Maybe (fromMaybe, isNothing)
  import qualified Data.Map as M
  import qualified Data.Text as T
  import Data.String (IsString)
  import GHC.Generics (Generic)
  import Text.HTML.TagSoup (Tag (TagOpen), fromAttrib, parseTags)
  import Data.Tree (Tree(..))
  import Data.Data (Data)
  import Data.Typeable (Typeable)
  import qualified Data.Set as Set
  import Control.Monad.State
  
  -- | Split 'Pandoc' into 'Chunk's, e.g. for conversion into
  -- a set of HTML pages or EPUB chapters.
  splitIntoChunks :: PathTemplate -- ^ Template for filepath
                  -> Bool -- ^ Number sections
                  -> Maybe Int -- ^ Base heading level
                  -> Int -- ^ Chunk level -- level of section to split at
                  -> Pandoc
                  -> ChunkedDoc
  splitIntoChunks pathTemplate numberSections mbBaseLevel
                  chunklev (Pandoc meta blocks) =
     addNav .
     fixInternalReferences .
     walk rmNavAttrs $
     ChunkedDoc{ chunkedMeta = meta
               , chunkedChunks = chunks
               , chunkedTOC = tocTree }
   where
    tocTree = fixTOCTreePaths chunks $ toTOCTree sections
    chunks = makeChunks chunklev pathTemplate meta $ sections
    sections = ensureIds $ makeSections numberSections mbBaseLevel blocks
  
  -- The TOC won't work if we don't have unique identifiers for all sections.
  ensureIds :: [Block] -> [Block]
  ensureIds bs = evalState (walkM go bs) mempty
   where
     go :: Block -> State (Set.Set Text) Block
     go b@(Div (ident,"section":cls,kvs) bs'@(Header _ _ ils : _))
       | T.null ident
         = do ids <- get
              let newid = uniqueIdent mempty ils ids
              modify $ Set.insert newid
              pure $ Div (newid,"section":cls,kvs) bs'
       | otherwise
         = do modify $ Set.insert ident
              pure b
     go b = pure b
  
  -- | Add chunkNext, chunkPrev, chunkUp
  addNav :: ChunkedDoc -> ChunkedDoc
  addNav chunkedDoc =
    chunkedDoc{ chunkedChunks =
       addNext . addPrev . addUp $ chunkedChunks chunkedDoc }
  
  addUp :: [Chunk] -> [Chunk]
  addUp (c : d : ds)
    | chunkLevel c < chunkLevel d
      = c : addUp (d{ chunkUp = Just c } : ds)
    | chunkLevel c == chunkLevel d
      = c : addUp (d{ chunkUp = chunkUp c} : ds)
  addUp (c:cs) = c : addUp cs
  addUp [] = []
  
  addNext :: [Chunk] -> [Chunk]
  addNext cs = zipWith go cs (map Just (drop 1 cs) ++ [Nothing])
   where
    go c nxt = c{ chunkNext = nxt }
  
  addPrev :: [Chunk] -> [Chunk]
  addPrev cs = zipWith go cs (Nothing : map Just cs)
   where
    go c prev = c{ chunkPrev = prev }
  
  -- | Fix internal references so they point to the path of the chunk.
  fixInternalReferences :: ChunkedDoc -> ChunkedDoc
  fixInternalReferences chunkedDoc = walk fixInternalRefs chunkedDoc
   where
    fixInternalRefs :: Inline -> Inline
    fixInternalRefs il@(Link attr ils (src,tit))
      = case T.uncons src of
          Just ('#', ident) -> Link attr ils (src', tit)
            where src' = case M.lookup ident refMap of
                           Just chunk -> T.pack (chunkPath chunk) <> src
                           Nothing -> src
          _ -> il
    fixInternalRefs il = il
  
    refMap = foldr chunkToRefs mempty (chunkedChunks chunkedDoc)
  
    chunkToRefs chunk m =
      let idents = chunkId chunk : getIdents (chunkContents chunk)
      in  foldr (\ident -> M.insert ident chunk) m idents
  
    getIdents bs = query getBlockIdent bs ++ query getInlineIdent bs
  
    getBlockIdent :: Block -> [Text]
    getBlockIdent (Div (ident, _, _) _)
      | not (T.null ident) = [ident]
    getBlockIdent (Header _ (ident, _, _) _)
      | not (T.null ident) = [ident]
    getBlockIdent (Table (ident,_,_) _ _ _ _ _)
      | not (T.null ident) = [ident]
    getBlockIdent (RawBlock fmt raw)
      | isHtmlFormat fmt
      = foldr (\tag ->
                  case tag of
                    TagOpen{} ->
                      case fromAttrib "id" tag of
                        "" -> id
                        x  -> (x:)
                    _ -> id)
          [] (parseTags raw)
    getBlockIdent _ = []
  
    getInlineIdent :: Inline -> [Text]
    getInlineIdent (Span (ident, _, _) _)
      | not (T.null ident) = [ident]
    getInlineIdent (Link (ident, _, _) _ _)
      | not (T.null ident) = [ident]
    getInlineIdent (Image (ident, _, _) _ _)
      | not (T.null ident) = [ident]
    getInlineIdent (RawInline fmt raw)
      | isHtmlFormat fmt
      = foldr (\tag ->
                  case tag of
                    TagOpen{} ->
                      case fromAttrib "id" tag of
                        "" -> id
                        x  -> (x:)
                    _ -> id)
          [] (parseTags raw)
    getInlineIdent _ = []
  
    isHtmlFormat :: Format -> Bool
    isHtmlFormat (Format "html") = True
    isHtmlFormat (Format "html4") = True
    isHtmlFormat (Format "html5") = True
    isHtmlFormat _ = False
  
  
  makeChunks :: Int -> PathTemplate -> Meta -> [Block] -> [Chunk]
  makeChunks chunklev pathTemplate meta = secsToChunks 1
   where
    isChunkHeader :: Block -> Bool
    isChunkHeader (Div (_,"section":_,_) (Header n _ _:_)) = n <= chunklev
    isChunkHeader _ = False
  
    secsToChunks :: Int -> [Block] -> [Chunk]
    secsToChunks chunknum bs =
      case break isChunkHeader bs of
        ([], []) -> []
        ([], (d@(Div attr@(_,"section":_,_) (h@(Header lvl _ _) : bs')) : rest))
          | chunklev == lvl ->
            -- If the header is of the same level as chunks, create a chunk
            toChunk chunknum d :
              secsToChunks (chunknum + 1) rest
          | chunklev > lvl ->
            case break isChunkHeader bs' of
              (xs, ys) -> toChunk chunknum (Div attr (h:xs)) :
                            secsToChunks (chunknum + 1) (ys ++ rest)
        (xs, ys) -> toChunk chunknum
                       (Div ("",["preamble"],[]) xs) :
                      secsToChunks (chunknum + 1) ys
  
    toChunk :: Int -> Block -> Chunk
    toChunk chunknum
      (Div (divid,"section":classes,kvs) (h@(Header lvl _ ils) : bs)) =
      Chunk
        { chunkHeading = ils
        , chunkId = divid
        , chunkLevel = lvl
        , chunkNumber = chunknum
        , chunkSectionNumber = secnum
        , chunkPath = chunkpath
        , chunkUp = Nothing
        , chunkNext = Nothing
        , chunkPrev = Nothing
        , chunkUnlisted = "unlisted" `elem` classes
        , chunkContents =
           [Div (divid,"section":classes,kvs') (h : bs)]
        }
       where kvs' = kvs ++ [("nav-path", T.pack chunkpath)]
             secnum = lookup "number" kvs
             chunkpath = resolvePathTemplate pathTemplate chunknum
                          (stringify ils)
                          divid
                          (fromMaybe "" secnum)
    toChunk chunknum (Div ("",["preamble"],[]) bs) =
        Chunk
        { chunkHeading = docTitle meta
        , chunkId = chunkid
        , chunkLevel = 0
        , chunkNumber = chunknum
        , chunkSectionNumber = Nothing
        , chunkPath = chunkpath
        , chunkUp = Nothing
        , chunkPrev = Nothing
        , chunkNext = Nothing
        , chunkUnlisted = False
        , chunkContents = bs
        }
      where
        chunkpath = resolvePathTemplate pathTemplate chunknum
                          (stringify (docTitle meta))
                          chunkid
                          "0"
        chunkid = inlineListToIdentifier mempty (docTitle meta) <>
                        "-" <> tshow chunknum
    toChunk _ b = error $ "toChunk called on inappropriate block " <> show b
    -- should not happen
  
  
  -- Remove some attributes we added just to construct chunkNext etc.
  rmNavAttrs :: Block -> Block
  rmNavAttrs (Div (ident,classes,kvs) bs) =
    Div (ident,classes,filter (not . isNavAttr) kvs) bs
   where
    isNavAttr (k,_) = "nav-" `T.isPrefixOf` k
  rmNavAttrs b = b
  
  resolvePathTemplate :: PathTemplate
                      -> Int -- ^ Chunk number
                      -> Text -- ^ Stringified heading text
                      -> Text -- ^ Section identifier
                      -> Text -- ^ Section number
                      -> FilePath
  resolvePathTemplate (PathTemplate templ) chunknum headingText ident secnum =
    T.unpack .
    T.filter (\c -> c /= '/' && c /= '\\') .
    T.replace "%n" (T.pack $ printf "%03d" chunknum) .
    T.replace "%s" secnum .
    T.replace "%h" headingText .
    T.replace "%i" ident $
    templ
  
  -- | A 'PathTemplate' is a FilePath in which certain codes
  -- will be substituted with information from a 'Chunk'.
  -- @%n@ will be replaced with the chunk number
  -- (padded with leading 0s to 3 digits),
  -- @%s@ with the section number of the heading,
  -- @%h@ with the (stringified) heading text,
  -- @%i@ with the section identifier.
  -- For example, @"section-%s-%i.html"@ might be resolved to
  -- @"section-1.2-introduction.html"@.
  newtype PathTemplate =
    PathTemplate { unPathTemplate :: Text }
    deriving (Show, IsString, Data, Typeable, Generic, ToJSON, FromJSON)
  
  -- | A part of a document (typically a chapter or section, or
  -- the part of a section before its subsections).
  data Chunk =
    Chunk
    { chunkHeading :: [Inline]
    , chunkId :: Text
    , chunkLevel :: Int
    , chunkNumber :: Int
    , chunkSectionNumber :: Maybe Text
    , chunkPath :: FilePath
    , chunkUp :: Maybe Chunk
    , chunkPrev :: Maybe Chunk
    , chunkNext :: Maybe Chunk
    , chunkUnlisted :: Bool
    , chunkContents :: [Block]
    }
    deriving (Show, Eq, Generic)
  
  instance Walkable Inline Chunk where
    query f chunk = query f (chunkContents chunk)
    walk f chunk = chunk{ chunkContents = walk f (chunkContents chunk) }
    walkM f chunk = do
      contents <- walkM f (chunkContents chunk)
      return chunk{ chunkContents = contents }
  
  instance Walkable Block Chunk where
    query f chunk = query f (chunkContents chunk)
    walk f chunk = chunk{ chunkContents = walk f (chunkContents chunk) }
    walkM f chunk = do
      contents <- walkM f (chunkContents chunk)
      return chunk{ chunkContents = contents }
  
  -- | A 'Pandoc' broken into 'Chunk's for writing to separate files.
  data ChunkedDoc =
    ChunkedDoc
    { chunkedMeta :: Meta
    , chunkedTOC :: Tree SecInfo
    , chunkedChunks :: [Chunk]
    } deriving (Show, Eq, Generic)
  
  instance Walkable Inline ChunkedDoc where
    query f doc = query f (chunkedChunks doc) <> query f (chunkedMeta doc)
    walk f doc = doc{ chunkedMeta = walk f (chunkedMeta doc)
                    , chunkedChunks = walk f (chunkedChunks doc)
                    }
    walkM f doc = do
      meta' <- walkM f (chunkedMeta doc)
      chunks' <- walkM f (chunkedChunks doc)
      return $ doc{ chunkedMeta = meta'
                  , chunkedChunks = chunks' }
  
  instance Walkable Block ChunkedDoc where
    query f doc = query f (chunkedChunks doc) <> query f (chunkedMeta doc)
    walk f doc = doc{ chunkedMeta = walk f (chunkedMeta doc)
                    , chunkedChunks = walk f (chunkedChunks doc)
                    }
    walkM f doc = do
      meta' <- walkM f (chunkedMeta doc)
      chunks' <- walkM f (chunkedChunks doc)
      return $ doc{ chunkedMeta = meta'
                  , chunkedChunks = chunks' }
  
  -- | Data for a section in a hierarchical document.
  data SecInfo =
    SecInfo
    { secTitle :: [Inline]
    , secNumber :: Maybe Text
    , secId :: Text
    , secPath :: Text -- including fragment, e.g. chunk001.html#section-one
    , secLevel :: Int
    } deriving (Show, Eq, Generic)
  
  instance Walkable Inline SecInfo where
    query f sec = query f (secTitle sec)
    walk f sec = sec{ secTitle = walk f (secTitle sec) }
    walkM f sec = do
      st <- walkM f (secTitle sec)
      return sec{ secTitle = st }
  
  -- | Create tree of sections with titles, links, and numbers,
  -- in a form that can be turned into a table of contents.
  -- Presupposes that the '[Block]' is the output of 'makeSections'.
  toTOCTree :: [Block] -> Tree SecInfo
  toTOCTree =
    Node SecInfo{ secTitle = []
                , secNumber = Nothing
                , secId = ""
                , secPath = ""
                , secLevel = 0 } . foldr go []
   where
    go :: Block -> [Tree SecInfo] -> [Tree SecInfo]
    go (Div (ident,_,_) (Header lev (_,classes,kvs) ils : subsecs))
      | not (isNothing (lookup "number" kvs) && "unlisted" `elem` classes)
      = ((Node SecInfo{ secTitle = ils
                      , secNumber = lookup "number" kvs
                      , secId = ident
                      , secPath = ""
                      , secLevel = lev } (foldr go [] subsecs)) :)
    go (Div _ [d@Div{}]) = go d -- #8402
    go _ = id
  
  -- | Adjusts paths in the TOC tree generated by 'toTOCTree'
  -- to reflect division into Chunks.
  fixTOCTreePaths :: [Chunk] -> Tree SecInfo -> Tree SecInfo
  fixTOCTreePaths chunks = go ""
   where
    idMap = foldr (\chunk m ->
                     let ids = filter (not . T.null)
                                 (chunkId chunk :
                                    query getIds (chunkContents chunk))
                      in foldr (\i -> M.insert i (chunkPath chunk)) m ids)
                  mempty chunks
    getIds :: Block -> [Text]
    getIds (Div (i,"section":_,_) _) = [i]
    getIds _ = []
    go :: FilePath -> Tree SecInfo -> Tree SecInfo
    go fp (Node secinfo subtrees) =
      let newpath = M.lookup (secId secinfo) idMap
          fp' = fromMaybe fp newpath
          fragment = "#" <> secId secinfo
       in Node secinfo{ secPath = T.pack fp' <> fragment }
               (map (go fp') subtrees)
  
  -- | Creates a TOC link to the respective document section.
  tocEntryToLink :: Bool -> SecInfo -> [Inline]
  tocEntryToLink includeNumbers secinfo = headerLink
   where
    addNumber  = case secNumber secinfo of
                   Just num | includeNumbers
                          -> (Span ("",["toc-section-number"],[])
                                 [Str num] :) . (Space :)
                   _ -> id
    clean (Link _ xs _) = xs
    clean (Note _) = []
    clean x = [x]
    anchor = if T.null (secPath secinfo)
                then if T.null (secId secinfo)
                        then ""
                        else "#" <> secId secinfo
                else secPath secinfo
    headerText = addNumber $ walk (concatMap clean) (secTitle secinfo)
    headerLink = if T.null anchor
                    then headerText
                    else [Link ((if T.null (secId secinfo)
                                    then ""
                                    else "toc-" <> secId secinfo), [], [])
                           headerText (anchor, "")]
  
  -- | Generate a table of contents of the given depth.
  tocToList :: Bool -> Int -> Tree SecInfo -> Block
  tocToList includeNumbers tocDepth (Node _ subtrees) = BulletList (toItems subtrees)
   where
    toItems = map go . filter isBelowTocDepth
    isBelowTocDepth (Node sec _) = secLevel sec <= tocDepth
    go (Node secinfo xs) =
      Plain (tocEntryToLink includeNumbers secinfo) :
        case toItems xs of
          [] -> []
          ys -> [BulletList ys]
  