{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE DeriveAnyClass #-}

import Control.DeepSeq (rnf)
import Control.Concurrent
import Control.Exception
import Control.Monad

import qualified Data.ByteString.Lazy.Char8 as BS
import Data.ByteString.Lazy.Char8 (pack, unpack)
import qualified Data.HashMap.Strict as HMap
import qualified Data.Text.Encoding as T
import qualified Data.Map.Strict as Map
import qualified Data.Text as T
import Data.Functor
import Data.Monoid
import Data.Binary
import Data.Maybe
import Data.Aeson
import Data.List
import Data.Char


import Hakyll.Core.Compiler.Internal
import Hakyll.Core.Compiler
import Hakyll.Web.Sass
import Hakyll

import qualified Network.URI.Encode as URI (encode)

import qualified Skylighting as Sky

import System.Directory
import System.Environment
import System.Process
import System.Exit
import System.IO

import Text.Pandoc.Walk (query, walkM, walk)
import Text.Pandoc.Highlighting
import Text.Pandoc.Definition
import Text.Pandoc.Options
import Text.Sass.Functions

import Data.Text (Text)
import Data.IORef
import Data.Hashable (Hashable (hashWithSalt))
import GHC.Stack
import Text.Read (readMaybe)
import GHC.Show (showCommaSpace)
import Data.Traversable
import qualified Data.Text.Lazy as LT

readerOpts :: ReaderOptions
readerOpts = def { readerExtensions = pandocExtensions
                 , readerIndentedCodeClasses = ["amulet"] }

writerOptions :: Compiler WriterOptions
writerOptions = do
  syntaxMap <- loadAllSnapshots "syntax/*.xml" "syntax"
           <&> foldr (Sky.addSyntaxDefinition . itemBody) Sky.defaultSyntaxMap

  pure $ defaultHakyllWriterOptions
    { writerExtensions = extensionsFromList
                         [ Ext_tex_math_dollars
                         , Ext_tex_math_double_backslash
                         , Ext_latex_macros
                         ] <> writerExtensions defaultHakyllWriterOptions
    , writerSyntaxMap = syntaxMap
    , writerHighlightStyle = Just kate
    }

rssfeed :: FeedConfiguration
rssfeed
  = FeedConfiguration { feedTitle = "Amelia's Blag: Latest articles"
                      , feedDescription = ""
                      , feedAuthorName = "Amélia"
                      , feedAuthorEmail = "me@amelia.how"
                      , feedRoot = "https://amelia.how"
                      }

conf :: Configuration
conf = def { destinationDirectory = ".site"
           , storeDirectory       = ".store"
           , tmpDirectory         = ".store/tmp"
           , deployCommand        = "./sync" }

katexFilter :: MVar KatexCache -> Pandoc -> Compiler Pandoc
katexFilter cacheVar (Pandoc meta doc) =
  do
    id <- compilerUnderlying <$> compilerAsk
    t <- getMetadata id

    invalidateCache id (abbrevs t) cacheVar

    doc <- Pandoc meta <$> walkM (go (show id) (abbrevs t)) doc
    unsafeCompiler $ flushCache cacheVar
    pure doc
  where
    abbrevs :: HMap.HashMap Text Value -> [String]
    abbrevs x =
      case HMap.lookup "abbreviations" x of
        Just (Object map) -> concat $ mapMaybe oneAbbrev (HMap.toList map)
        _ -> []
     
    oneAbbrev (x, String t) = Just ["-m", '\\':T.unpack x ++ ':':T.unpack t]
    oneAbbrev _ = Nothing

    go :: String -> [String] -> Inline -> Compiler Inline
    go id abbrevs (Math kind math) = unsafeCompiler $ do
      cache <- readMVar cacheVar
      case HMap.lookup (id, kind, math) (spanMap cache) of
        Just x -> pure (RawInline "html" x)
        Nothing -> do
          let args = flip (:) abbrevs $ case kind of { DisplayMath -> "-td"; InlineMath  -> "-t" }

          (contents, _) <- readProcessBS "katex" args . BS.fromStrict . T.encodeUtf8 $ math
          let text = T.init . T.init . T.decodeUtf8 . BS.toStrict $ contents

          modifyMVar cacheVar (\m -> pure (bumpCacheEntry cache id abbrevs kind text math, ()))
          pure $ RawInline "html" text

    go id _ x = pure x

    bumpCacheEntry (KatexCache spans depends abbrevVars) id abbrevs kind text math =
      let
        str = T.unpack text
        usedAbbrevs = map (\x -> (T.pack (takeWhile (/= ':') (tail x)), T.pack (tail (dropWhile (/= ':') (tail x)))))
                    $ filter (\x -> (takeWhile (/= ':') (tail x)) `isInfixOf` str)
                    $ filter (not . isPrefixOf "-") abbrevs
        addDeps [] x = x
        addDeps ((k, _):xs) vl = HMap.alter (\v -> Just (maybe [(id, kind, math)] ((id, kind, math):) v)) (id, k)  $ addDeps xs vl

        recordVars [] x = x
        recordVars ((k, v):xs) x = HMap.insert (id, k) v (recordVars xs x)
       in
         case usedAbbrevs of
           [] -> KatexCache (HMap.insert (id, kind, math) text spans) depends abbrevVars
           xs -> KatexCache (HMap.insert (id, kind, math) text spans) (addDeps xs depends) (recordVars xs abbrevVars)

abbreviationFilter :: Pandoc -> Compiler Pandoc
abbreviationFilter (Pandoc meta doc) =
  do
    id <- compilerUnderlying <$> compilerAsk
    t <- getMetadata id
    case HMap.lookup "abbreviations" t of
      Just (Object map) -> do
        pure (Pandoc meta (walk (replace map) doc))
      _ -> pure (Pandoc meta doc)
  where
    replace map x =
      case x of
        Str t | Just (e, r) <- entity t -> fromMaybe (Str t) (toInline r =<< HMap.lookup e map)
        x -> x

    toInline r (String t) = Just (Str (t <> r))
    toInline _ _ = Nothing

    entity x
      | T.isPrefixOf "&" x && T.length x >= 3 =
        let
          (name, rest') = T.span (/= ';') (T.tail x)
          rest = T.tail rest'
        in pure (name, rest)
      | otherwise = Nothing


addLanguageTag :: Pandoc -> Pandoc
addLanguageTag (Pandoc meta blocks) = Pandoc meta (map go blocks) where
  go :: Block -> Block
  go block@(CodeBlock (identifier, classes@(language:classes'), kv) text) =
    Div
      ( mempty
      , "code-container":if haskv then "custom-tag":classes' else classes'
      , []
      )
      $ [block] ++ maybe [Div (mempty, ["code-tag"], []) [Plain [Span (mempty, [], []) [Str tag]]]]
    where
      language' = case T.uncons language of
        Nothing -> mempty
        Just (c, cs) -> T.cons (toUpper c) cs
      tag   = fromMaybe language' (lookup "tag" kv)
      haskv = fromMaybe False (True <$ lookup "tag" kv)
      maybe
        | "notag" `elem` classes' = const []
        | otherwise = id

  go block@(CodeBlock (identifier, [], kv) text) =
    Div (mempty, ["code-container"], [])
      [block, Div (mempty, ["empty-code-tag"], []) []]
  go x = x

saveSynopsys :: Pandoc -> Compiler Pandoc
saveSynopsys (Pandoc meta doc) =
  do
    id <- getUnderlying
    n <- fromMaybe (1 :: Int) . readMaybe . fromMaybe "" . lookupString "synopsys" <$> getMetadata id

    case dropWhile (not . isParagraph) doc of
      p:ps -> do
        saveSnapshot "synopsys-block" =<< makeItem (map removeFootnotes (take n (p:ps)))
        pure ()
      [] -> pure ()
    pure $ Pandoc meta doc
  where
    isParagraph Para{} = True
    isParagraph _ = False

    removeFootnotes (Para xs) = Para $ filter (\case { Note _ -> False; _ -> True }) xs
    removeFootnotes x = x

saveWordCount :: Pandoc -> Compiler Pandoc
saveWordCount (Pandoc meta doc) =
  do
    saveSnapshot "wc" =<< makeItem wordCount
    pure $ Pandoc meta doc
  where
    wordCount = show (getSum (query inlineLen doc))

    inlineLen (Str s) = Sum (length (T.words s))
    inlineLen _ = mempty

saveTableOfContents :: Pandoc -> Compiler Pandoc
saveTableOfContents (Pandoc meta input) =
  do
    saveSnapshot "table-of-contents" =<< makeItem toc
    pure $ Pandoc meta (fixHeaders 0 doc)
  where
    headers = filter (\case { Header _ _ _ -> True; _ -> False }) doc

    doc = fixHeaders 0 input

    fixHeaders n (Header l (_, ms, mt) x:bs) =
      Header l (anchor, ms, mt) (Link (anchor, ms, mt) [] (T.singleton '#' <> anchor, mempty):x):fixHeaders (n + 1) bs where
        anchor = T.pack ("h" ++ show n)
    fixHeaders k (x:bs) = x:fixHeaders k bs
    fixHeaders _ [] = []

    into :: [Block] -> [[Block]]
    into (Header l m@(anchor, _, _) x:ms) =
      let
        contained (Header l' _ _) = l' > l
        contained _ = undefined

        (ours, rest) = span contained ms
      in [Para [Link (mempty, mempty, mempty) (tail x) (T.singleton '#' <> anchor, mempty)], list (into ours)]:into rest
    into [] = []
    into _ = undefined
    
    list = BulletList

    toc :: Block
    toc = list (into headers)

setup :: IO (MVar KatexCache)
setup = do
  setEnv "AMC_LIBRARY_PATH" "/usr/lib/amuletml/lib/"
  loadCache

compiler :: MVar KatexCache -> Compiler (Item String)
compiler katexCache = do
  wops <- writerOptions
  pandocCompilerWithTransformM readerOpts wops $
          katexFilter katexCache
      >=> abbreviationFilter
      >=> saveSynopsys
      >=> saveWordCount
      >=> saveTableOfContents

main :: IO ()
main = setup >>= \katexCache -> hakyllWith conf $ do
  match "static/*" do
    route idRoute
    compile copyFileCompiler

  match "static/**/*" $ do
    route idRoute
    compile copyFileCompiler

  match "css/**/*" $ do
    route idRoute
    compile copyFileCompiler

  match "css/*.css" $ do
    route idRoute
    compile copyFileCompiler

  match "css/*.scss" $ do
    route $ setExtension "css"
    compile $ do
      imports <- unsafeCompiler $ newMVar ([] :: [(String, String)])
      let add f p = modifyMVar imports (\x -> pure ((f, p):x, []))

      body <- sassCompilerWith def
                { sassOutputStyle = SassStyleCompressed
                , sassImporters = Just [ SassImporter 0 add ]
                }        

      list <- unsafeCompiler $ takeMVar imports
      for list $ \(req, path) -> do
        load (fromFilePath ("css/" ++ reverse (dropWhile (/= '.') (reverse req)) ++ "scss"))
          :: Compiler (Item String)

      pure body

  match "diagrams/**/*.tex" $ do
    route $ setExtension "svg"
    compile $ getResourceBody
          >>= loadAndApplyTemplate "templates/tikz.tex" (bodyField "body")
          >>= withItemBody (return . pack
                        >=> unixFilterLBS "rubber-pipe" ["--pdf"]
                        >=> unixFilterLBS "pdftocairo" ["-svg", "-", "-"]
                        >=> return . unpack)

  match "pages/posts/*" do
    route $ metadataRoute pathFromTitle

    compile $ do
      wops <- writerOptions
      id <- getUnderlying
      r <- compiler katexCache
        >>= loadAndApplyTemplate "templates/post.html"    postCtx
        >>= saveSnapshot "content"
        >>= loadAndApplyTemplate "templates/default.html" postCtx
        >>= relativizeUrls

      loadSnapshot id "synopsys-block"
        >>= saveSnapshot "synopsys-text"
          . writePandocWith wops
          . fmap (Pandoc mempty)

      pure r

  match "pages/posts/*.lhs" $ version "raw" $ do
    route idRoute
    compile copyFileCompiler

  create ["archive.html"] $ do
    route idRoute
    compile $ do
      posts <- recentFirst =<< onlyPublic =<< loadAll ("pages/posts/*" .&&. hasNoVersion)
      let archiveCtx =
            listField "posts" postCtx (return posts) <>
            constField "title" "Archives"            <>
            defaultContext
      makeItem ""
        >>= loadAndApplyTemplate "templates/archive.html" archiveCtx
        >>= loadAndApplyTemplate "templates/default.html" archiveCtx
        >>= relativizeUrls

  match "pages/*.html" $ do
    route $ gsubRoute "pages/" (const "")
    compile $ do
      posts <- fmap (take 5) . recentFirst =<< onlyPublic =<< loadAll ("pages/posts/*" .&&. hasNoVersion)
      let indexCtx =
            listField "posts" postCtx (return posts) <>
            constField "title" "Home"                <>
            defaultContext
      getResourceBody
        >>= applyAsTemplate indexCtx
        >>= loadAndApplyTemplate "templates/default.html" indexCtx
        >>= relativizeUrls

  match "pages/*.md" $ do
    route $ setExtension "html" <> gsubRoute "pages/" (const "")
    compile $ compiler katexCache
      >>= loadAndApplyTemplate "templates/page.html" defaultContext
      >>= loadAndApplyTemplate "templates/default.html" defaultContext
      >>= relativizeUrls

  match "syntax/*.xml" $ compile $ do
    path <- toFilePath <$> getUnderlying
    contents <- itemBody <$> getResourceBody
    debugCompiler ("Loaded syntax definition from " ++ show path)
    let res = Sky.parseSyntaxDefinitionFromText path (LT.pack contents)
    _ <- saveSnapshot "syntax" =<< case res of
      Left e -> fail e
      Right x -> makeItem x
    makeItem contents

  match "templates/*" $ compile templateBodyCompiler

  create ["feed.xml"] $ do
    route idRoute
    compile $ do
      let feedCtx = postCtx <> bodyField "description"
      posts <- fmap (take 10) . recentFirst =<< onlyPublic =<< loadAllSnapshots ("pages/posts/*" .&&. hasNoVersion) "synopsys-text"
      renderRss rssfeed feedCtx posts

onlyPublic :: [Item String] -> Compiler [Item String]
onlyPublic = filterM isPublic where
  isPublic item = do
    t <- getMetadata (itemIdentifier item)
    case HMap.lookup "public" t of
      Just (Bool False) -> pure False
      _ -> pure True

postCtx :: Context String
postCtx =
       dateField "date" "%B %e, %Y"
    <> snapshotField "synopsys" "synopsys-text"
    <> snapshotField "words" "wc"
    <> snapshotField' render "toc" "table-of-contents"
    <> defaultContext
  where
    snapshotField = snapshotField' pure
    snapshotField' f key snap = field key $ \x -> do
      let id = itemIdentifier x
      fmap itemBody . f =<< loadSnapshot id snap
      
    render x = do
      wops <- writerOptions
      pure . writePandocWith wops . fmap (Pandoc mempty . pure) $ x

readProcessBS :: FilePath -> [String] -> BS.ByteString -> IO (BS.ByteString, String)
readProcessBS path args input =
  let process = (proc path args)
        { std_in = CreatePipe
        , std_out = CreatePipe
        , std_err = CreatePipe
        }
  in withCreateProcess process $ \stdin stdout stderr ph ->
    case (stdin, stdout, stderr) of
      (Nothing, _, _) -> fail "Failed to get a stdin handle."
      (_, Nothing, _) -> fail "Failed to get a stdout handle."
      (_, _, Nothing) -> fail "Failed to get a stderr handle."
      (Just stdin, Just stdout, Just stderr) -> do
        out <- BS.hGetContents stdout
        err <- hGetContents stderr

        withForkWait (evaluate $ rnf out) $ \waitOut ->
          withForkWait (evaluate $ rnf err) $ \waitErr -> do
            -- Write input and close.
            BS.hPutStr stdin input
            hClose stdin

            -- wait on the output
            waitOut
            waitErr

            hClose stdout
            hClose stderr

            -- wait on the process
            ex <- waitForProcess ph
            case ex of
              ExitSuccess -> pure (out, err)
              ExitFailure ex -> fail (err ++ "Exited with " ++ show ex)

  where
    withForkWait :: IO () -> (IO () ->  IO a) -> IO a
    withForkWait async body = do
      waitVar <- newEmptyMVar :: IO (MVar (Either SomeException ()))
      mask $ \restore -> do
        tid <- forkIO $ try (restore async) >>= putMVar waitVar
        let wait = takeMVar waitVar >>= either throwIO return
        restore (body wait) `onException` killThread tid

pathFromTitle :: Metadata -> Routes
pathFromTitle meta =
  let
    declaredCategory =
      case lookupString "category" meta of
        Just s -> ((s ++ "/") ++)
        Nothing -> ("posts/" <>)

    !titleString =
      case lookupString "title" meta of
        Just s -> s
        Nothing -> error "post has no title?"

    title = filter (/= "") . map (filter isAlphaNum . map toLower) . words $ titleString

    (category, title') =
      if | "or" `elem` title                   -> (declaredCategory, takeWhile (/= "or") title)
         | ["a", "quickie"] `isPrefixOf` title -> (("quick/" ++), drop 2 title)
         | otherwise                           -> (declaredCategory, title)
  in
    case lookupString "path" meta of
      Just p -> constRoute (category (p <> ".html"))
      Nothing -> constRoute (category (intercalate "-" title' <> ".html"))

foldMapM :: (Monad w, Monoid m, Foldable f) => (a -> w m) -> f a -> w m
foldMapM k = foldr (\x y -> do { m <- k x; (m <>) <$> y }) (pure mempty)

loadCache :: HasCallStack => IO (MVar KatexCache)
loadCache = do
  t <- doesFileExist ".katex_cache"
  let fixup (a, b, c) = KatexCache (HMap.fromList a) (HMap.fromList b) (HMap.fromList c)
  map <- if t
    then (fixup <$> decodeFile ".katex_cache") `catch` \e ->
            const (print e *> pure (KatexCache mempty mempty mempty)) (e :: SomeException)
    else pure (KatexCache mempty mempty mempty)
  var <- newMVar map
  pure var

flushCache :: MVar KatexCache -> IO ()
flushCache var = do
  withMVar var $ \(KatexCache x y z) -> do
    Data.Binary.encodeFile ".katex_cache" (HMap.toList x, HMap.toList y, HMap.toList z)

invalidateCache :: Identifier -> [String] -> MVar KatexCache -> Compiler KatexCache
invalidateCache id abbrevs cacheVar = unsafeCompiler $ modifyMVar cacheVar (\x -> pure (go x, go x)) where
  currentValues = map (\x -> (T.pack (takeWhile (/= ':') (tail x)), T.pack (tail (dropWhile (/= ':') (tail x)))))
                $ filter (not . isPrefixOf "-") abbrevs
  ident = show id

  go (KatexCache spanMap abbrKeys abbrText) =
    let
      go (abbr, val) (spanMap, abbrKeys, abbrText) =
        case HMap.lookup (ident, abbr) abbrText of
          Just vl | vl /= val ->
            let l = HMap.lookupDefault [] (ident, abbr) abbrKeys
             in (foldr HMap.delete spanMap l, abbrKeys, abbrText)
          _ -> (spanMap, abbrKeys, abbrText)

      (a, b, c) = foldr go (spanMap, abbrKeys, abbrText) currentValues
     in KatexCache a b c

data KatexCache
  = KatexCache { spanMap          :: HMap.HashMap (String, MathType, Text) Text
               , abbreviationKeys :: HMap.HashMap (String, Text) [(String, MathType, Text)]
               , abbreviationText :: HMap.HashMap (String, Text) Text
               } 

instance Hashable MathType where
  hashWithSalt s DisplayMath = hashWithSalt s (0 :: Int)
  hashWithSalt s InlineMath = hashWithSalt s (1 :: Int)

deriving instance Binary Block
deriving instance Binary Inline
deriving instance Binary Format
deriving instance Binary ListNumberStyle
deriving instance Binary ListNumberDelim
deriving instance Binary Caption
deriving instance Binary Alignment
deriving instance Binary ColWidth
deriving instance Binary TableHead
deriving instance Binary TableBody
deriving instance Binary TableFoot
deriving instance Binary QuoteType
deriving instance Binary Citation
deriving instance Binary Row
deriving instance Binary MathType
deriving instance Binary RowHeadColumns
deriving instance Binary CitationMode
deriving instance Binary Cell
deriving instance Binary RowSpan
deriving instance Binary ColSpan