import Testing
import Foundation
@testable import Oppi

@Suite("TimelineReducer — Tools")
@MainActor
struct TimelineReducerToolTests {

    @Test func toolCallSequence() {
        let reducer = TimelineReducer()
        let toolId = "tool-1"

        reducer.process(.agentStart(sessionId: "s1"))
        reducer.process(.toolStart(sessionId: "s1", toolEventId: toolId, tool: "bash", args: ["command": "ls"]))
        reducer.process(.toolOutput(sessionId: "s1", toolEventId: toolId, output: "file1.txt\nfile2.txt", isError: false))
        reducer.process(.toolEnd(sessionId: "s1", toolEventId: toolId))
        reducer.process(.agentEnd(sessionId: "s1"))

        let toolItems = reducer.items.filter {
            if case .toolCall = $0 { return true }
            return false
        }
        #expect(toolItems.count == 1)

        guard case .toolCall(_, let tool, _, let preview, let bytes, let isError, let isDone) = toolItems[0] else {
            Issue.record("Expected toolCall")
            return
        }
        #expect(tool == "bash")
        #expect(preview.contains("file1.txt"))
        #expect(bytes > 0)
        #expect(!isError)
        #expect(isDone)
    }

    @Test func assistantTextIsSplitAroundToolCall() {
        let reducer = TimelineReducer()
        let toolId = "tool-1"

        reducer.process(.agentStart(sessionId: "s1"))
        reducer.process(.textDelta(sessionId: "s1", delta: "before"))
        reducer.process(.toolStart(sessionId: "s1", toolEventId: toolId, tool: "bash", args: ["command": "pwd"]))
        reducer.process(.toolEnd(sessionId: "s1", toolEventId: toolId))
        reducer.process(.textDelta(sessionId: "s1", delta: "after"))
        reducer.process(.agentEnd(sessionId: "s1"))

        #expect(reducer.items.count == 3)

        guard case .assistantMessage(_, let before, _) = reducer.items[0] else {
            Issue.record("Expected first assistant message")
            return
        }
        #expect(before == "before")

        guard case .toolCall = reducer.items[1] else {
            Issue.record("Expected tool call between assistant chunks")
            return
        }

        guard case .assistantMessage(_, let after, _) = reducer.items[2] else {
            Issue.record("Expected second assistant message")
            return
        }
        #expect(after == "after")
    }

    @Test func whitespaceOnlyTextBeforeToolDiscarded() {
        let reducer = TimelineReducer()

        reducer.process(.agentStart(sessionId: "s1"))
        reducer.process(.textDelta(sessionId: "s1", delta: "\n\n"))
        reducer.process(.toolStart(sessionId: "s1", toolEventId: "t1", tool: "bash", args: [:]))
        reducer.process(.toolEnd(sessionId: "s1", toolEventId: "t1"))
        reducer.process(.textDelta(sessionId: "s1", delta: "Done!"))
        reducer.process(.agentEnd(sessionId: "s1"))

        #expect(reducer.items.count == 2)
        guard case .toolCall = reducer.items[0] else {
            Issue.record("Expected toolCall at [0], got \(reducer.items[0])")
            return
        }
        guard case .assistantMessage(_, let text, _) = reducer.items[1] else {
            Issue.record("Expected assistantMessage at [1], got \(reducer.items[1])")
            return
        }
        #expect(text == "Done!")
    }

    @Test func orphanedToolIsClosedOnAgentEnd() {
        let reducer = TimelineReducer()

        reducer.process(.agentStart(sessionId: "s1"))
        reducer.process(.toolStart(sessionId: "s1", toolEventId: "t1", tool: "read", args: [:]))
        // No toolEnd before agentEnd
        reducer.process(.agentEnd(sessionId: "s1"))

        guard case .toolCall(_, _, _, _, _, _, let isDone) = reducer.items[0] else {
            Issue.record("Expected toolCall")
            return
        }
        #expect(isDone, "Orphaned tool should be marked done on agentEnd")
    }

