import Testing
import Foundation
@testable import Oppi

@Suite("ToolCallCorrelator")
@MainActor
struct ToolCallCorrelatorTests {

    // MARK: - Normal sequence

    @Test func startOutputEndSequenceSharesToolId() {
        let mapper = ToolCallCorrelator()

        let startEvent = mapper.start(sessionId: "s1", tool: "bash", args: ["command": "ls"])
        guard case .toolStart(_, let startId, let tool, let args, _) = startEvent else {
            Issue.record("Expected toolStart")
            return
        }
        #expect(tool == "bash")
        #expect(args["command"] == .string("ls"))
        #expect(!startId.isEmpty)

        let outputEvent = mapper.output(sessionId: "s1", output: "file.txt", isError: false)
        guard case .toolOutput(let payload) = outputEvent else {
            Issue.record("Expected toolOutput")
            return
        }
        #expect(payload.toolEventId == startId, "output should reuse start's toolEventId")
        #expect(payload.output == "file.txt")
        #expect(!payload.isError)

        let endEvent = mapper.end(sessionId: "s1")
        guard case .toolEnd(_, let endId, _, _, _) = endEvent else {
            Issue.record("Expected toolEnd")
            return
        }
        #expect(endId == startId, "end should reuse start's toolEventId")
    }

    @Test func updateOutputEndSequenceSharesToolIdWithoutStartingExecution() {
        let mapper = ToolCallCorrelator()

        let updateEvent = mapper.update(sessionId: "s1", tool: "edit", args: ["path": "README.md"])
        guard case .toolUpdate(_, let updateId, let tool, let args, _) = updateEvent else {
            Issue.record("Expected toolUpdate")
            return
        }
        #expect(tool == "edit")
        #expect(args["path"] == .string("README.md"))
        #expect(!updateId.isEmpty)

        let outputEvent = mapper.output(sessionId: "s1", output: "preview", isError: false)
        guard case .toolOutput(let payload) = outputEvent else {
            Issue.record("Expected toolOutput")
            return
        }
        #expect(payload.toolEventId == updateId)

        let endEvent = mapper.end(sessionId: "s1")
        guard case .toolEnd(_, let endId, _, _, _) = endEvent else {
            Issue.record("Expected toolEnd")
            return
        }
        #expect(endId == updateId)
    }

    // MARK: - Sequential tools get distinct IDs

    @Test func sequentialToolsGetDistinctIds() {
        let mapper = ToolCallCorrelator()

        let start1 = mapper.start(sessionId: "s1", tool: "bash", args: [:])
        guard case .toolStart(_, let id1, _, _, _) = start1 else {
            Issue.record("Expected toolStart")
            return
        }
        _ = mapper.end(sessionId: "s1")

        let start2 = mapper.start(sessionId: "s1", tool: "read", args: [:])
        guard case .toolStart(_, let id2, _, _, _) = start2 else {
            Issue.record("Expected toolStart")
            return
        }

        #expect(id1 != id2, "Sequential tools should get distinct IDs")
    }

    // MARK: - Output without start (orphan)

    @Test func outputWithoutStartGeneratesNewId() {
        let mapper = ToolCallCorrelator()

        // No start called — output should still produce a valid event
        let event = mapper.output(sessionId: "s1", output: "orphan", isError: true)
        guard case .toolOutput(let payload) = event else {
            Issue.record("Expected toolOutput")
            return
        }
        #expect(!payload.toolEventId.isEmpty)
        #expect(payload.output == "orphan")
        #expect(payload.isError)
    }

    // MARK: - End without start (orphan)

    @Test func endWithoutStartGeneratesNewId() {
        let mapper = ToolCallCorrelator()

        let event = mapper.end(sessionId: "s1")
        guard case .toolEnd(_, let id, _, _, _) = event else {
            Issue.record("Expected toolEnd")
            return
        }
        #expect(!id.isEmpty)
    }

    // MARK: - End clears current ID

    @Test func endClearsCurrentId() {
        let mapper = ToolCallCorrelator()

        let start = mapper.start(sessionId: "s1", tool: "bash", args: [:])
        guard case .toolStart(_, let startId, _, _, _) = start else {
            Issue.record("Expected toolStart")
            return
        }
        _ = mapper.end(sessionId: "s1")

        // After end, output should get a new (different) ID
        let orphanOutput = mapper.output(sessionId: "s1", output: "stray", isError: false)
        guard case .toolOutput(let payload) = orphanOutput else {
            Issue.record("Expected toolOutput")
            return
        }
        #expect(payload.toolEventId != startId, "After end, new events should get fresh IDs")
    }

    // MARK: - Reset

