//
//  GatewayBehaviorTests.swift
//  Zyquo Router
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Phase 7.4 — gateway behavior against a mock OpenAI-compatible upstream
//  (built from the same HTTPServer): client-disconnect cancels the upstream
//  stream, fallback chains report the actually-used model, transient 429s
//  retry, and graceful shutdown ends an active stream.
//

import XCTest
@testable import ZyquoRouter

final class GatewayBehaviorTests: XCTestCase {
    private let upstreamPort = 18901
    private let routerPort = 18902

    /// Signals observed inside the mock upstream.
    private actor UpstreamState {
        var attempts = 0
        var streamCancelled = false

        func recordAttempt() -> Int {
            attempts += 1
            return attempts
        }

        func markCancelled() {
            streamCancelled = true
        }
    }

    private func model(_ id: String, provider: ProviderID) -> AIModel {
        AIModel(
            id: id, provider: provider, displayName: id,
            contextWindow: 128_000, maxOutputTokens: 512,
            capabilities: ModelCapabilities(tools: true),
            pricing: nil,
            parameterSupport: ParameterSupport(),
            customBaseURL: URL(string: "http://127.0.0.1:\(upstreamPort)/v1")
        )
    }

    /// Mock upstream: model "slow" streams forever (marks cancellation when
    /// the client goes away), "fail-500" always 500s, "flaky-429" 429s once
    /// then succeeds, "good" answers immediately.
    private func startMockUpstream(state: UpstreamState) async -> Task<Void, Error> {
        let handler: @Sendable (RouteRequest) async -> RouteResult = { request in
            let body = (try? JSONSerialization.jsonObject(with: request.body)) as? [String: Any] ?? [:]
            let model = body["model"] as? String ?? ""
            let stream = body["stream"] as? Bool ?? false

            func completion(_ text: String) -> Data {
                ChunkEmitter.serialize([
                    "id": "mock-1", "object": "chat.completion", "created": 1, "model": model,
                    "choices": [["index": 0, "message": ["role": "assistant", "content": text],
                                 "finish_reason": "stop"] as [String: Any]],
                    "usage": ["prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2],
                ])
            }

            switch model {
            case "fail-500":
                return .complete(status: .internalServerError, headers: [], body: Data("{\"error\":{\"message\":\"boom\"}}".utf8))
            case "flaky-429":
                let attempt = await state.recordAttempt()
                if attempt == 1 {
                    return .complete(status: .tooManyRequests, headers: [("Retry-After", "0")], body: Data("{}".utf8))
                }
                return .complete(status: .ok, headers: [("Content-Type", "application/json")], body: completion("second try"))
            case "slow" where stream:
                return .stream(status: .ok, headers: []) { writer in
                    do {
                        for index in 0..<600 {
                            try await writer.send(raw: ChunkEmitter.serialize([
                                "id": "mock-s", "object": "chat.completion.chunk", "created": 1, "model": model,
                                "choices": [["index": 0, "delta": ["content": "tick\(index) "],
                                             "finish_reason": NSNull()] as [String: Any]],
                            ]))
                            try await Task.sleep(nanoseconds: 50_000_000)
                        }
                    } catch {
                        await state.markCancelled()
                        throw error
                    }
                }
            default:
                return .complete(status: .ok, headers: [("Content-Type", "application/json")], body: completion("hello from good"))
            }
        }

        let server = HTTPServer(host: "127.0.0.1", port: upstreamPort, handler: handler)
        let started = expectation(description: "mock upstream up")
        let task = Task { try await server.run { started.fulfill() } }
        await fulfillment(of: [started], timeout: 5)
        return task
    }

    private func startRouter(catalog: [AIModel], chains: [String: [String]] = [:]) async -> Task<Void, Error> {
        let routes = Routes(
            router: RequestRouter(catalog: catalog, fallbackChains: chains),
            providerKey: { _ in "mock-key" }
        )
        let server = HTTPServer(host: "127.0.0.1", port: routerPort) { request in
            await routes.handle(request)
        }
        let started = expectation(description: "router up")
        let task = Task { try await server.run { started.fulfill() } }
        await fulfillment(of: [started], timeout: 5)
        return task
    }

    private func post(_ body: [String: Any]) async throws -> (Int, [String: Any]) {
        var request = URLRequest(url: URL(string: "http://127.0.0.1:\(routerPort)/v1/chat/completions")!)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.httpBody = try JSONSerialization.data(withJSONObject: body)
        let (data, response) = try await URLSession.shared.data(for: request)
        let json = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any] ?? [:]
        return ((response as! HTTPURLResponse).statusCode, json)
    }

