{-# LANGUAGE OverloadedStrings #-}

module Main where

import Text.Pandoc
import Text.Pandoc.Error
import Text.Pandoc.Class (runIOorExplode)
import Text.Pandoc.PDF (makePDF)
import System.Environment (getArgs)
import System.FilePath (replaceExtension)
import System.Process (callProcess)
import System.IO (writeFile)
import System.Directory (doesFileExist)
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.IO as TIO
import Data.Hashable (hash)
import Control.Monad (when, void)

-- Transform Mermaid code blocks into image embeds
processMermaid :: Block -> IO Block
processMermaid (CodeBlock (id', classes, _) contents)
  | "mermaid" `elem` classes = do
      let baseName = if null id' then "diagram-" ++ take 6 (show (abs (hash contents))) else id'
          mmdFile = baseName ++ ".mmd"
          svgFile = baseName ++ ".svg"

      writeFile mmdFile contents
      exists <- doesFileExist svgFile
      when (not exists) $ void $ callProcess "mmdc" ["-i", mmdFile, "-o", svgFile]

      return $ Para [Image nullAttr [] (T.pack svgFile, "Mermaid diagram")]
processMermaid x = return x

-- Walk the Pandoc AST and process blocks
transformDoc :: Pandoc -> IO Pandoc
transformDoc = walkM processMermaid

main :: IO ()
main = do
  args <- getArgs
  case args of
    ["-pdf", path] -> compileToPDF path
    ["-html", path] -> compileToHTML path
    _ -> putStrLn "Usage: docster -pdf|-html <file.md>"

compileToPDF :: FilePath -> IO ()
compileToPDF path = do
  content <- TIO.readFile path
  pandoc <- runIOorExplode $ readMarkdown def content
  transformed <- transformDoc pandoc

  let outputPath = replaceExtension path "pdf"
  result <- runIOorExplode $ makePDF "pdflatex" [] transformed
  case result of
    Left err -> error $ "PDF error: " ++ show err
    Right bs -> writeFile outputPath bs >> putStrLn ("✅ PDF written to " ++ outputPath)

compileToHTML :: FilePath -> IO ()
compileToHTML path = do
  content <- TIO.readFile path
  pandoc <- runIOorExplode $ readMarkdown def content
  transformed <- transformDoc pandoc

  let outputPath = replaceExtension path "html"
  html <- runIOorExplode $ writeHtml5String def transformed
  TIO.writeFile outputPath html
  putStrLn ("✅ HTML written to " ++ outputPath)