{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Main (main) where

import Text.Pandoc
import Text.Pandoc.PDF (makePDF)
import System.Environment (getArgs)
import System.FilePath (replaceExtension, takeDirectory, takeFileName, (</>))
import System.Process (callProcess)
import System.Directory (removeFile)
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.IO as TIO
import Data.Hashable (hash)
import Control.Monad (void)
import Control.Exception (Exception, throwIO, bracket, catch, SomeException)
import qualified Data.ByteString.Lazy as BL

-- | Custom error types for better error handling
data DocsterError 
  = InvalidUsage Text
  | FileError Text
  | PDFGenerationError Text
  | ProcessError Text
  deriving (Show)

instance Exception DocsterError

-- | Output format for explicit handling instead of file extension checking
data OutputFormat = PDF | HTML deriving (Show, Eq)

-- | Type-safe wrappers for better domain modeling
newtype SourceDir = SourceDir FilePath deriving (Show, Eq)
newtype OutputPath = OutputPath FilePath deriving (Show, Eq)
newtype DiagramId = DiagramId Text deriving (Show, Eq)

-- | Configuration for diagram generation
data DiagramConfig = DiagramConfig
  { dcSourceDir :: SourceDir
  , dcOutputFormat :: OutputFormat
  } deriving (Show)

-- | Constants for the application (using Text for consistency)
mermaidCommand :: Text
mermaidCommand = "mmdc"

diagramPrefix :: Text
diagramPrefix = "diagram-"

successEmoji :: Text
successEmoji = "✅"

-- | Generate a diagram ID from content hash or explicit ID
generateDiagramId :: Text -> Text -> DiagramId
generateDiagramId explicitId contents
  | T.null explicitId = DiagramId $ diagramPrefix <> T.take 6 (T.pack . show . abs . hash $ T.unpack contents)
  | otherwise = DiagramId explicitId

-- | Transform Mermaid code blocks into image embeds with resource cleanup
processMermaidBlock :: DiagramConfig -> Block -> IO (Either DocsterError Block)
processMermaidBlock config (CodeBlock (id', classes, _) contents)
  | "mermaid" `elem` classes = do
      let diagId = generateDiagramId id' contents
      result <- renderMermaidDiagram config diagId contents
      case result of
        Left err -> return $ Left err
        Right imagePath -> return $ Right $ createImageBlock imagePath
processMermaidBlock _ block = return $ Right block

-- | Generate file paths for diagram based on format
generateDiagramPaths :: DiagramConfig -> DiagramId -> (FilePath, Text)
generateDiagramPaths (DiagramConfig (SourceDir sourceDir) format) (DiagramId diagId) =
  let diagIdStr = T.unpack diagId
  in case format of
       HTML -> let svgFile = sourceDir </> diagIdStr <> ".svg"
               in (svgFile, T.pack $ takeFileName svgFile)
       PDF  -> let pngFile = sourceDir </> diagIdStr <> ".png"
               in (pngFile, T.pack pngFile)

-- | Render Mermaid diagram to appropriate format
renderMermaidDiagram :: DiagramConfig -> DiagramId -> Text -> IO (Either DocsterError Text)
renderMermaidDiagram config@(DiagramConfig (SourceDir sourceDir) format) diagId contents = do
  let diagIdStr = T.unpack $ (\(DiagramId d) -> d) diagId
      mmdFile = sourceDir </> diagIdStr <> ".mmd"
      (outputFile, imagePath) = generateDiagramPaths config diagId
  
  -- Use bracket to ensure cleanup of temporary mermaid file
  result <- bracket
    (TIO.writeFile mmdFile contents >> return mmdFile)
    (\file -> removeFile file `catch` \(_ :: SomeException) -> return ())
    (\_ -> do
      processResult <- callMermaidProcess format mmdFile outputFile
      case processResult of
        Left err -> return $ Left err
        Right _ -> do
          putStrLn $ T.unpack $ successEmoji <> " Generated " <> T.pack outputFile
          return $ Right imagePath)
  return result

-- | Call mermaid process with appropriate arguments
callMermaidProcess :: OutputFormat -> FilePath -> FilePath -> IO (Either DocsterError ())
callMermaidProcess format mmdFile outputFile = do
  let args = case format of
        HTML -> ["-i", mmdFile, "-o", outputFile]
        PDF  -> ["-i", mmdFile, "-o", outputFile, "--scale", "3"]
  
  result <- catch
    (callProcess (T.unpack mermaidCommand) args >> return (Right ()))
    (\(e :: SomeException) -> return $ Left $ ProcessError $ "Mermaid process failed: " <> T.pack (show e))
  return result

-- | Create Pandoc image block
createImageBlock :: Text -> Block
createImageBlock imagePath = Para [Image nullAttr [] (imagePath, "Mermaid diagram")]

-- | Walk the Pandoc AST and process blocks using walkM with proper error handling
transformDocument :: DiagramConfig -> Pandoc -> IO (Either DocsterError Pandoc)
transformDocument config doc = do
  result <- walkMEither (processMermaidBlock config) doc
  return result

-- | Walk with error handling - transforms Either into IO Either
walkMEither :: Monad m => (Block -> m (Either e Block)) -> Pandoc -> m (Either e Pandoc)
walkMEither f (Pandoc meta blocks) = do
  results <- mapM f blocks
  case sequence results of
    Left err -> return $ Left err
    Right newBlocks -> return $ Right $ Pandoc meta newBlocks

-- | LaTeX template with comprehensive package support
latexTemplate :: Text -> Text
latexTemplate bodyContent = T.unlines
  [ "\\documentclass{article}"
  , "\\usepackage[utf8]{inputenc}"
  , "\\usepackage{fontspec}"
  , "\\usepackage{graphicx}"
  , "\\usepackage{adjustbox}"
  , "\\usepackage{geometry}"
  , "\\geometry{margin=1in}"
  , "\\usepackage{hyperref}"
  , "\\usepackage{enumitem}"
  , "\\usepackage{amsmath}"
  , "\\usepackage{amssymb}"
  , "\\usepackage{fancyvrb}"
  , "\\usepackage{color}"
  , "\\DefineVerbatimEnvironment{Highlighting}{Verbatim}{commandchars=\\\\\\{\\}}"
  , "\\newenvironment{Shaded}{}{}"
  , syntaxHighlightingCommands
  , "\\providecommand{\\tightlist}{%"
  , "  \\setlength{\\itemsep}{0pt}\\setlength{\\parskip}{0pt}}"
  , "% Auto-scale oversized images to fit page"
  , "\\makeatletter"
  , "\\def\\maxwidth{\\ifdim\\Gin@nat@width>\\linewidth\\linewidth\\else\\Gin@nat@width\\fi}"
  , "\\def\\maxheight{\\ifdim\\Gin@nat@height>\\textheight\\textheight\\else\\Gin@nat@height\\fi}"
  , "\\makeatother"
  , "\\setkeys{Gin}{width=\\maxwidth,height=\\maxheight,keepaspectratio}"
  , "\\begin{document}"
  , bodyContent
  , "\\end{document}"
  ]

-- | Syntax highlighting commands for LaTeX
syntaxHighlightingCommands :: Text
syntaxHighlightingCommands = T.unlines
  [ "\\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}}}}"
  ]

-- | Parse command line arguments
parseArgs :: [String] -> Either DocsterError (IO ())
parseArgs ["-pdf", path] = Right (compileToPDF path)
parseArgs ["-html", path] = Right (compileToHTML path)
parseArgs _ = Left $ InvalidUsage "Usage: docster -pdf|-html <file.md>"

-- | Compile markdown to PDF using XeLaTeX
compileToPDF :: FilePath -> IO ()
compileToPDF path = do
  let sourceDir = SourceDir $ takeDirectory path
      outputPath = OutputPath $ replaceExtension path "pdf"
  
  result <- compileToPDFSafe sourceDir (OutputPath path) outputPath
  case result of
    Left err -> throwIO err
    Right _ -> return ()

-- | Safe PDF compilation with proper error handling
compileToPDFSafe :: SourceDir -> OutputPath -> OutputPath -> IO (Either DocsterError ())
compileToPDFSafe sourceDir (OutputPath inputPath) (OutputPath outputPathStr) = do
  content <- TIO.readFile inputPath
  let readerOptions = def { readerExtensions = getDefaultExtensions "markdown" }
      config = DiagramConfig sourceDir PDF
  
  pandocResult <- runIO $ readMarkdown readerOptions content
  case pandocResult of
    Left err -> return $ Left $ FileError $ "Failed to parse markdown: " <> T.pack (show err)
    Right pandoc -> do
      transformResult <- transformDocument config pandoc
      case transformResult of
        Left err -> return $ Left err
        Right transformed -> do
          latexResult <- runIO $ writeLaTeX def transformed
          case latexResult of
            Left err -> return $ Left $ PDFGenerationError $ "LaTeX generation failed: " <> T.pack (show err)
            Right latexOutput -> do
              let completeLatex = latexTemplate latexOutput
              pdfResult <- runIO $ makePDF "xelatex" [] (\_ _ -> return completeLatex) def transformed
              case pdfResult of
                Left err -> return $ Left $ PDFGenerationError $ T.pack $ show err
                Right (Left err) -> return $ Left $ PDFGenerationError $ T.pack $ show err
                Right (Right bs) -> do
                  BL.writeFile outputPathStr bs
                  putStrLn $ T.unpack $ successEmoji <> " PDF written to " <> T.pack outputPathStr
                  return $ Right ()

-- | Compile markdown to HTML
compileToHTML :: FilePath -> IO ()
compileToHTML path = do
  let sourceDir = SourceDir $ takeDirectory path
      outputPath = OutputPath $ replaceExtension path "html"
  
  result <- compileToHTMLSafe sourceDir (OutputPath path) outputPath
  case result of
    Left err -> throwIO err
    Right _ -> return ()

-- | Safe HTML compilation with proper error handling
compileToHTMLSafe :: SourceDir -> OutputPath -> OutputPath -> IO (Either DocsterError ())
compileToHTMLSafe sourceDir (OutputPath inputPath) (OutputPath outputPathStr) = do
  content <- TIO.readFile inputPath
  let readerOptions = def { readerExtensions = getDefaultExtensions "markdown" }
      config = DiagramConfig sourceDir HTML
  
  pandocResult <- runIO $ readMarkdown readerOptions content
  case pandocResult of
    Left err -> return $ Left $ FileError $ "Failed to parse markdown: " <> T.pack (show err)
    Right pandoc -> do
      transformResult <- transformDocument config pandoc
      case transformResult of
        Left err -> return $ Left err
        Right transformed -> do
          htmlResult <- runIO $ writeHtml5String def transformed
          case htmlResult of
            Left err -> return $ Left $ FileError $ "HTML generation failed: " <> T.pack (show err)
            Right html -> do
              TIO.writeFile outputPathStr html
              putStrLn $ T.unpack $ successEmoji <> " HTML written to " <> T.pack outputPathStr
              
              -- Open the generated HTML file in browser
              putStrLn $ "🌐 Opening " <> outputPathStr <> " in browser for error checking..."
              void $ callProcess "open" [outputPathStr]
              
              return $ Right ()

-- | Main entry point
main :: IO ()
main = do
  args <- getArgs
  case parseArgs args of
    Left err -> throwIO err
    Right action -> action