    func testClientDisconnectCancelsUpstream() async throws {
        let state = UpstreamState()
        let upstream = await startMockUpstream(state: state)
        let router = await startRouter(catalog: [model("slow", provider: .together)])
        defer { upstream.cancel(); router.cancel() }

        // Open a streaming request, read a couple of chunks, then drop it.
        var request = URLRequest(url: URL(string: "http://127.0.0.1:\(routerPort)/v1/chat/completions")!)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.httpBody = try JSONSerialization.data(withJSONObject: [
            "model": "together/slow",
            "messages": [["role": "user", "content": "go"]],
            "stream": true,
        ])
        let client = Task {
            let (bytes, _) = try await URLSession.shared.bytes(for: request)
            var seen = 0
            for try await _ in bytes.lines {
                seen += 1
                if seen >= 4 { break } // abandon mid-stream
            }
        }
        _ = try? await client.value

        // The router must cancel its upstream call promptly after the client
        // vanishes (write failure propagates → upstream stream torn down).
        let deadline = Date().addingTimeInterval(8)
        while Date() < deadline {
            if await state.streamCancelled { break }
            try await Task.sleep(nanoseconds: 100_000_000)
        }
        let cancelled = await state.streamCancelled
        XCTAssertTrue(cancelled, "upstream stream was not cancelled after client disconnect")
    }

    func testFallbackChainReportsActuallyUsedModel() async throws {
        let state = UpstreamState()
        let upstream = await startMockUpstream(state: state)
        let router = await startRouter(
            catalog: [model("fail-500", provider: .together), model("good", provider: .deepinfra)],
            chains: ["together/fail-500": ["deepinfra/good"]]
        )
        defer { upstream.cancel(); router.cancel() }

        let (status, json) = try await post([
            "model": "together/fail-500",
            "messages": [["role": "user", "content": "go"]],
        ])
        XCTAssertEqual(status, 200)
        XCTAssertEqual(json["model"] as? String, "deepinfra/good", "must report the model that answered")
    }

    func testRetryOnTransient429() async throws {
        let state = UpstreamState()
        let upstream = await startMockUpstream(state: state)
        let router = await startRouter(catalog: [model("flaky-429", provider: .together)])
        defer { upstream.cancel(); router.cancel() }

        let (status, json) = try await post([
            "model": "together/flaky-429",
            "messages": [["role": "user", "content": "go"]],
        ])
        XCTAssertEqual(status, 200, "transient 429 must be retried: \(json)")
        let attempts = await state.attempts
        XCTAssertEqual(attempts, 2)
        let content = (((json["choices"] as? [[String: Any]])?.first?["message"] as? [String: Any])?["content"] as? String)
        XCTAssertEqual(content, "second try")
    }

    func testGracefulShutdownWithActiveStream() async throws {
        let state = UpstreamState()
        let upstream = await startMockUpstream(state: state)
        let router = await startRouter(catalog: [model("slow", provider: .together)])
        defer { upstream.cancel() }

        var request = URLRequest(url: URL(string: "http://127.0.0.1:\(routerPort)/v1/chat/completions")!)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.httpBody = try JSONSerialization.data(withJSONObject: [
            "model": "together/slow",
            "messages": [["role": "user", "content": "go"]],
            "stream": true,
        ])
        let client = Task { () -> Int in
            let (bytes, _) = try await URLSession.shared.bytes(for: request)
            var seen = 0
            for try await _ in bytes.lines { seen += 1 }
            return seen
        }

        // Give the stream time to start flowing, then stop the router.
        try await Task.sleep(nanoseconds: 700_000_000)
        router.cancel()
        _ = try? await router.value

        // The client's connection must terminate (not hang) once the server
        // shuts down, and the port must be immediately rebindable.
        let seen = (try? await client.value) ?? -1
        XCTAssertNotEqual(seen, -1, "client saw some chunks then a clean termination")

        let reborn = await startRouter(catalog: [model("good", provider: .together)])
        reborn.cancel()
        _ = try? await reborn.value
    }
}