{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- | Prompt templating system for Ava agents.
--
-- Prompts are stored as markdown files with mustache templating in
-- @$AVA_DATA_ROOT/prompts/@. They support:
--
-- - Variables: @{{var_name}}@
-- - Sections: @{{#maybe_field}}...{{/maybe_field}}@
-- - Partials: @{{> shared/memory}}@
--
-- Directory structure:
--   prompts/
--   ├── shared/           -- Reusable partials
--   │   ├── memory.md
--   │   └── formatting/
--   │       └── telegram.md
--   ├── agents/
--   │   └── telegram/
--   │       └── system.md
--   └── subagents/
--       ├── generic/
--       │   └── system.md
--       └── coder/
--           └── system.md
--
-- : out omni-agent-prompts
-- : dep aeson
-- : dep directory
-- : dep mustache
-- : dep text
module Omni.Agent.Prompts
  ( -- * Types
    PromptMetadata (..),
    PromptTemplate (..),

    -- * Loading
    loadPrompt,
    loadPromptPure,

    -- * Rendering
    renderPrompt,
    renderPromptWith,

    -- * Listing
    listPrompts,

    -- * Paths
    promptsDir,
    promptPath,

    -- * Testing
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.=))
import qualified Data.Aeson as Aeson
import qualified Data.List as List
import qualified Data.Text as Text
import qualified Data.Text.IO as TextIO
import qualified Omni.Agent.Paths as Paths
import qualified Omni.Test as Test
import qualified System.Directory as Dir
import System.FilePath ((</>))
import qualified System.FilePath as FilePath
import qualified Text.Mustache as Mustache
import Text.Mustache.Types (Template)

main :: IO ()
main = Test.run test

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Prompts"
    [ Test.unit "promptPath constructs correct path" <| do
        let path = promptPath "agents/telegram/system"
        ("agents/telegram/system.mustache" `List.isSuffixOf` path) Test.@=? True,
      Test.unit "parsePromptMetadata parses frontmatter" <| do
        let content =
              "---\n\
              \name: test-prompt\n\
              \description: A test prompt\n\
              \---\n\
              \Hello {{name}}!"
        case splitFrontmatter content of
          (yaml, body) -> do
            ("name:" `Text.isInfixOf` yaml) Test.@=? True
            ("Hello" `Text.isInfixOf` body) Test.@=? True,
      Test.unit "simple mustache substitution works" <| do
        let tpl = "Hello {{name}}!"
        case Mustache.compileTemplate "test" tpl of
          Left _ -> Test.assertFailure "Failed to compile template"
          Right t -> do
            let result = Mustache.substitute t (Aeson.object ["name" .= ("World" :: Text)])
            result Test.@=? "Hello World!"
    ]

promptsDir :: FilePath
promptsDir = Paths.promptsDir

promptPath :: Text -> FilePath
promptPath pid = promptsDir </> Text.unpack pid FilePath.<.> "mustache"

data PromptMetadata = PromptMetadata
  { pmName :: Maybe Text,
    pmDescription :: Maybe Text,
    pmVars :: [Text]
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON PromptMetadata where
  toJSON m =
    Aeson.object
      [ "name" .= pmName m,
        "description" .= pmDescription m,
        "vars" .= pmVars m
      ]

data PromptTemplate = PromptTemplate
  { ptId :: Text,
    ptPath :: FilePath,
    ptMetadata :: PromptMetadata,
    ptTemplate :: Template
  }
  deriving (Show)

splitFrontmatter :: Text -> (Text, Text)
splitFrontmatter content =
  let stripped = Text.strip content
   in if Text.isPrefixOf "---" stripped
        then
          let afterFirst = Text.drop 3 stripped
              (yamlPart, rest) = Text.breakOn "---" (Text.stripStart afterFirst)
           in if Text.null rest
                then ("", content)
                else (Text.strip yamlPart, Text.strip (Text.drop 3 rest))
        else ("", content)

parsePromptMetadata :: Text -> PromptMetadata
parsePromptMetadata yaml =
  let kvPairs = parseKvLines (Text.lines yaml)
      getName = List.lookup "name" kvPairs
      getDesc = List.lookup "description" kvPairs
   in PromptMetadata
        { pmName = getName,
          pmDescription = getDesc,
          pmVars = []
        }
  where
    parseKvLines :: [Text] -> [(Text, Text)]
    parseKvLines = mapMaybe parseKvLine

    parseKvLine :: Text -> Maybe (Text, Text)
    parseKvLine line = do
      let (key, rest) = Text.breakOn ":" line
      guard (not (Text.null rest))
      let value = Text.strip (Text.drop 1 rest)
      guard (not (Text.null key))
      pure (Text.strip key, value)

loadPrompt :: Text -> IO (Either Text PromptTemplate)
loadPrompt pid = do
  let path = promptPath pid
  exists <- Dir.doesFileExist path
  if not exists
    then pure <| Left ("Prompt not found: " <> pid <> " (" <> Text.pack path <> ")")
    else do
      content <- TextIO.readFile path
      pure <| loadPromptPure pid path content

loadPromptPure :: Text -> FilePath -> Text -> Either Text PromptTemplate
loadPromptPure pid path content =
  let (yamlPart, body) = splitFrontmatter content
      meta = parsePromptMetadata yamlPart
   in case Mustache.compileTemplate (Text.unpack pid) body of
        Left err -> Left ("Mustache compile error: " <> Text.pack (show err))
        Right tpl ->
          Right
            PromptTemplate
              { ptId = pid,
                ptPath = path,
                ptMetadata = meta,
                ptTemplate = tpl
              }

renderPromptWith :: PromptTemplate -> Aeson.Value -> Text
renderPromptWith pt = Mustache.substitute (ptTemplate pt)

renderPrompt :: Text -> Aeson.Value -> IO (Either Text Text)
renderPrompt pid ctx = do
  let path = promptPath pid
      relativePath = Text.unpack pid FilePath.<.> "mustache"
  exists <- Dir.doesFileExist path
  if not exists
    then pure <| Left ("Prompt not found: " <> pid <> " (" <> Text.pack path <> ")")
    else do
      absPromptsDir <- Dir.makeAbsolute promptsDir
      eTpl <- Mustache.automaticCompile [absPromptsDir] relativePath
      case eTpl of
        Left err -> pure <| Left ("Mustache compile error: " <> Text.pack (show err))
        Right tpl -> pure <| Right <| Mustache.substitute tpl ctx

listPrompts :: IO [(Text, PromptMetadata)]
listPrompts = listPromptsInDir promptsDir ""

listPromptsInDir :: FilePath -> Text -> IO [(Text, PromptMetadata)]
listPromptsInDir baseDir prefix = do
  exists <- Dir.doesDirectoryExist baseDir
  if not exists
    then pure []
    else do
      entries <- Dir.listDirectory baseDir
      results <-
        forM entries <| \entry -> do
          let fullPath = baseDir </> entry
          isDir <- Dir.doesDirectoryExist fullPath
          if isDir
            then listPromptsInDir fullPath (if Text.null prefix then Text.pack entry else prefix <> "/" <> Text.pack entry)
            else
              if FilePath.takeExtension entry == ".mustache"
                then do
                  let name = FilePath.dropExtension entry
                      pid = if Text.null prefix then Text.pack name else prefix <> "/" <> Text.pack name
                  content <- TextIO.readFile fullPath
                  let (yamlPart, _) = splitFrontmatter content
                      meta = parsePromptMetadata yamlPart
                  pure [(pid, meta)]
                else pure []
      pure (concat results)