{-# LANGUAGE OverloadedStrings #-}

module Main where

import Text.Pandoc
import Text.Pandoc.Error
import Text.Pandoc.Class (runIOorExplode)
import Text.Pandoc.PDF (makePDF)
import Text.Pandoc.Walk (walkM)
import Text.Pandoc.Extensions (Extension(..), enableExtension, getDefaultExtensions)
import System.Environment (getArgs)
import System.FilePath (replaceExtension)
import System.Process (callProcess)
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)
import qualified Data.ByteString.Lazy as BL

-- Transform Mermaid code blocks into image embeds
processMermaid :: Block -> IO Block
processMermaid block@(CodeBlock (id', classes, _) contents)
  | "mermaid" `elem` classes = do
      putStrLn $ "🎯 Found Mermaid block with classes: " ++ show classes
      let baseName = if T.null id' then "diagram-" ++ take 6 (show (abs (hash (T.unpack contents)))) else T.unpack id'
          mmdFile = baseName ++ ".mmd"
          pngFile = baseName ++ ".png"

      putStrLn $ "📝 Writing to " ++ mmdFile ++ " and generating " ++ pngFile
      writeFile mmdFile (T.unpack contents)
      void $ callProcess "mmdc" ["-i", mmdFile, "-o", pngFile]
      putStrLn $ "✅ Generated " ++ pngFile

      return $ Para [Image nullAttr [] (T.pack pngFile, "Mermaid diagram")]
processMermaid x = do
  -- Debug: show what blocks we're processing
  case x of
    CodeBlock (_, classes, _) _ -> putStrLn $ "📄 Found code block with classes: " ++ show classes
    _ -> return ()
  return x

-- Walk the Pandoc AST and process blocks using walkM
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>"

pdfTemplate :: T.Text
pdfTemplate = T.unlines [
  "\\documentclass{article}",
  "\\usepackage[utf8]{inputenc}",
  "\\usepackage{graphicx}",
  "\\usepackage{geometry}",
  "\\geometry{margin=1in}",
  "\\usepackage{hyperref}",
  "\\usepackage{enumitem}",
  "\\providecommand{\\tightlist}{%",
  "  \\setlength{\\itemsep}{0pt}\\setlength{\\parskip}{0pt}}",
  "\\title{$title$}",
  "\\author{$author$}",
  "\\date{$date$}",
  "\\begin{document}",
  "$if(title)$\\maketitle$endif$",
  "$body$",
  "\\end{document}"
  ]

compileToPDF :: FilePath -> IO ()
compileToPDF path = do
  content <- TIO.readFile path
  let readerOptions = def { readerExtensions = getDefaultExtensions "markdown" }
  pandoc <- runIOorExplode $ readMarkdown readerOptions content
  transformed <- transformDoc pandoc

  let outputPath = replaceExtension path "pdf"
      writerOptions = def
  -- First generate LaTeX with proper document structure
  latexOutput <- runIOorExplode $ writeLaTeX writerOptions transformed
  let latexWithHeader = "\\documentclass{article}\n\\usepackage[utf8]{inputenc}\n\\usepackage{graphicx}\n\\usepackage{geometry}\n\\geometry{margin=1in}\n\\begin{document}\n" <> latexOutput <> "\n\\end{document}"
  result <- runIOorExplode $ makePDF "pdflatex" [] (\_ _ -> return latexWithHeader) def transformed
  case result of
    Left err -> error $ "PDF error: " ++ show err
    Right bs -> BL.writeFile outputPath bs >> putStrLn ("✅ PDF written to " ++ outputPath)

compileToHTML :: FilePath -> IO ()
compileToHTML path = do
  content <- TIO.readFile path
  putStrLn $ "📖 Reading: " ++ path
  putStrLn $ "📝 Content: " ++ T.unpack content
  let readerOptions = def { readerExtensions = getDefaultExtensions "markdown" }
  pandoc <- runIOorExplode $ readMarkdown readerOptions content
  putStrLn $ "🔍 Parsed AST: " ++ show pandoc
  transformed <- transformDoc pandoc

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