-- compile with: ghc -Wall -O2 --make footnotes.hs

module Main (main, doFiles, putFootnotes) where

import System.Environment (getArgs)
import System.Console.GetOpt

import Control.Exception (Exception, catchJust)
import Control.Monad (foldM, guard)
import Data.Array.IArray (Array, array, (!))
import qualified Data.ByteString.Lazy.Char8 as B
import qualified Data.IntMap as M
import Data.List (isSuffixOf)
import Data.Maybe (mapMaybe)

data Flag = Appearance deriving (Show, Eq)
options :: [OptDescr Flag]
options = [Option ['a'] ["reorder-by-appearance"] (NoArg Appearance)
                      "Reorder footnotes by order of appearence"]

-- main: process the command-line options for footnotes.
main :: IO ()
main = do
  args <- getArgs
  case getOpt Permute options args of
    (opts, files, []) ->
        doFiles files $ Appearance `elem` opts
    (_, _, errors)       ->
        ioError $ userError $ concat errors ++ usageInfo header options
            where header = "Usage: footnotes [OPTION...] [files...]"

-- doFiles: read input either from given files or from stdin,
-- if none were given, and print the output to stdout.
doFiles :: [FilePath] -> Bool -> IO ()
doFiles []    useBody = B.getContents >>= putFootnotes useBody
doFiles files useBody = mapM_ doFile files where
    doFile file = B.readFile file >>= putFootnotes useBody

type RefsMap      = M.IntMap Int     -- maps old to new refs.
data RefsMapState = RMS !RefsMap Int -- keeps track of map and counter.
type FootsMap     = Array Int (IO ()) -- sorted footnotes.

sepLine :: B.ByteString         -- the seperator line.
sepLine = B.pack "@footnote:"

emptyRMS :: RefsMapState        -- an empty RefsMap and its counter.
emptyRMS = RMS M.empty 0

-- putFootnotes: print s with renumbered footnotes.
putFootnotes :: Bool -> B.ByteString -> IO ()
putFootnotes True  s = do       -- by order of appearance in body.
  let (body, _ : foots) = break (== sepLine) $ B.lines s
  (RMS m n) <- foldM (putLineSubRefs putNewRef) emptyRMS body
  B.putStrLn sepLine
  let fm = array (1, n) $ mapMaybe (subRef m) foots
  mapM_ (lookupPutFoot fm) [1..n]

putFootnotes False s = do       -- by order of the footnotes.
  let (_ : foots) = dropWhile (/= sepLine) $ B.lines s
  let st@(RMS m _) = foldl insertFootRef emptyRMS foots
  let (body, _ : foots') = break (== sepLine) $ B.lines s
  mapM_ (putLineSubRefs putUnknownRef st) body
  B.putStrLn sepLine
  mapM_ (\l -> maybe (B.putStrLn l) snd $ subRef m l) foots'

-- putLineSubRefs: print body line l with substituted references,
-- evaluate doUnknownRef for references not in m, return new state.
putLineSubRefs :: (Int -> RefsMapState -> IO RefsMapState)
               -> RefsMapState -> B.ByteString -> IO RefsMapState
putLineSubRefs doUnknownRef st@(RMS m _) l =
    case B.elemIndex '[' l of
      Nothing -> B.putStrLn l >> return st
      Just i  -> do
        let (before, after) = B.splitAt (i + 1) l
        B.putStr before
        let parse = do (n, rest) <- B.readInt after
                       (']',  _) <- B.uncons rest
                       return (n, rest)
        case parse of
          Nothing        ->
              putLineSubRefs doUnknownRef st after
          Just (n, rest) ->
              case M.lookup n m of
                Just newN -> do putStr $ show newN
                                putLineSubRefs doUnknownRef st rest
                Nothing   -> do newSt <- doUnknownRef n st
                                putLineSubRefs doUnknownRef newSt rest

-- insertRef: return new RefsMapState with oldN inserted.
insertRef :: RefsMapState -> Int -> RefsMapState
insertRef (RMS m n) oldN = RMS (M.insert oldN n' m) n' where n' = n+1

-- subRef: substitute the reference in footnote line l according to m,
-- return its number and an IO action for printing the new l.
subRef :: RefsMap -> B.ByteString -> Maybe (Int, IO ())
subRef m l = do (oldN, rest) <- parseFoot l
                n <- M.lookup oldN m
                return (n, putStr ('[' : show n) >> B.putStrLn rest)

-- parseFoot: return number and the rest after it from footnote line l.
parseFoot :: B.ByteString -> Maybe (Int, B.ByteString)
parseFoot l = do ('[', l') <- B.uncons l
                 (n, rest) <- B.readInt l'
                 (']', _)  <- B.uncons rest
                 return (n, rest)

-- putNewRef :: print new ref for oldN and return new RefsMapState.
putNewRef :: Int -> RefsMapState -> IO RefsMapState
putNewRef oldN st = putStr (show n) >> return newSt 
    where newSt@(RMS _ n) = insertRef st oldN

-- lookupPutFoot: return the IO action for footnote number n in fm.
lookupPutFoot :: FootsMap -> Int -> IO ()
lookupPutFoot fm n = catchJust undefElements (fm ! n) $
                     \_ -> do putStr ('[' : show n)
                              putStrLn "] ### missing footnote ###"

-- undefElements: exception predicate for undefined array elements.
undefElements :: Exception -> Maybe ()
undefElements e = guard $ "undefined array element" `isSuffixOf` show e

-- insertFootRef: insert the ref from footnote line l into st.
insertFootRef :: RefsMapState -> B.ByteString -> RefsMapState
insertFootRef st l = maybe st (insertRef st . fst) $ parseFoot l

-- putUnknownRef: print dummy for a reference without a footnote.
putUnknownRef :: Int -> RefsMapState -> IO RefsMapState
putUnknownRef _ st = putStr "?" >> return st