{-# 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, takeDirectory, (</>))
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
processMermaidInDir :: FilePath -> Block -> IO Block
processMermaidInDir sourceDir block@(CodeBlock (id', classes, _) contents)
  | "mermaid" `elem` classes = do
      let baseName = if T.null id' then "diagram-" ++ take 6 (show (abs (hash (T.unpack contents)))) else T.unpack id'
          mmdFile = sourceDir </> baseName ++ ".mmd"
          pngFile = sourceDir </> baseName ++ ".png"

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

      return $ Para [Image nullAttr [] (T.pack pngFile, "Mermaid diagram")]
processMermaidInDir _ x = return x


-- Walk the Pandoc AST and process blocks using walkM
transformDoc :: FilePath -> Pandoc -> IO Pandoc
transformDoc sourceDir = walkM (processMermaidInDir sourceDir)

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" }
      sourceDir = takeDirectory path
  pandoc <- runIOorExplode $ readMarkdown readerOptions content
  transformed <- transformDoc sourceDir pandoc

  let outputPath = replaceExtension path "pdf"
      writerOptions = def
  -- Generate LaTeX and add proper header with tightlist definition
  latexOutput <- runIOorExplode $ writeLaTeX writerOptions transformed
  let latexWithProperHeader = T.unlines [
        "\\documentclass{article}",
        "\\usepackage[utf8]{inputenc}",
        "\\usepackage{fontspec}",
        "\\usepackage{graphicx}",
        "\\usepackage{geometry}",
        "\\geometry{margin=1in}",
        "\\usepackage{hyperref}",
        "\\usepackage{enumitem}",
        "\\usepackage{amsmath}",
        "\\usepackage{amssymb}",
        "\\usepackage{fancyvrb}",
        "\\usepackage{color}",
        "\\DefineVerbatimEnvironment{Highlighting}{Verbatim}{commandchars=\\\\\\{\\}}",
        "\\newenvironment{Shaded}{}{}",
        "\\newcommand{\\AlertTok}[1]{\\textcolor[rgb]{1.00,0.00,0.00}{\\textbf{#1}}}",
        "\\newcommand{\\AnnotationTok}[1]{\\textcolor[rgb]{0.38,0.63,0.69}{\\textbf{\\textit{#1}}}}",
        "\\newcommand{\\AttributeTok}[1]{\\textcolor[rgb]{0.49,0.56,0.16}{#1}}",
        "\\newcommand{\\BaseNTok}[1]{\\textcolor[rgb]{0.25,0.63,0.44}{#1}}",
        "\\newcommand{\\BuiltInTok}[1]{#1}",
        "\\newcommand{\\CharTok}[1]{\\textcolor[rgb]{0.25,0.44,0.63}{#1}}",
        "\\newcommand{\\CommentTok}[1]{\\textcolor[rgb]{0.38,0.63,0.69}{\\textit{#1}}}",
        "\\newcommand{\\CommentVarTok}[1]{\\textcolor[rgb]{0.38,0.63,0.69}{\\textbf{\\textit{#1}}}}",
        "\\newcommand{\\ConstantTok}[1]{\\textcolor[rgb]{0.53,0.00,0.00}{#1}}",
        "\\newcommand{\\ControlFlowTok}[1]{\\textcolor[rgb]{0.00,0.44,0.13}{\\textbf{#1}}}",
        "\\newcommand{\\DataTypeTok}[1]{\\textcolor[rgb]{0.56,0.13,0.00}{#1}}",
        "\\newcommand{\\DecValTok}[1]{\\textcolor[rgb]{0.25,0.63,0.44}{#1}}",
        "\\newcommand{\\DocumentationTok}[1]{\\textcolor[rgb]{0.73,0.13,0.13}{\\textit{#1}}}",
        "\\newcommand{\\ErrorTok}[1]{\\textcolor[rgb]{1.00,0.00,0.00}{\\textbf{#1}}}",
        "\\newcommand{\\ExtensionTok}[1]{#1}",
        "\\newcommand{\\FloatTok}[1]{\\textcolor[rgb]{0.25,0.63,0.44}{#1}}",
        "\\newcommand{\\FunctionTok}[1]{\\textcolor[rgb]{0.02,0.16,0.49}{#1}}",
        "\\newcommand{\\ImportTok}[1]{#1}",
        "\\newcommand{\\InformationTok}[1]{\\textcolor[rgb]{0.38,0.63,0.69}{\\textbf{\\textit{#1}}}}",
        "\\newcommand{\\KeywordTok}[1]{\\textcolor[rgb]{0.00,0.44,0.13}{\\textbf{#1}}}",
        "\\newcommand{\\NormalTok}[1]{#1}",
        "\\newcommand{\\OperatorTok}[1]{\\textcolor[rgb]{0.40,0.40,0.40}{#1}}",
        "\\newcommand{\\OtherTok}[1]{\\textcolor[rgb]{0.00,0.44,0.13}{#1}}",
        "\\newcommand{\\PreprocessorTok}[1]{\\textcolor[rgb]{0.74,0.48,0.00}{#1}}",
        "\\newcommand{\\RegionMarkerTok}[1]{#1}",
        "\\newcommand{\\SpecialCharTok}[1]{\\textcolor[rgb]{0.25,0.44,0.63}{#1}}",
        "\\newcommand{\\SpecialStringTok}[1]{\\textcolor[rgb]{0.73,0.40,0.53}{#1}}",
        "\\newcommand{\\StringTok}[1]{\\textcolor[rgb]{0.25,0.44,0.63}{#1}}",
        "\\newcommand{\\VariableTok}[1]{\\textcolor[rgb]{0.10,0.09,0.49}{#1}}",
        "\\newcommand{\\VerbatimStringTok}[1]{\\textcolor[rgb]{0.25,0.44,0.63}{#1}}",
        "\\newcommand{\\WarningTok}[1]{\\textcolor[rgb]{0.38,0.63,0.69}{\\textbf{\\textit{#1}}}}",
        "\\providecommand{\\tightlist}{%",
        "  \\setlength{\\itemsep}{0pt}\\setlength{\\parskip}{0pt}}",
        "\\begin{document}"
        ] <> latexOutput <> "\n\\end{document}"
  result <- runIOorExplode $ makePDF "xelatex" [] (\_ _ -> return latexWithProperHeader) 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
  let readerOptions = def { readerExtensions = getDefaultExtensions "markdown" }
      sourceDir = takeDirectory path
  pandoc <- runIOorExplode $ readMarkdown readerOptions content
  transformed <- transformDoc sourceDir pandoc

  let outputPath = replaceExtension path "html"
  html <- runIOorExplode $ writeHtml5String def transformed
  TIO.writeFile outputPath html
  putStrLn ("✅ HTML written to " ++ outputPath)
  
  -- Open the generated HTML file in browser for Claude Code to check errors
  putStrLn $ "🌐 Opening " ++ outputPath ++ " in browser for error checking..."
  void $ callProcess "open" [outputPath]