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

-- | CLI tool for previewing and debugging prompt templates.
--
-- Usage:
--   omni-agent-prompt list
--   omni-agent-prompt agents/telegram/system
--   omni-agent-prompt subagents/coder/system --var namespace=Omni/Jr --var task="Add feature"
--   omni-agent-prompt subagents/generic/system --json '{"task":"Research X","role_description":"research"}'
--
-- : out omni-agent-prompt
-- : dep aeson
-- : dep bytestring
-- : dep text
-- : dep mustache
-- : dep unordered-containers
module Omni.Agent.Prompts.Cli (main) where

import Alpha
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Key as Key
import qualified Data.Aeson.KeyMap as KeyMap
import qualified Data.ByteString.Lazy as BL
import qualified Data.Text as Text
import qualified Data.Text.Encoding as TE
import qualified Data.Text.IO as TextIO
import qualified Omni.Agent.Prompts as Prompts
import qualified System.Environment as Env
import qualified System.Exit as Exit

main :: IO ()
main = do
  args <- Env.getArgs
  case args of
    [] -> usage
    ["--help"] -> usage
    ["-h"] -> usage
    ["list"] -> listPrompts
    (pid : rest) -> renderPromptCmd (Text.pack pid) rest

usage :: IO ()
usage = do
  putText "omni-agent-prompt - Preview and debug prompt templates"
  putText ""
  putText "Usage:"
  putText "  omni-agent-prompt list"
  putText "  omni-agent-prompt <prompt-id> [--var key=value]... [--json '{...}']"
  putText ""
  putText "Examples:"
  putText "  omni-agent-prompt agents/telegram/system"
  putText "  omni-agent-prompt subagents/coder/system --var namespace=Omni/Jr --var task='Add feature'"
  putText "  omni-agent-prompt subagents/generic/system --json '{\"task\":\"Research X\"}'"
  putText ""
  putText <| "Prompts directory: " <> Text.pack Prompts.promptsDir

listPrompts :: IO ()
listPrompts = do
  prompts <- Prompts.listPrompts
  if null prompts
    then do
      putText "No prompts found."
      putText <| "Prompts directory: " <> Text.pack Prompts.promptsDir
    else do
      putText "Available prompts:"
      putText ""
      forM_ prompts <| \(pid, meta) -> do
        let desc = fromMaybe "" (Prompts.pmDescription meta)
        if Text.null desc
          then putText <| "  " <> pid
          else putText <| "  " <> pid <> " - " <> desc

renderPromptCmd :: Text -> [String] -> IO ()
renderPromptCmd pid args = do
  let (jsonCtx, kvPairs) = parseArgs args
  ctx <- case jsonCtx of
    Nothing -> pure <| buildContext kvPairs
    Just jsonStr -> case Aeson.decode (BL.fromStrict (TE.encodeUtf8 jsonStr)) of
      Nothing -> do
        putText <| "Error: Invalid JSON: " <> jsonStr
        Exit.exitFailure
      Just obj -> pure <| mergeContext obj kvPairs

  result <- Prompts.renderPrompt pid ctx
  case result of
    Left err -> do
      putText <| "Error: " <> err
      Exit.exitFailure
    Right rendered -> TextIO.putStrLn rendered

parseArgs :: [String] -> (Maybe Text, [(Text, Text)])
parseArgs = go Nothing []
  where
    go json kvs [] = (json, reverse kvs)
    go _json kvs ("--json" : val : rest) = go (Just (Text.pack val)) kvs rest
    go json kvs ("--var" : kv : rest) =
      case Text.breakOn "=" (Text.pack kv) of
        (k, v)
          | not (Text.null v) -> go json ((k, Text.drop 1 v) : kvs) rest
          | otherwise -> go json kvs rest
    go json kvs (_ : rest) = go json kvs rest

buildContext :: [(Text, Text)] -> Aeson.Value
buildContext kvs =
  Aeson.Object <| KeyMap.fromList [(Key.fromText k, Aeson.String v) | (k, v) <- kvs]

mergeContext :: Aeson.Value -> [(Text, Text)] -> Aeson.Value
mergeContext (Aeson.Object obj) kvs =
  Aeson.Object <| foldr (\(k, v) m -> KeyMap.insert (Key.fromText k) (Aeson.String v) m) obj kvs
mergeContext _ kvs = buildContext kvs