    @Test func toolStartStoresArgs() {
        let reducer = TimelineReducer()
        let args: [String: JSONValue] = ["command": .string("echo hello")]

        reducer.process(.agentStart(sessionId: "s1"))
        reducer.process(.toolStart(sessionId: "s1", toolEventId: "t1", tool: "bash", args: args))

        let stored = reducer.toolArgsStore.args(for: "t1")
        #expect(stored?["command"] == .string("echo hello"))
    }

    @Test func toolStartEmptyArgsNotStored() {
        let reducer = TimelineReducer()

        reducer.process(.agentStart(sessionId: "s1"))
        reducer.process(.toolStart(sessionId: "s1", toolEventId: "t1", tool: "bash", args: [:]))

        let stored = reducer.toolArgsStore.args(for: "t1")
        #expect(stored == nil, "Empty args should not be stored")
    }

    @Test func toolEndStoresDetails() {
        let reducer = TimelineReducer()

        reducer.process(.agentStart(sessionId: "s1"))
        reducer.process(.toolStart(sessionId: "s1", toolEventId: "t1", tool: "plot", args: [:]))
        reducer.process(.toolEnd(
            sessionId: "s1",
            toolEventId: "t1",
            details: .object([
                "ui": .array([
                    .object([
                        "id": .string("chart-1"),
                        "kind": .string("chart"),
                        "version": .number(1),
                    ]),
                ]),
            ])
        ))

        let stored = reducer.toolDetailsStore.details(for: "t1")
        #expect(stored?.objectValue?["ui"]?.arrayValue?.count == 1)
    }

    @Test func toolArgsStoreClearAll() {
        let store = ToolArgsStore()
        store.set(["key": .string("val")], for: "t1")
        #expect(store.args(for: "t1") != nil)

        store.clearAll()
        #expect(store.args(for: "t1") == nil)
    }

    @Test func toolOutputForUnknownIdIsStoredButNoItemCreated() {
        let reducer = TimelineReducer()

        reducer.process(.agentStart(sessionId: "s1"))
        reducer.process(.toolOutput(sessionId: "s1", toolEventId: "orphan", output: "data", isError: false))
        reducer.process(.agentEnd(sessionId: "s1"))

        let toolItems = reducer.items.filter {
            if case .toolCall = $0 { return true }
            return false
        }
        #expect(toolItems.isEmpty)
        #expect(reducer.toolOutputStore.fullOutput(for: "orphan") == "data")
    }

    @Test func toolEndForUnknownIdIsIgnored() {
        let reducer = TimelineReducer()

        reducer.process(.agentStart(sessionId: "s1"))
        reducer.process(.toolEnd(sessionId: "s1", toolEventId: "nonexistent"))
        reducer.process(.agentEnd(sessionId: "s1"))

        #expect(reducer.items.isEmpty)
    }

    @Test func traceToolResultDetailsPopulatesToolDetailsStore() {
        let reducer = TimelineReducer()
        let toolId = "tc-ext-1"
        let details: JSONValue = .object([
            "expandedText": .string("# My Todo\n\nCreated successfully"),
            "presentationFormat": .string("markdown"),
        ])

        let events: [TraceEvent] = [
            TraceEvent(
                id: "u1", type: .user, timestamp: "2026-01-01T00:00:00Z",
                text: "create a todo", tool: nil, args: nil,
                output: nil, toolCallId: nil, toolName: nil, isError: nil, thinking: nil
            ),
            TraceEvent(
                id: toolId, type: .toolCall, timestamp: "2026-01-01T00:00:01Z",
                text: nil, tool: "todo",
                args: ["action": .string("create"), "title": .string("test")],
                output: nil, toolCallId: nil, toolName: nil, isError: nil, thinking: nil
            ),
            TraceEvent(
                id: "r1", type: .toolResult, timestamp: "2026-01-01T00:00:02Z",
                text: nil, tool: nil, args: nil,
                output: "Todo created", toolCallId: toolId, toolName: "todo",
                isError: false, details: details, thinking: nil
            ),
        ]

        reducer.loadSession(events)

        // Details should be stored keyed by toolCallId (matchId)
        let stored = reducer.toolDetailsStore.details(for: toolId)
        #expect(stored == details)
    }