    @Test func resetClearsState() {
        let mapper = ToolCallCorrelator()

        let start = mapper.start(sessionId: "s1", tool: "bash", args: [:])
        guard case .toolStart(_, let startId, _, _, _) = start else {
            Issue.record("Expected toolStart")
            return
        }

        mapper.reset()

        // After reset, output should get a new ID (not the start's ID)
        let output = mapper.output(sessionId: "s1", output: "after-reset", isError: false)
        guard case .toolOutput(let payload) = output else {
            Issue.record("Expected toolOutput")
            return
        }
        #expect(payload.toolEventId != startId, "After reset, should not reuse old toolEventId")
    }

    // MARK: - Session ID passthrough

    @Test func sessionIdIsPassedThrough() {
        let mapper = ToolCallCorrelator()

        let start = mapper.start(sessionId: "session-42", tool: "read", args: [:])
        guard case .toolStart(let sid, _, _, _, _) = start else {
            Issue.record("Expected toolStart")
            return
        }
        #expect(sid == "session-42")

        let output = mapper.output(sessionId: "session-42", output: "data", isError: false)
        guard case .toolOutput(let payload) = output else {
            Issue.record("Expected toolOutput")
            return
        }
        #expect(payload.sessionId == "session-42")

        let end = mapper.end(sessionId: "session-42")
        guard case .toolEnd(let sid3, _, _, _, _) = end else {
            Issue.record("Expected toolEnd")
            return
        }
        #expect(sid3 == "session-42")
    }

    // MARK: - Args are preserved

    @Test func argsArePreserved() {
        let mapper = ToolCallCorrelator()
        let args: [String: JSONValue] = [
            "command": .string("echo hello"),
            "timeout": .number(30),
        ]

        let event = mapper.start(sessionId: "s1", tool: "bash", args: args)
        guard case .toolStart(_, _, _, let resultArgs, _) = event else {
            Issue.record("Expected toolStart")
            return
        }
        #expect(resultArgs["command"] == .string("echo hello"))
        #expect(resultArgs["timeout"] == .number(30))
    }

    // MARK: - Error output flag

    @Test func errorOutputFlagIsPreserved() {
        let mapper = ToolCallCorrelator()
        _ = mapper.start(sessionId: "s1", tool: "bash", args: [:])

        let event = mapper.output(sessionId: "s1", output: "stderr stuff", isError: true)
        guard case .toolOutput(let payload) = event else {
            Issue.record("Expected toolOutput")
            return
        }
        #expect(payload.isError)
    }

    // MARK: - Server-provided toolCallId

    @Test func serverProvidedToolCallIdIsUsed() {
        let mapper = ToolCallCorrelator()

        let startEvent = mapper.start(sessionId: "s1", tool: "bash", args: [:], toolCallId: "server-tc-1")
        guard case .toolStart(_, let startId, _, _, _) = startEvent else {
            Issue.record("Expected toolStart")
            return
        }
        #expect(startId == "server-tc-1", "Should use server-provided toolCallId")

        let outputEvent = mapper.output(sessionId: "s1", output: "data", isError: false, toolCallId: "server-tc-1")
        guard case .toolOutput(let payload) = outputEvent else {
            Issue.record("Expected toolOutput")
            return
        }
        #expect(payload.toolEventId == "server-tc-1", "Output should use server-provided toolCallId")

        let endEvent = mapper.end(sessionId: "s1", toolCallId: "server-tc-1")
        guard case .toolEnd(_, let endId, _, _, _) = endEvent else {
            Issue.record("Expected toolEnd")
            return
        }
        #expect(endId == "server-tc-1", "End should use server-provided toolCallId")
    }

    @Test func outputFallsBackToCurrentToolWhenNoServerToolCallId() {
        let mapper = ToolCallCorrelator()

        // Start with server-provided ID
        _ = mapper.start(sessionId: "s1", tool: "bash", args: [:], toolCallId: "server-tc-1")

        // Output without server-provided ID falls back to current tool's ID
        let outputEvent = mapper.output(sessionId: "s1", output: "data", isError: false)
        guard case .toolOutput(let payload) = outputEvent else {
            Issue.record("Expected toolOutput")
            return
        }
        #expect(payload.toolEventId == "server-tc-1", "Should fall back to current tool's server-provided ID")
    }

    @Test func noServerIdGeneratesSyntheticUUID() {
        let mapper = ToolCallCorrelator()

        let startEvent = mapper.start(sessionId: "s1", tool: "bash", args: [:])
        guard case .toolStart(_, let id, _, _, _) = startEvent else {
            Issue.record("Expected toolStart")
            return
        }
        #expect(!id.isEmpty, "Should generate synthetic UUID when no server ID")
        #expect(id != "server-tc-1", "Should not be a server-style ID")
    }
}