import Foundation
import Testing

@testable import NucleicCore

struct ClaudeTokenProxyTests {
    // MARK: - In-memory byte-stream doubles

    // NSLock is unavailable from async contexts, so each double's locked step is a synchronous helper.
    final class ScriptedReader: ByteReader, @unchecked Sendable {
        private let lock = NSLock()
        private var chunks: [Data]
        init(_ chunks: [Data]) { self.chunks = chunks }
        func read() async -> Data? { next() }
        private func next() -> Data? {
            lock.lock(); defer { lock.unlock() }
            return chunks.isEmpty ? nil : chunks.removeFirst()
        }
    }

    final class CollectingWriter: ByteWriter, @unchecked Sendable {
        private let lock = NSLock()
        private var buffer = Data()
        func write(_ d: Data) async { append(d) }
        func close() async {}
        private func append(_ d: Data) { lock.lock(); buffer.append(d); lock.unlock() }
        var text: String { lock.lock(); defer { lock.unlock() }; return String(decoding: buffer, as: UTF8.self) }
    }

    /// Captures the request the proxy forwarded upstream, and replays a fixed response.
    final class UpstreamStub: @unchecked Sendable {
        private let lock = NSLock()
        private var captured: ClaudeTokenProxy.UpstreamRequest?
        var request: ClaudeTokenProxy.UpstreamRequest? {
            lock.lock(); defer { lock.unlock() }; return captured
        }
        private func capture(_ r: ClaudeTokenProxy.UpstreamRequest) {
            lock.lock(); captured = r; lock.unlock()
        }
        let respond: @Sendable (ClaudeTokenProxy.UpstreamRequest) throws -> ClaudeTokenProxy.UpstreamResponse
        init(respond: @escaping @Sendable (ClaudeTokenProxy.UpstreamRequest) throws -> ClaudeTokenProxy.UpstreamResponse) {
            self.respond = respond
        }
        func upstream() -> ClaudeTokenProxy.Upstream {
            { req in
                self.capture(req)
                return try self.respond(req)
            }
        }
    }

    private static func response(
        status: Int, headers: [(String, String)] = [], chunks: [String] = []
    ) -> ClaudeTokenProxy.UpstreamResponse {
        let stream = AsyncThrowingStream<Data, Error> { c in
            for ch in chunks { c.yield(Data(ch.utf8)) }
            c.finish()
        }
        return .init(status: status, headers: headers.map { (name: $0.0, value: $0.1) }, body: stream)
    }

    private static func requestBytes(
        method: String = "POST", path: String = "/v1/messages",
        headers: [(String, String)], body: String = ""
    ) -> Data {
        var text = "\(method) \(path) HTTP/1.1\r\n"
        for (n, v) in headers { text += "\(n): \(v)\r\n" }
        text += "\r\n\(body)"
        return Data(text.utf8)
    }

    // MARK: - Auth injection

    @Test func injectsFreshBearerOverClientAuthorization() async throws {
        let stub = UpstreamStub { _ in Self.response(status: 200, chunks: ["ok"]) }
        let proxy = ClaudeTokenProxy(tokenProvider: { "fresh-token" }, upstream: stub.upstream())
        let body = #"{"a":"b"}"#
        let reader = ScriptedReader([Self.requestBytes(
            path: "/v1/messages?beta=true",
            headers: [
                ("Host", "127.0.0.1"),
                ("Authorization", "Bearer stale-token"),
                ("anthropic-beta", "oauth-2025-04-20"),
                ("Content-Type", "application/json"),
                ("Content-Length", "\(body.utf8.count)"),
            ],
            body: body)])
        await proxy.serve(reader: reader, writer: CollectingWriter())

        let forwarded = try #require(stub.request)
        #expect(forwarded.method == "POST")
        #expect(forwarded.path == "/v1/messages?beta=true")
        #expect(String(decoding: forwarded.body, as: UTF8.self) == body)
        let auth = forwarded.headers.filter { $0.name.caseInsensitiveCompare("authorization") == .orderedSame }
        #expect(auth.count == 1)
        #expect(auth.first?.value == "Bearer fresh-token")  // stale bearer replaced
        // Passthrough headers survive; Host is stripped (re-derived by the HTTP client).
        #expect(forwarded.headers.contains { $0.name == "anthropic-beta" && $0.value == "oauth-2025-04-20" })
        #expect(!forwarded.headers.contains { $0.name.caseInsensitiveCompare("host") == .orderedSame })
    }