    // MARK: - Write tool args persistence

    @Test func writeToolArgsContentSurvivesSecondToolStart() {
        // The server sends two tool_start messages for the write tool:
        // 1. From streaming (message_update) — has full args including content
        // 2. From tool_execution_start — also has full args
        // Verify args (especially content) survive the second tool_start.
        let reducer = TimelineReducer()
        let toolId = "write-1"
        let writeArgs: [String: JSONValue] = [
            "path": .string("docs/guide.md"),
            "content": .string("# Guide\n\nSome **bold** text."),
        ]

        reducer.process(.agentStart(sessionId: "s1"))

        // First tool_start (streaming phase)
        reducer.process(.toolStart(sessionId: "s1", toolEventId: toolId, tool: "write", args: writeArgs))
        #expect(reducer.toolArgsStore.args(for: toolId)?["content"]?.stringValue == "# Guide\n\nSome **bold** text.")

        // Second tool_start (execution phase) with same args
        reducer.process(.toolStart(sessionId: "s1", toolEventId: toolId, tool: "write", args: writeArgs))
        #expect(reducer.toolArgsStore.args(for: toolId)?["content"]?.stringValue == "# Guide\n\nSome **bold** text.")

        // Tool completes
        reducer.process(.toolOutput(sessionId: "s1", toolEventId: toolId, output: "Wrote 28 bytes", isError: false))
        reducer.process(.toolEnd(sessionId: "s1", toolEventId: toolId))

        // Args should still be available after tool completion
        #expect(reducer.toolArgsStore.args(for: toolId)?["content"]?.stringValue == "# Guide\n\nSome **bold** text.")
        #expect(reducer.toolArgsStore.args(for: toolId)?["path"]?.stringValue == "docs/guide.md")
    }

    @Test func writeToolArgsContentSurvivesTraceRoundTrip() {
        // When loading a session from trace, the write tool's args (including
        // content) should be available for rendering as markdown.
        let reducer = TimelineReducer()
        let toolId = "write-md"
        let markdownContent = "# Research\n\n## Problem\n\nHeaders too large."
        let events: [TraceEvent] = [
            TraceEvent(
                id: "u1", type: .user, timestamp: "2026-01-01T00:00:00Z",
                text: "write a research doc", tool: nil, args: nil,
                output: nil, toolCallId: nil, toolName: nil, isError: nil, thinking: nil
            ),
            TraceEvent(
                id: toolId, type: .toolCall, timestamp: "2026-01-01T00:00:01Z",
                text: nil, tool: "write",
                args: [
                    "path": .string("research/typography.md"),
                    "content": .string(markdownContent),
                ],
                output: nil, toolCallId: nil, toolName: nil, isError: nil, thinking: nil
            ),
            TraceEvent(
                id: "r1", type: .toolResult, timestamp: "2026-01-01T00:00:02Z",
                text: nil, tool: nil, args: nil,
                output: "Wrote 45 bytes to research/typography.md",
                toolCallId: toolId, toolName: "write",
                isError: false, thinking: nil
            ),
        ]

        reducer.loadSession(events)

        // After loading from trace, args should have content for markdown rendering
        let storedArgs = reducer.toolArgsStore.args(for: toolId)
        #expect(storedArgs?["content"]?.stringValue == markdownContent,
                "Write tool content must survive trace round-trip for markdown rendering")
        #expect(storedArgs?["path"]?.stringValue == "research/typography.md")

