{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE BlockArguments #-}

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)

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 = "Abigail's Blag: Latest articles"
                      , feedDescription = ""
                      , feedAuthorName = "Abigail Magalhães"
                      , feedAuthorEmail = "[email protected]"
                      , feedRoot = "https://abby.how"
                      }

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

katexFilter :: IORef KatexCache -> Pandoc -> Compiler Pandoc
katexFilter cacheVar (Pandoc meta doc) =
  do
    initCache <- unsafeCompiler (readIORef cacheVar)
    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 <- readIORef 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 "node_modules/.bin/katex" args . BS.fromStrict . T.encodeUtf8 $ math
          let text = T.init . T.init . T.decodeUtf8 . BS.toStrict $ contents

          atomicModifyIORef' cacheVar (\m -> (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

estimateReadingTime :: Pandoc -> Pandoc
estimateReadingTime (Pandoc meta doc) = Pandoc meta doc' where
  wordCount = T.pack (show (getSum (query inlineLen doc)))

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

  doc' = RawBlock "html" ("<span id=reading-length>" <> wordCount <> "</span>")
       : doc

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, 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)

  go block@(CodeBlock (identifier, [], kv) text) = Div (mempty, ["code-container"], []) [block]
  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" =<< makeItem (take n (p:ps))
        pure ()
      [] -> pure ()
    pure $ Pandoc meta doc
  where
    isParagraph Para{} = True
    isParagraph _ = False

sassImporter :: SassImporter
sassImporter = SassImporter 0 go where
  go "normalize" _ = do
    c <- readFile "node_modules/normalize.css/normalize.css"
    pure [ SassImport { importPath = Nothing
                      , importAbsolutePath = Nothing
                      , importSource = Just c
                      , importSourceMap = Nothing
                      } ]
  go _ _ = pure []

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

compiler :: IORef KatexCache -> Compiler (Item String)
compiler katexCache = do
  wops <- writerOptions
  pandocCompilerWithTransformM readerOpts wops $
          katexFilter katexCache
      >=> abbreviationFilter
      >=> saveSynopsys
      >=> pure . estimateReadingTime
      >=> pure . addLanguageTag

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 $ sassCompilerWith def { sassOutputStyle = SassStyleCompressed
                                   , sassImporters = Just [ sassImporter ]
                                   }

  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" >>= saveSnapshot "synopsys" . 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/default.html" defaultContext
      >>= relativizeUrls

  match "syntax/*.xml" $ compile $ do
    path <- toFilePath <$> getUnderlying
    contents <- itemBody <$> getResourceBody
    debugCompiler ("Loaded syntax definition from " ++ show path)
    res <- unsafeCompiler (Sky.parseSyntaxDefinitionFromString path 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"
      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"
    <> synopsysField
    <> defaultContext
  where
    synopsysField = field "synopsys" $ \x -> do
      let id = itemIdentifier x
      itemBody <$> loadSnapshot id "synopsys"

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 (IORef 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 <- newIORef map
  pure var

flushCache :: IORef KatexCache -> IO ()
flushCache var = do
  KatexCache x y z <- readIORef var
  Data.Binary.encodeFile ".katex_cache" (HMap.toList x, HMap.toList y, HMap.toList z)

invalidateCache :: Identifier -> [String] -> IORef KatexCache -> Compiler KatexCache
invalidateCache id abbrevs cacheVar = unsafeCompiler $ atomicModifyIORef' cacheVar (\x -> (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)