{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}

{- | Module for printing session files in markdown format.

This module provides functionality to load and format session JSON files
as readable markdown output. It is used by the @session-print@ CLI command.
-}
module System.Agents.SessionPrint (
    -- * Options
    SessionPrintOptions (..),
    OrderPreference (..),
    PrintVisibility (..),
    PrintAmount (..),

    -- * Main handler
    handleSessionPrint,

    -- * Formatting functions
    formatSessionAsMarkdown,

    -- * Statistics
    SessionStatistics (..),
    calculateStatistics,

    -- * Elision
    elideDocument,
) where

import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Encode.Pretty as AesonPretty
import qualified Data.Aeson.KeyMap as KeyMap
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import Data.List (intersperse, sortOn)
import Data.Map (Map)
import qualified Data.Map as Map
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Text
import qualified Data.Text.IO as Text
import qualified Data.Text.Lazy as TextLazy
import qualified Data.Text.Lazy.Encoding as TextLazyEncoding
import System.IO (stderr)

import System.Agents.LLMs.OpenAI (TokenUsage (..))
import System.Agents.Media.Types (ContentPart (..), MediaAttachment (..))
import qualified System.Agents.Session.Base as Session
import System.Agents.Session.Signals (
    calculateTrajectorySignals,
 )
import System.Agents.Session.Types (
    EnvironmentSignals (..),
    ExecutionSignals (..),
    InteractionSignals (..),
    StepByteUsage (..),
    TrajectorySignals (..),
    UserToolResponse (..),
 )

-- | Preference for ordering session steps.
data OrderPreference
    = -- | Oldest first (earlier steps shown first)
      Chronological
    | -- | Newest first (recent steps shown first)
      Antichronological
    deriving (Show, Eq)

-- | Amount of content to print (either in lines or characters).
data PrintAmount
    = -- | Number of lines to keep
      Lines Int
    | -- | Number of characters to keep
      Chars Int
    deriving (Show, Eq)

{- | Visibility preference for displaying content.

- 'Hidden': Don't show the content at all
- 'Elided leading trailing': Show leading and trailing portions, eliding the middle
- 'ShownFull': Show the complete content
-}
data PrintVisibility
    = Hidden
    | -- | Keep leading amount and trailing amount
      Elided PrintAmount PrintAmount
    | ShownFull
    deriving (Show, Eq)

-- | Options for controlling the session print output.
data SessionPrintOptions = SessionPrintOptions
    { sessionPrintFile :: FilePath
    -- ^ Path to the session JSON file to print
    , showToolCallResults :: PrintVisibility
    -- ^ How to display tool call results in the output
    , showToolCallArguments :: PrintVisibility
    -- ^ How to display tool call arguments in the output
    , nTurns :: Maybe Int
    -- ^ Optional limit on the number of turns to display
    , repeatSystemPrompt :: Bool
    -- ^ Whether to repeat the system prompt at each turn
    , repeatTools :: Bool
    -- ^ Whether to repeat the available tools at each turn
    , orderPreference :: OrderPreference
    -- ^ Order preference for displaying session steps
    , noFunnyStamp :: Bool
    -- ^ Whether to skip the ASCII art logo stamp in the header
    }

-- | Token usage statistics aggregated across a session.
data TokenUsageStats = TokenUsageStats
    { statPromptTokens :: Int
    -- ^ Tokens in the prompt/context (input)
    , statCompletionTokens :: Int
    -- ^ Tokens in the completion response (output)
    , statTotalTokens :: Int
    -- ^ Total tokens used
    , statCachedTokens :: Maybe Int
    -- ^ Cached tokens (if provider supports it)
    , statThinkingTokens :: Maybe Int
    -- ^ Reasoning/thinking tokens (if model supports it)
    }
    deriving (Show, Eq)

-- | Statistics about a session.
data SessionStatistics = SessionStatistics
    { statTotalTurns :: Int
    , statUserTurns :: Int
    , statLlmTurns :: Int
    , statTotalToolCalls :: Int
    , statToolCallsByName :: Map Text.Text Int
    , statInputBytes :: Int
    , statOutputBytes :: Int
    , statReasoningBytes :: Int
    , statTotalBytes :: Int
    , statTokenUsage :: Maybe TokenUsageStats
    -- ^ Token usage statistics if available
    , statTrajectorySignals :: Maybe TrajectorySignals
    -- ^ Signal-based trajectory metrics
    }
    deriving (Show, Eq)

-- | Default ASCII art logo for the session report header.
defaultLogo :: Text.Text
defaultLogo =
    Text.intercalate
        "\n"
        [ "```"
        , "    λλ"
        , "   λ⊙λ⊙"
        , "  λλλλλλ"
        , " λλλλλλλλ"
        , "    λλ"
        , "    λλ"
        , "    λλ"
        , "    λλ"
        , "    λλ"
        , "    λλ"
        , "    λλ"
        , "    λλ"
        , "```"
        ]

-- | Handle the session-print command: load a session file and output it as markdown.
handleSessionPrint :: SessionPrintOptions -> IO ()
handleSessionPrint opts = do
    result <- Aeson.eitherDecodeFileStrict' opts.sessionPrintFile
    case result of
        Left err -> do
            Text.hPutStrLn stderr $ "Error loading session file: " <> Text.pack err
        Right session -> do
            let markdown = formatSessionAsMarkdown opts session
            Text.putStr markdown