        // The tool should be done
        guard case .toolCall(_, let tool, _, _, _, _, let isDone) = reducer.items.first(where: {
            if case .toolCall(let id, _, _, _, _, _, _) = $0 { return id == toolId }
            return false
        }) else {
            Issue.record("Expected toolCall item for write tool")
            return
        }
        #expect(tool == "write")
        #expect(isDone)
    }

    // MARK: - Replace mode (shell preview)

    @Test func toolOutputReplaceModeOverwritesStore() {
        let reducer = TimelineReducer()
        let toolId = "tool-shell"

        reducer.process(.agentStart(sessionId: "s1"))
        reducer.process(.toolStart(sessionId: "s1", toolEventId: toolId, tool: "bash", args: ["command": "find /"]))
        // Normal append
        reducer.process(.toolOutput(sessionId: "s1", toolEventId: toolId, output: "line1\nline2\n", isError: false))
        #expect(reducer.toolOutputStore.fullOutput(for: toolId) == "line1\nline2\n")

        // Replace mode (server switched to tail preview)
        reducer.process(.toolOutput(sessionId: "s1", toolEventId: toolId, output: "line99\nline100\n", isError: false, mode: .replace, truncated: true, totalBytes: 50000))
        #expect(reducer.toolOutputStore.fullOutput(for: toolId) == "line99\nline100\n")

        reducer.process(.toolEnd(sessionId: "s1", toolEventId: toolId))
        reducer.process(.agentEnd(sessionId: "s1"))

        let toolItems = reducer.items.filter {
            if case .toolCall = $0 { return true }
            return false
        }
        #expect(toolItems.count == 1)

        guard case .toolCall(_, let tool, _, let preview, let outputByteCount, _, let isDone) = toolItems[0] else {
            Issue.record("Expected toolCall")
            return
        }
        #expect(tool == "bash")
        #expect(preview.contains("line99") || preview.contains("line100"))
        #expect(outputByteCount == 50000)
        #expect(reducer.toolOutputStore.hasPreviewOnlyOutput(for: toolId))
        #expect(isDone)
    }

    @Test func toolOutputReplaceModeInBatch() {
        let reducer = TimelineReducer()
        let toolId = "tool-batch"

        reducer.process(.agentStart(sessionId: "s1"))
        reducer.process(.toolStart(sessionId: "s1", toolEventId: toolId, tool: "bash", args: [:]))

        // Batch with mixed append and replace — replace should win
        reducer.processBatch([
            .toolOutput(sessionId: "s1", toolEventId: toolId, output: "early\n", isError: false),
            .toolOutput(sessionId: "s1", toolEventId: toolId, output: "tail preview\n", isError: false, mode: .replace, truncated: true, totalBytes: 10000),
        ])

        #expect(reducer.toolOutputStore.fullOutput(for: toolId) == "tail preview\n")
        #expect(reducer.toolOutputStore.outputByteCount(for: toolId) == 10000)
        #expect(reducer.toolOutputStore.hasPreviewOnlyOutput(for: toolId))
    }

    @Test func traceToolResultWithoutDetailsLeavesStoreEmpty() {
        let reducer = TimelineReducer()
        let toolId = "tc-bash-1"

        let events: [TraceEvent] = [
            TraceEvent(
                id: "u1", type: .user, timestamp: "2026-01-01T00:00:00Z",
                text: "list files", tool: nil, args: nil,
                output: nil, toolCallId: nil, toolName: nil, isError: nil, thinking: nil
            ),
            TraceEvent(
                id: toolId, type: .toolCall, timestamp: "2026-01-01T00:00:01Z",
                text: nil, tool: "bash", args: ["command": .string("ls")],
                output: nil, toolCallId: nil, toolName: nil, isError: nil, thinking: nil
            ),
            TraceEvent(
                id: "r1", type: .toolResult, timestamp: "2026-01-01T00:00:02Z",
                text: nil, tool: nil, args: nil,
                output: "file1.txt", toolCallId: toolId, toolName: "bash",
                isError: false, thinking: nil
            ),
        ]

        reducer.loadSession(events)

        let stored = reducer.toolDetailsStore.details(for: toolId)
        #expect(stored == nil)
    }
}