    @Test func leavesClientAuthWhenBrokerHasNoToken() async throws {
        // API-key mode / signed out: the broker returns nil, so an x-api-key request is untouched.
        let stub = UpstreamStub { _ in Self.response(status: 200, chunks: ["ok"]) }
        let proxy = ClaudeTokenProxy(tokenProvider: { nil }, upstream: stub.upstream())
        let reader = ScriptedReader([Self.requestBytes(
            method: "GET", path: "/v1/models",
            headers: [("Host", "x"), ("x-api-key", "sk-ant-123")])])
        await proxy.serve(reader: reader, writer: CollectingWriter())

        let forwarded = try #require(stub.request)
        #expect(forwarded.headers.contains { $0.name == "x-api-key" && $0.value == "sk-ant-123" })
        #expect(!forwarded.headers.contains { $0.name.caseInsensitiveCompare("authorization") == .orderedSame })
    }

    // MARK: - Response streaming

    @Test func streamsResponseHeadersAndBodyBackToClient() async {
        let stub = UpstreamStub { _ in
            Self.response(
                status: 200,
                headers: [("content-type", "text/event-stream"), ("content-length", "999")],
                chunks: ["event: a\ndata: 1\n\n", "event: b\ndata: 2\n\n"])
        }
        let proxy = ClaudeTokenProxy(tokenProvider: { "t" }, upstream: stub.upstream())
        let writer = CollectingWriter()
        await proxy.serve(
            reader: ScriptedReader([Self.requestBytes(headers: [("Host", "x"), ("Content-Length", "0")])]),
            writer: writer)

        let out = writer.text
        #expect(out.hasPrefix("HTTP/1.1 200 OK\r\n"))
        #expect(out.contains("content-type: text/event-stream"))
        #expect(out.contains("Connection: close"))
        // We stream + close, so the upstream framing headers are dropped (they'd mislead the client).
        #expect(!out.lowercased().contains("content-length: 999"))
        #expect(out.contains("event: a\ndata: 1\n\n"))
        #expect(out.contains("event: b\ndata: 2\n\n"))
    }

    @Test func mapsUpstreamFailureTo502() async {
        struct Boom: Error {}
        let stub = UpstreamStub { _ in throw Boom() }
        let proxy = ClaudeTokenProxy(tokenProvider: { "t" }, upstream: stub.upstream())
        let writer = CollectingWriter()
        await proxy.serve(
            reader: ScriptedReader([Self.requestBytes(headers: [("Host", "x"), ("Content-Length", "0")])]),
            writer: writer)
        #expect(writer.text.hasPrefix("HTTP/1.1 502 Bad Gateway\r\n"))
    }

    @Test func malformedRequestGives400() async {
        let proxy = ClaudeTokenProxy(tokenProvider: { "t" }, upstream: { _ in
            Self.response(status: 200)
        })
        let writer = CollectingWriter()
        // Header terminator never arrives, then EOF.
        await proxy.serve(reader: ScriptedReader([Data("GET / HTTP/1.1\r\nHost: x".utf8)]), writer: writer)
        #expect(writer.text.hasPrefix("HTTP/1.1 400 Bad Request\r\n"))
    }

    // MARK: - Request parsing

    @Test func reassemblesBodyAcrossReads() async throws {
        // Content-Length body split across multiple reads is fully assembled before forwarding.
        let stub = UpstreamStub { _ in Self.response(status: 200) }
        let proxy = ClaudeTokenProxy(tokenProvider: { "t" }, upstream: stub.upstream())
        let head = Data("POST /v1/messages HTTP/1.1\r\nHost: x\r\nContent-Length: 10\r\n\r\n12345".utf8)
        let rest = Data("67890".utf8)
        await proxy.serve(reader: ScriptedReader([head, rest]), writer: CollectingWriter())
        #expect(String(decoding: try #require(stub.request).body, as: UTF8.self) == "1234567890")
    }
}