{- | Elide a document by keeping leading and trailing portions.

This function intelligently handles overlapping regions. For example, if a
document is 7 lines long and we request @Elided (Lines 5) (Lines 5)@, all
7 lines will be displayed (not duplicated).

>>> elideDocument (Lines 3) (Lines 3) "line1\nline2\nline3\nline4\nline5\nline6\nline7"
"line1\nline2\nline3\nline4\nline5\nline6\nline7"

>>> elideDocument (Lines 2) (Lines 2) "line1\nline2\nline3\nline4\nline5"
"line1\nline2\n... (1 line elided) ...\nline4\nline5"

>>> elideDocument (Chars 10) (Chars 5) "Hello, world! This is a test."
"Hello, wor... (elided 14 chars) ...test."

>>> elideDocument (Lines 100) (Lines 100) "short text"
"short text"
-}
elideDocument :: PrintAmount -> PrintAmount -> Text.Text -> Text.Text
elideDocument _ _ "" = ""
elideDocument leading trailing content =
    case (leading, trailing) of
        (Lines n, Lines m) -> elideLines n m content
        (Chars n, Chars m) -> elideChars n m content
        (Lines n, Chars m) ->
            -- Convert leading lines to chars for consistent handling
            let lines' = Text.lines content
                leadingLines = take n lines'
                leadingChars = sum (map Text.length leadingLines) + length leadingLines
             in elideChars leadingChars m content
        (Chars n, Lines m) ->
            -- Convert trailing lines to chars for consistent handling
            let lines' = Text.lines content
                trailingLines = take m (reverse lines')
                trailingChars = sum (map Text.length trailingLines) + length trailingLines
             in elideChars n trailingChars content

-- | Elide by lines, handling overlaps intelligently.
elideLines :: Int -> Int -> Text.Text -> Text.Text
elideLines leadingCount trailingCount content =
    let lines' = Text.lines content
        totalLines = length lines'
        -- If the total would exceed document length, show everything
        actualLeading = min leadingCount totalLines
        actualTrailing = min trailingCount totalLines
        -- Check for overlap
        overlap = actualLeading + actualTrailing - totalLines
     in if overlap >= 0
            then Text.intercalate "\n" lines' -- No elision needed
            else
                let leading = take actualLeading lines'
                    trailing = drop (totalLines - actualTrailing) lines'
                    elidedCount = totalLines - actualLeading - actualTrailing
                 in Text.intercalate "\n" leading
                        <> "\n... ("
                        <> Text.pack (show elidedCount)
                        <> " line"
                        <> (if elidedCount == 1 then "" else "s")
                        <> " elided) ...\n"
                        <> Text.intercalate "\n" trailing

-- | Elide by characters, handling overlaps intelligently.
elideChars :: Int -> Int -> Text.Text -> Text.Text
elideChars leadingCount trailingCount content =
    let totalChars = Text.length content
        -- If the total would exceed document length, show everything
        actualLeading = min leadingCount totalChars
        actualTrailing = min trailingCount totalChars
        -- Check for overlap
        overlap = actualLeading + actualTrailing - totalChars
     in if overlap >= 0
            then content -- No elision needed
            else
                let leading = Text.take actualLeading content
                    trailing = Text.drop (totalChars - actualTrailing) content
                    elidedCount = totalChars - actualLeading - actualTrailing
                 in leading
                        <> "... (elided "
                        <> Text.pack (show elidedCount)
                        <> " char"
                        <> (if elidedCount == 1 then "" else "s")
                        <> ") ..."
                        <> trailing

{- | Format a Session as markdown text.

The logic for processing steps is:
1. First, number the steps (zip with [1..]) - preserves original step numbers
2. Then, cut the n-latest steps if n-turns limit is specified
3. Then, reverse or not depending on the order preference
4. Finally, traverse and print
-}
formatSessionAsMarkdown :: SessionPrintOptions -> Session.Session -> Text.Text
formatSessionAsMarkdown opts session =
    let
        -- Session turns are stored newest-first (last turn is first in the list)
        -- Step 1: Number the steps with their original position
        -- We reverse first to get chronological order for numbering, then zip
        numberedChronological = zip [(1 :: Int) ..] (reverse session.turns)

        -- Step 2: Apply n-turns limit (takes the last N steps in chronological order = most recent)
        limitedTurns = case opts.nTurns of
            Just n -> takeLast n numberedChronological
            Nothing -> numberedChronological

        -- Step 3: Apply ordering preference
        orderedTurns = case opts.orderPreference of
            Chronological -> limitedTurns -- Already in chronological order
            Antichronological -> reverse limitedTurns -- Newest first

        -- Determine the first step number in display order for "is first turn" checks
        firstDisplayStepNum = case orderedTurns of
            [] -> Nothing
            ((n, _) : _) -> Just n

        -- Calculate statistics for the full session
        stats = calculateStatistics session

        mdHeader = formatHeader opts.noFunnyStamp
        sessionInfo = formatSessionInfo session
        statsSection = formatStatistics stats
        turnsSection = formatTurns opts firstDisplayStepNum orderedTurns
        limitNotice = case opts.nTurns of
            Just n -> "\n\n_(Showing last " <> Text.pack (show n) <> " turns)_\n"
            Nothing -> ""
     in
        mdHeader <> sessionInfo <> "\n---\n\n" <> statsSection <> "\n---\n\n" <> turnsSection <> limitNotice
  where
    -- \| Take the last n elements from a list
    takeLast :: Int -> [a] -> [a]
    takeLast n xs = drop (max 0 (length xs - n)) xs

-- | Format the header with optional ASCII art logo.
formatHeader :: Bool -> Text.Text
formatHeader skipLogo =
    if skipLogo
        then "# 📋 Session Report\n\n"
        else Text.intercalate "\n" (zipWith combineLogoLine logoLines titleLines) <> "\n\n"
  where
    logoLines = Text.lines defaultLogo
    titleLines = "# 📋 Session Report" : repeat ""

    -- Combine a logo line with a title line, padding the logo to align nicely
    combineLogoLine :: Text.Text -> Text.Text -> Text.Text
    combineLogoLine logoLine titleLine =
        let paddedLogo = Text.justifyLeft 20 ' ' logoLine
         in if Text.null titleLine
                then paddedLogo
                else paddedLogo <> "  " <> titleLine

-- | Calculate statistics for a session.
calculateStatistics :: Session.Session -> SessionStatistics
calculateStatistics session =
    let turns = session.turns
        totalTurns = length turns
        userTurns = length [() | Session.UserTurn _ _ <- turns]
        llmTurns = length [() | Session.LlmTurn _ _ <- turns]

        -- Collect all LLM turn contents to extract tool calls
        llmTurnContents = [tc | Session.LlmTurn tc _ <- turns]
        allToolCalls = concatMap Session.llmToolCalls llmTurnContents
        toolCallCount = length allToolCalls
        toolCallsByName = countToolCallsByName allToolCalls

        -- Calculate byte counts
        (inputBytes, outputBytes, reasoningBytes) = calculateByteCounts turns
        totalBytes = inputBytes + outputBytes + reasoningBytes

        -- Calculate token usage from StepByteUsage data
        tokenStats = calculateTokenUsageStats turns

        -- Calculate trajectory signals
        signals = calculateTrajectorySignals session
     in SessionStatistics
            { statTotalTurns = totalTurns
            , statUserTurns = userTurns
            , statLlmTurns = llmTurns
            , statTotalToolCalls = toolCallCount
            , statToolCallsByName = toolCallsByName
            , statInputBytes = inputBytes
            , statOutputBytes = outputBytes
            , statReasoningBytes = reasoningBytes
            , statTotalBytes = totalBytes
            , statTokenUsage = tokenStats
            , statTrajectorySignals = Just signals
            }

-- | Calculate aggregated token usage statistics from all turns.
calculateTokenUsageStats :: [Session.Turn] -> Maybe TokenUsageStats
calculateTokenUsageStats turns =
    let allTokenUsages = collectTokenUsages turns
     in if null allTokenUsages
            then Nothing
            else Just $ aggregateTokenUsages allTokenUsages
  where
    collectTokenUsages :: [Session.Turn] -> [TokenUsage]
    collectTokenUsages ts =
        [ usage
        | turn <- ts
        , usage <- extractUsageFromTurn turn
        ]

    -- Extract token usage from a turn's StepByteUsage if present
    extractUsageFromTurn :: Session.Turn -> [TokenUsage]
    extractUsageFromTurn turn =
        let mUsage = case turn of
                Session.UserTurn _ mStepUsage -> mStepUsage
                Session.LlmTurn _ mStepUsage -> mStepUsage
                Session.PartialUserTurn _ mStepUsage -> mStepUsage
         in case mUsage of
                Just stepUsage -> maybeToList (stepTokenUsage stepUsage)
                Nothing -> []

    maybeToList :: Maybe a -> [a]
    maybeToList Nothing = []
    maybeToList (Just x) = [x]

    aggregateTokenUsages :: [TokenUsage] -> TokenUsageStats
    aggregateTokenUsages usages =
        TokenUsageStats
            { statPromptTokens = sum $ map tokenPromptTokens usages
            , statCompletionTokens = sum $ map tokenCompletionTokens usages
            , statTotalTokens = sum $ map tokenTotalTokens usages
            , statCachedTokens = sumMaybe $ map tokenCachedTokens usages
            , statThinkingTokens = sumMaybe $ map tokenThinkingTokens usages
            }

    sumMaybe :: [Maybe Int] -> Maybe Int
    sumMaybe ms =
        let vs = [v | Just v <- ms]
         in if null vs then Nothing else Just (sum vs)

-- | Count tool calls by name.
countToolCallsByName :: [Session.LlmToolCall] -> Map Text.Text Int
countToolCallsByName toolCalls =
    let toolNames = map extractToolCallName toolCalls
     in Map.fromListWith (+) [(name, 1) | name <- toolNames, name /= "(unnamed)"]

-- | Calculate byte counts for input, output, and reasoning.
calculateByteCounts :: [Session.Turn] -> (Int, Int, Int)
calculateByteCounts turns =
    let (inputBytes, outputBytes, reasoningBytes) = foldr countTurn (0, 0, 0) turns
     in (inputBytes, outputBytes, reasoningBytes)
  where
    countTurn :: Session.Turn -> (Int, Int, Int) -> (Int, Int, Int)
    countTurn (Session.UserTurn utc _mUsage) (inp, out, reas) =
        let userBytes = textBytes $ case utc.userPrompt of
                Session.SystemPrompt sp -> sp
            queryBytes = maybe 0 (textBytes . unwrapQuery) utc.userQuery
            toolRespBytes = sum [userToolResponseBytes r | (_, r) <- utc.userToolResponses]
         in (inp + userBytes + queryBytes + toolRespBytes, out, reas)
    countTurn (Session.LlmTurn ltc _mUsage) (inp, out, reas) =
        let response = ltc.llmResponse
            outBytes = maybe 0 textBytes response.responseText
            reasBytes = maybe 0 textBytes response.responseThinking
         in (inp, out + outBytes, reas + reasBytes)
    countTurn (Session.PartialUserTurn _ _) (inp, out, reas) =
        -- Partial user turns don't contribute to byte counts
        (inp, out, reas)
    textBytes :: Text.Text -> Int
    textBytes = BS.length . Text.encodeUtf8

    unwrapQuery :: Session.UserQuery -> Text.Text
    unwrapQuery (Session.UserQuery q _) = q

    -- Calculate bytes for a UserToolResponse
    userToolResponseBytes :: UserToolResponse -> Int
    userToolResponseBytes (TextResponse txt) = Text.length txt * 4
    userToolResponseBytes (JsonResponse val) = fromIntegral (BSL.length (Aeson.encode val))
    userToolResponseBytes (MediaResponse media) = Text.length (media.mediaBase64Data)
    userToolResponseBytes (MixedResponse parts) = sum (map contentPartBytes parts)

    -- Calculate bytes for a ContentPart
    contentPartBytes :: ContentPart -> Int
    contentPartBytes (TextPart txt) = Text.length txt * 4
    contentPartBytes (MediaPart media) = Text.length (media.mediaBase64Data)

-- | Format statistics as markdown.
formatStatistics :: SessionStatistics -> Text.Text
formatStatistics stats =
    "## 📊 Statistics\n\n"
        <> "### Turn Counts\n\n"
        <> "| Metric | Value |\n"
        <> "|--------|-------|\n"
        <> "| Total Turns | "
        <> Text.pack (show (statTotalTurns stats :: Int))
        <> " |\n"
        <> "| 👤 User Turns | "
        <> Text.pack (show (statUserTurns stats :: Int))
        <> " |\n"
        <> "| 🤖 LLM Turns | "
        <> Text.pack (show (statLlmTurns stats :: Int))
        <> " |\n"
        <> "\n### 🔧 Tool Calls\n\n"
        <> "**Total Tool Calls:** "
        <> Text.pack (show (statTotalToolCalls stats :: Int))
        <> "\n\n"
        <> formatToolCallStats stats.statToolCallsByName
        <> "\n### 💾 Byte Usage\n\n"
        <> formatByteChart stats
        <> case statTokenUsage stats of
            Nothing -> ""
            Just tokenStats -> "\n\n### 🎫 Token Usage\n\n" <> formatTokenChart tokenStats
        <> case statTrajectorySignals stats of
            Nothing -> ""
            Just signals -> "\n\n### 📡 Signal Metrics\n\n" <> formatSignalMetrics signals

-- | Format signal metrics as markdown.
formatSignalMetrics :: TrajectorySignals -> Text.Text
formatSignalMetrics signals =
    let is = trajInteraction signals
        es = trajExecution signals
        env = trajEnvironment signals
        score = trajInformativenessScore signals
     in "**Informativeness Score:** "
            <> showScore score
            <> "\n\n"
            <> "#### Interaction Signals\n\n"
            <> formatInteractionSignals is
            <> "\n\n#### Execution Signals\n\n"
            <> formatExecutionSignals es
            <> "\n\n#### Environment Signals\n\n"
            <> formatEnvironmentSignals env

-- | Show informativeness score with visual indicator.
showScore :: Int -> Text.Text
showScore score
    | score >= 70 = Text.pack (show score) <> "/100 🟢"
    | score >= 40 = Text.pack (show score) <> "/100 🟡"
    | otherwise = Text.pack (show score) <> "/100 🔴"

-- | Format interaction signals as markdown table.
formatInteractionSignals :: InteractionSignals -> Text.Text
formatInteractionSignals is =
    "| Signal | Status | Details |\n"
        <> "|--------|--------|---------|\n"
        <> "| Misalignment | "
        <> showCount (sigMisalignmentCount is)
        <> " | "
        <> if sigMisalignmentCount is > 0
            then Text.pack (show (sigMisalignmentCount is)) <> " instances detected"
            else
                "None detected"
                    <> " |\n"
                    <> "| Stagnation | "
                    <> showCount (sigStagnationCount is)
                    <> " | "
                    <> if sigStagnationCount is > 0
                        then Text.pack (show (sigStagnationCount is)) <> " instances detected"
                        else
                            "None detected"
                                <> " |\n"
                                <> "| Disengagement | "
                                <> showBool (sigDisengagementDetected is)
                                <> " | "
                                <> if sigDisengagementDetected is
                                    then "⚠️ Exit/withdrawal markers found"
                                    else
                                        "Not detected"
                                            <> " |\n"
                                            <> "| Satisfaction | "
                                            <> showBool (sigSatisfactionDetected is)
                                            <> " | "
                                            <> if sigSatisfactionDetected is
                                                then "✓ Success markers found"
                                                else
                                                    "Not detected"
                                                        <> " |"

-- | Format execution signals as markdown table.
formatExecutionSignals :: ExecutionSignals -> Text.Text
formatExecutionSignals es =
    "| Signal | Status | Details |\n"
        <> "|--------|--------|---------|\n"
        <> "| Failures | "
        <> showCount (sigFailureCount es)
        <> " | "
        <> if sigFailureCount es > 0
            then Text.pack (show (sigFailureCount es)) <> " tool/action failures"
            else
                "No failures detected"
                    <> " |\n"
                    <> "| Loop | "
                    <> showBool (sigLoopDetected es)
                    <> " | "
                    <> if sigLoopDetected es
                        then "🔴 Loop detected: " <> Text.intercalate " → " (sigLoopToolSequence es)
                        else
                            "No loops detected"
                                <> " |"

-- | Format environment signals as markdown table.
formatEnvironmentSignals :: EnvironmentSignals -> Text.Text
formatEnvironmentSignals env =
    "| Signal | Status | Details |\n"
        <> "|--------|--------|---------|\n"
        <> "| Exhaustion | "
        <> showCount (sigExhaustionCount env)
        <> " | "
        <> if sigExhaustionCount env > 0
            then
                Text.pack (show (sigExhaustionCount env))
                    <> " resource boundary hits: "
                    <> Text.intercalate ", " (sigExhaustionTypes env)
            else
                "No resource exhaustion detected"
                    <> " |"

-- | Show count with visual indicator.
showCount :: Int -> Text.Text
showCount 0 = "✓ 0"
showCount n = "⚠️ " <> Text.pack (show n)

-- | Show boolean status.
showBool :: Bool -> Text.Text
showBool True = "⚠️ Yes"
showBool False = "✓ No"

-- | Format tool call statistics with bar chart.
formatToolCallStats :: Map Text.Text Int -> Text.Text
formatToolCallStats toolMap
    | Map.null toolMap = "_No tool calls recorded_\n"
    | otherwise =
        let maxCount = maximum (Map.elems toolMap)
            sortedTools = sortOn (negate . snd) (Map.toList toolMap)
         in Text.intercalate "\n\n" $ map (formatToolBar maxCount) sortedTools <> [""]

-- | Format a single tool bar in the chart.
formatToolBar :: Int -> (Text.Text, Int) -> Text.Text
formatToolBar maxCount (name, count) =
    let barWidth = 60 -- Maximum width of the bar in characters
        filled = if maxCount == 0 then 0 else (count * barWidth) `div` maxCount
        bar = Text.replicate filled "█"
        label = Text.justifyLeft 20 ' ' name
        valueStr = Text.justifyRight 4 ' ' (Text.pack $ show count)
     in "`" <> label <> "` " <> valueStr <> " " <> bar

-- | Format byte usage as bar chart.
formatByteChart :: SessionStatistics -> Text.Text
formatByteChart stats =
    let categories =
            [ ("Input", statInputBytes stats)
            , ("Output", statOutputBytes stats)
            , ("Reasoning", statReasoningBytes stats)
            ]
        maxBytes = maximum (map snd categories)
        total = statTotalBytes stats
     in if total == 0
            then "_No byte data recorded_\n"
            else
                let chart = Text.intercalate "\n\n" $ map (formatByteBar maxBytes) categories
                    totalLine = "\n**Total:** " <> formatBytes total <> "\n"
                 in chart <> totalLine

-- | Format token usage as bar chart.
formatTokenChart :: TokenUsageStats -> Text.Text
formatTokenChart stats =
    let
        -- Build list of categories, only including non-zero ones
        allCategories =
            [ ("Prompt", statPromptTokens stats)
            , ("Completion", statCompletionTokens stats)
            ]
        categories = filter ((> 0) . snd) allCategories
        -- Add cached tokens if available
        categoriesWithCached = case statCachedTokens stats of
            Just n | n > 0 -> categories ++ [("Cached", n)]
            _ -> categories
        -- Add thinking tokens if available
        categoriesWithThinking = case statThinkingTokens stats of
            Just n | n > 0 -> categoriesWithCached ++ [("Thinking", n)]
            _ -> categoriesWithCached
        maxTokens = if null categoriesWithThinking then 1 else maximum (map snd categoriesWithThinking)
        total = statTotalTokens stats
     in
        if total == 0
            then "_No token data recorded_\n"
            else
                let chart = Text.intercalate "\n\n" $ map (formatTokenBar maxTokens) categoriesWithThinking
                    totalLine = "\n**Total:** " <> formatTokenCount total <> "\n"
                    cachedNote = case statCachedTokens stats of
                        Just n | n > 0 -> "\n_(" <> formatTokenCount n <> " cached)_\n"
                        _ -> ""
                    thinkingNote = case statThinkingTokens stats of
                        Just n | n > 0 -> "\n_(" <> formatTokenCount n <> " thinking tokens)_\n"
                        _ -> ""
                 in chart <> totalLine <> cachedNote <> thinkingNote

-- | Format a single token bar in the chart.
formatTokenBar :: Int -> (Text.Text, Int) -> Text.Text
formatTokenBar maxTokens (label, tokens) =
    let barWidth = 60
        filled = if maxTokens == 0 then 0 else (tokens * barWidth) `div` maxTokens
        bar = Text.replicate filled "█"
        label' = Text.justifyLeft 12 ' ' label
        valueStr = Text.justifyRight 10 ' ' (formatTokenCount tokens)
     in "`" <> label' <> "` " <> valueStr <> " " <> bar

-- | Format token count with thousands separator.
formatTokenCount :: Int -> Text.Text
formatTokenCount n =
    let txt = Text.pack (show n)
     in addThousandSeparators txt

-- | Add thousand separators to a numeric text.
addThousandSeparators :: Text.Text -> Text.Text
addThousandSeparators txt =
    let digits = Text.unpack txt
        grouped = reverse $ group3 (reverse digits)
     in Text.pack $ concat (intersperse "," grouped)
  where
    group3 :: String -> [String]
    group3 [] = []
    group3 s = take 3 s : group3 (drop 3 s)

-- | Format a single byte bar in the chart.
formatByteBar :: Int -> (Text.Text, Int) -> Text.Text
formatByteBar maxBytes (label, bytes) =
    let barWidth = 60
        filled = if maxBytes == 0 then 0 else (bytes * barWidth) `div` maxBytes
        bar = Text.replicate filled "█"
        label' = Text.justifyLeft 10 ' ' label
        valueStr = Text.justifyRight 10 ' ' (formatBytes bytes)
     in "`" <> label' <> "` " <> valueStr <> " " <> bar

-- | Format bytes in human-readable form.
formatBytes :: Int -> Text.Text
formatBytes n
    | n >= 1024 * 1024 * 1024 = Text.pack (show (n `div` (1024 * 1024 * 1024))) <> " GiB"
    | n >= 1024 * 1024 = Text.pack (show (n `div` (1024 * 1024))) <> " MiB"
    | n >= 1024 = Text.pack (show (n `div` 1024)) <> " KiB"
    | otherwise = Text.pack (show n) <> " B"

-- | Format session metadata.
formatSessionInfo :: Session.Session -> Text.Text
formatSessionInfo session =
    "**🔑 Session ID:** "
        <> formatSessionId session.sessionId
        <> "\n"
        <> maybe "" (\sid -> "**🍴 Forked from:** " <> formatSessionId sid <> "\n") session.forkedFromSessionId

-- | Extract UUID text from SessionId.
formatSessionId :: Session.SessionId -> Text.Text
formatSessionId (Session.SessionId uuid) = Text.pack $ show uuid

{- | Format all turns (already numbered and ordered).
The firstDisplayStepNum indicates which step number appears first in the output
(used for determining whether to show system prompt/tools on the first displayed turn).
-}
formatTurns :: SessionPrintOptions -> Maybe Int -> [(Int, Session.Turn)] -> Text.Text
formatTurns opts firstDisplayStepNum turns =
    Text.intercalate "\n\n---\n\n" $ map (formatTurn opts firstDisplayStepNum) turns

-- | Format a single turn with its original step number.
formatTurn :: SessionPrintOptions -> Maybe Int -> (Int, Session.Turn) -> Text.Text
formatTurn opts firstDisplayStepNum (stepNum, turn) =
    let isFirstTurn = Just stepNum == firstDisplayStepNum
     in case turn of
            Session.UserTurn content _mUsage ->
                "## 👤 Step "
                    <> Text.pack (show (stepNum :: Int))
                    <> ": User Turn\n\n"
                    <> formatUserTurn opts isFirstTurn content
            Session.LlmTurn content _mUsage ->
                "## 🤖 Step "
                    <> Text.pack (show (stepNum :: Int))
                    <> ": LLM Turn\n\n"
                    <> formatLlmTurn opts content
            Session.PartialUserTurn _content _mUsage ->
                "## ⏸️ Step "
                    <> Text.pack (show (stepNum :: Int))
                    <> ": Partial Turn (in progress)\n\n"
                    <> "_(Turn execution paused - async mode)_\n"

{- | Format user turn content.
The 'isFirstTurn' parameter ensures that system prompt and tools are shown
at least once, even when 'repeatSystemPrompt' or 'repeatTools' is False.
-}
formatUserTurn :: SessionPrintOptions -> Bool -> Session.UserTurnContent -> Text.Text
formatUserTurn opts isFirstTurn content =
    let
        -- Show system prompt if repeatSystemPrompt is True, or if it's the first turn
        systemPromptSection = case content.userPrompt of
            Session.SystemPrompt sp
                | opts.repeatSystemPrompt || isFirstTurn ->
                    "### 📝 System Prompt\n\n```\n" <> sp <> "```\n"
            _ -> ""
        querySection = case content.userQuery of
            Just (Session.UserQuery q _) -> "\n### 💬 User Query\n\n" <> q <> "\n"
            Nothing -> ""
        -- Show tools if repeatTools is True, or if it's the first turn (and tools exist)
        toolsSection =
            if null content.userTools
                then ""
                else
                    if opts.repeatTools || isFirstTurn
                        then "\n### 🛠️ Available Tools\n\n" <> formatAvailableTools content.userTools
                        else ""
        toolResponsesSection = case opts.showToolCallResults of
            Hidden -> ""
            _ | null content.userToolResponses -> ""
            _ -> "\n### 📥 Tool Responses\n\n" <> formatToolResponses opts.showToolCallResults content.userToolResponses
     in
        systemPromptSection <> querySection <> toolsSection <> toolResponsesSection

-- | Format LLM turn content.
formatLlmTurn :: SessionPrintOptions -> Session.LlmTurnContent -> Text.Text
formatLlmTurn opts content =
    let thinkingSection = case content.llmResponse.responseThinking of
            Just thinking ->
                "### 💭 Thinking Process\n\n"
                    <> "<details>\n"
                    <> "<summary>Click to expand reasoning</summary>\n\n"
                    <> thinking
                    <> "\n"
                    <> "</details>\n\n"
            Nothing -> ""
        responseSection = case content.llmResponse.responseText of
            Just txt -> "### 💬 Response\n\n" <> txt <> "\n"
            Nothing -> "### 💬 Response\n\n_(No text response)_\n"
        toolCallsSection =
            if null content.llmToolCalls
                then ""
                else "\n### 🔧 Tool Calls\n\n" <> formatLlmToolCalls opts.showToolCallArguments content.llmToolCalls
     in thinkingSection <> responseSection <> toolCallsSection

-- | Format available tools (just names and descriptions).
formatAvailableTools :: [Session.SystemTool] -> Text.Text
formatAvailableTools tools =
    Text.intercalate "\n" $ map formatSystemTool tools

-- | Format a single system tool.
formatSystemTool :: Session.SystemTool -> Text.Text
formatSystemTool (Session.SystemTool toolDef) = case toolDef of
    Session.V0 val -> "- **Tool (V0):** `" <> formatJsonAsText val <> "`"
    Session.V1 def ->
        "- **"
            <> def.name
            <> "** (`"
            <> def.llmName
            <> "`)\n"
            <> "  - Description: "
            <> def.description

-- | Format LLM tool calls (names and optionally arguments).
formatLlmToolCalls :: PrintVisibility -> [Session.LlmToolCall] -> Text.Text
formatLlmToolCalls visibility calls =
    Text.intercalate "\n\n" $ map (formatLlmToolCall visibility) calls

{- | Format a single LLM tool call, extracting the name and optionally arguments.

Tool calls from OpenAI have the structure:
@
{
  "id": "call_xxx",
  "type": "function",
  "function": {
    "name": "tool_name",
    "arguments": "{\"key\": \"value\"}"  -- Note: this is a STRING, not an object
  }
}
@
-}
formatLlmToolCall :: PrintVisibility -> Session.LlmToolCall -> Text.Text
formatLlmToolCall visibility (Session.LlmToolCall val) =
    case val of
        Aeson.Object obj ->
            -- Try to extract the function name from the tool call structure
            case KeyMap.lookup "function" obj of
                Just (Aeson.Object funcObj) ->
                    let toolName = case KeyMap.lookup "name" funcObj of
                            Just (Aeson.String n) -> n
                            _ -> "(unnamed)"
                        args = formatToolArguments visibility funcObj
                     in "- **" <> toolName <> "**" <> args
                _ -> case KeyMap.lookup "name" obj of
                    Just (Aeson.String toolName) ->
                        let args = formatToolArgumentsDirect visibility obj
                         in "- **" <> toolName <> "**" <> args
                    _ -> "- (unnamed tool call): `" <> formatJsonAsText val <> "`"
        _ -> "- (unnamed tool call): `" <> formatJsonAsText val <> "`"

{- | Format tool arguments from the function object.

The 'arguments' field is stored as a JSON string (because that's how the OpenAI API
returns it), so we need to parse it first before pretty-printing.
-}
formatToolArguments :: PrintVisibility -> KeyMap.KeyMap Aeson.Value -> Text.Text
formatToolArguments Hidden _ = ""
formatToolArguments visibility funcObj =
    case KeyMap.lookup "arguments" funcObj of
        Just argsVal ->
            let
                -- Parse arguments: if it's a string, try to parse as JSON; otherwise use as-is
                parsedVal = parseArgumentsValue argsVal
                formatted = formatJsonValuePretty parsedVal
                elided = applyVisibility visibility formatted
                byteCount = BS.length $ Text.encodeUtf8 formatted
             in
                if Text.null elided
                    then ""
                    else "\n\n  ```json\n  " <> Text.replace "\n" "\n  " elided <> "\n  ```\n  _(" <> formatBytes byteCount <> ")_"
        Nothing -> ""

{- | Format tool arguments directly from the tool call object.

Similar to 'formatToolArguments', handles the case where arguments are stored
as a JSON string rather than a JSON object.
-}
formatToolArgumentsDirect :: PrintVisibility -> KeyMap.KeyMap Aeson.Value -> Text.Text
formatToolArgumentsDirect Hidden _ = ""
formatToolArgumentsDirect visibility obj =
    case KeyMap.lookup "arguments" obj of
        Just argsVal ->
            let
                -- Parse arguments: if it's a string, try to parse as JSON; otherwise use as-is
                parsedVal = parseArgumentsValue argsVal
                formatted = formatJsonValuePretty parsedVal
                elided = applyVisibility visibility formatted
                byteCount = BS.length $ Text.encodeUtf8 formatted
             in
                if Text.null elided
                    then ""
                    else "\n\n  ```json\n  " <> Text.replace "\n" "\n  " elided <> "\n  ```\n  _(" <> formatBytes byteCount <> ")_"
        Nothing -> ""

-- | Apply visibility settings to format content.
applyVisibility :: PrintVisibility -> Text.Text -> Text.Text
applyVisibility ShownFull content = content
applyVisibility (Elided leading trailing) content = elideDocument leading trailing content
applyVisibility Hidden _ = ""

{- | Parse an arguments value.

The OpenAI API returns tool call arguments as a JSON string (e.g., "{\"key\": \"value\"}").
This function attempts to parse such strings into proper JSON values for pretty-printing.
If parsing fails or the value is not a string, it returns the original value.
-}
parseArgumentsValue :: Aeson.Value -> Aeson.Value
parseArgumentsValue (Aeson.String txt) =
    -- Try to parse the string as JSON
    case Aeson.decode (BSL.fromStrict $ Text.encodeUtf8 txt) of
        Just val -> val
        Nothing -> Aeson.String txt -- Keep as string if parsing fails
parseArgumentsValue other = other

-- | Format a JSON value as pretty-printed text.
formatJsonValuePretty :: Aeson.Value -> Text.Text
formatJsonValuePretty val =
    TextLazy.toStrict $ TextLazyEncoding.decodeUtf8 $ AesonPretty.encodePretty val

-- | Format tool responses.
formatToolResponses :: PrintVisibility -> [(Session.LlmToolCall, UserToolResponse)] -> Text.Text
formatToolResponses visibility responses =
    Text.intercalate "\n\n" $ map (formatToolResponse visibility) responses

-- | Format a single tool response.
formatToolResponse :: PrintVisibility -> (Session.LlmToolCall, UserToolResponse) -> Text.Text
formatToolResponse visibility (call, response) =
    let callName = extractToolCallName call
     in case response of
            TextResponse txt ->
                let elided = applyVisibility visibility txt
                    byteCount = BS.length $ Text.encodeUtf8 txt
                 in if Text.null elided
                        then ""
                        else "**" <> callName <> "** response:\n```\n" <> elided <> "\n```\n_(" <> formatBytes byteCount <> ")_"
            JsonResponse val ->
                let formatted = formatJsonValuePretty val
                    elided = applyVisibility visibility formatted
                    byteCount = BS.length $ Text.encodeUtf8 formatted
                 in if Text.null elided
                        then ""
                        else "**" <> callName <> "** response:\n```json\n" <> elided <> "\n```\n_(" <> formatBytes byteCount <> ")_"
            MediaResponse media ->
                let filename = maybe "" (\f -> ", " <> f) (media.mediaFilename)
                    bytes = Text.length media.mediaBase64Data
                 in "**" <> callName <> "** response:\n```\n[Media: " <> media.mediaMimeType <> filename <> ", " <> formatBytes bytes <> "]\n```"
            MixedResponse parts ->
                let formatted = Text.intercalate "\n" $ map formatContentPart parts
                    elided = applyVisibility visibility formatted
                    byteCount = BS.length $ Text.encodeUtf8 formatted
                 in if Text.null elided
                        then ""
                        else "**" <> callName <> "** response:\n```\n" <> elided <> "\n```\n_(" <> formatBytes byteCount <> ")_"

-- | Format a content part for display.
formatContentPart :: ContentPart -> Text.Text
formatContentPart (TextPart txt) = txt
formatContentPart (MediaPart media) =
    let filename = maybe "" (\f -> ", " <> f) (media.mediaFilename)
        bytes = Text.length media.mediaBase64Data
     in "[Media: " <> media.mediaMimeType <> filename <> ", " <> formatBytes bytes <> "]"

-- | Extract tool name from a tool call.
extractToolCallName :: Session.LlmToolCall -> Text.Text
extractToolCallName (Session.LlmToolCall val) =
    case val of
        Aeson.Object obj ->
            case KeyMap.lookup "function" obj of
                Just (Aeson.Object funcObj) ->
                    case KeyMap.lookup "name" funcObj of
                        Just (Aeson.String toolName) -> toolName
                        _ -> "(unnamed)"
                _ -> case KeyMap.lookup "name" obj of
                    Just (Aeson.String toolName) -> toolName
                    _ -> "(unnamed)"
        _ -> "(unnamed)"

-- | Format a JSON value as compact text.
formatJsonAsText :: Aeson.Value -> Text.Text
formatJsonAsText = Text.pack . show . Aeson.encode
