//
//  HTTPServer.swift
//  Zyquo Router
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Embedded HTTP/1.1 server on SwiftNIO's structured-concurrency APIs
//  (NIOAsyncChannel). One task per connection, one request at a time per
//  connection (keep-alive), spec-exact SSE via SSEWriter. Cancelling the
//  server task stops the accept loop and cancels in-flight connections,
//  which propagates into upstream provider calls.
//

import Foundation
import NIOCore
import NIOFoundationCompat
import NIOHTTP1
import NIOPosix

/// What the server does with one parsed request.
/// `.complete` answers with a full buffered body; `.stream` hands an SSEWriter
/// to the route (Phase 3 streaming completions).
enum RouteResult {
    case complete(status: HTTPResponseStatus, headers: [(String, String)], body: Data)
    case stream(status: HTTPResponseStatus, headers: [(String, String)], body: (SSEWriter) async throws -> Void)
}

/// One parsed inbound HTTP request, as handed to the route layer.
struct RouteRequest {
    let method: HTTPMethod
    let uri: String
    let headers: HTTPHeaders
    let body: Data

    /// URI path without the query string.
    var path: String {
        uri.split(separator: "?", maxSplits: 1, omittingEmptySubsequences: false)
            .first.map(String.init) ?? uri
    }
}

enum ServerError: LocalizedError {
    case portInUse(host: String, port: Int)
    case bindFailed(String)

    var errorDescription: String? {
        switch self {
        case .portInUse(_, let port):
            return "Port \(port) is already in use — try \(port + 1), or stop the other process."
        case .bindFailed(let detail):
            return "Could not start the server: \(detail)"
        }
    }
}

/// The embedded server. Create one per Start; it is single-use.
final class HTTPServer: Sendable {
    /// Maximum buffered request body (base64 images make chat bodies large).
    static let maxBodyBytes = 32 * 1024 * 1024

    let host: String
    let port: Int
    private let handler: @Sendable (RouteRequest) async -> RouteResult

    init(host: String, port: Int, handler: @escaping @Sendable (RouteRequest) async -> RouteResult) {
        self.host = host
        self.port = port
        self.handler = handler
    }

    /// Binds and serves until the surrounding task is cancelled.
    /// `onRunning` fires once the socket is bound and accepting.
    func run(onRunning: @escaping @Sendable () -> Void) async throws {
        let serverChannel: NIOAsyncChannel<NIOAsyncChannel<HTTPServerRequestPart, HTTPServerResponsePart>, Never>
        do {
            serverChannel = try await ServerBootstrap(group: MultiThreadedEventLoopGroup.singleton)
                .serverChannelOption(ChannelOptions.backlog, value: 64)
                .serverChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1)
                .childChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1)
                .bind(host: host, port: port) { channel in
                    channel.eventLoop.makeCompletedFuture {
                        try channel.pipeline.syncOperations.configureHTTPServerPipeline(withErrorHandling: true)
                        return try NIOAsyncChannel(wrappingChannelSynchronously: channel)
                    }
                }
        } catch let error as IOError where error.errnoCode == EADDRINUSE {
            throw ServerError.portInUse(host: host, port: port)
        } catch let error as NIOCore.ChannelError {
            throw ServerError.bindFailed(String(describing: error))
        }

        onRunning()

        // Cancellation of the surrounding task ends the accept iteration; the
        // task group then cancels and awaits every in-flight connection, so
        // shutdown is clean end-to-end. (withThrowingTaskGroup rather than a
        // discarding group: the package targets macOS 13.)
        try await serverChannel.executeThenClose { inbound, _ in
            try await withThrowingTaskGroup(of: Void.self) { group in
                for try await connection in inbound {
                    group.addTask { [handler] in
                        await Self.serve(connection: connection, handler: handler)
                    }
                }
            }
        }
    }

    // MARK: - Per-connection loop

    private static func serve(
        connection: NIOAsyncChannel<HTTPServerRequestPart, HTTPServerResponsePart>,
        handler: @Sendable (RouteRequest) async -> RouteResult
    ) async {
        try? await connection.executeThenClose { inbound, outbound in
            var iterator = inbound.makeAsyncIterator()
            while let part = try await iterator.next() {
                guard case .head(let head) = part else { continue }

                var bodyBuffer = ByteBuffer()
                var tooLarge = false
                readLoop: while let next = try await iterator.next() {
                    switch next {
                    case .body(var chunk):
                        if bodyBuffer.readableBytes + chunk.readableBytes > maxBodyBytes {
                            tooLarge = true
                        } else {
                            bodyBuffer.writeBuffer(&chunk)
                        }
                    case .end:
                        break readLoop
                    case .head:
                        return  // protocol violation; drop the connection
                    }
                }

                let keepAlive = head.isKeepAlive
                if tooLarge {
                    try await Self.write(
                        result: OpenAIError.response(
                            status: .payloadTooLarge,
                            message: "Request body exceeds the \(maxBodyBytes / (1024 * 1024)) MB limit.",
                            type: "invalid_request_error"
                        ),
                        version: head.version, keepAlive: false, outbound: outbound
                    )
                    return
                }

                let request = RouteRequest(
                    method: head.method,
                    uri: head.uri,
                    headers: head.headers,
                    body: bodyBuffer.readData(length: bodyBuffer.readableBytes) ?? Data()
                )
                let result = await handler(request)
                try await Self.write(result: result, version: head.version, keepAlive: keepAlive, outbound: outbound)
                if !keepAlive { return }
            }
        }
    }

    private static func write(
        result: RouteResult,
        version: HTTPVersion,
        keepAlive: Bool,
        outbound: NIOAsyncChannelOutboundWriter<HTTPServerResponsePart>
    ) async throws {
        switch result {
        case .complete(let status, let extraHeaders, let body):
            var headers = HTTPHeaders(extraHeaders)
            headers.replaceOrAdd(name: "Content-Length", value: String(body.count))
            headers.replaceOrAdd(name: "Connection", value: keepAlive ? "keep-alive" : "close")
            try await outbound.write(.head(HTTPResponseHead(version: version, status: status, headers: headers)))
            if !body.isEmpty {
                try await outbound.write(.body(.byteBuffer(ByteBuffer(bytes: body))))
            }
            try await outbound.write(.end(nil))

        case .stream(let status, let extraHeaders, let body):
            var headers = HTTPHeaders(extraHeaders)
            headers.replaceOrAdd(name: "Content-Type", value: "text/event-stream")
            headers.replaceOrAdd(name: "Cache-Control", value: "no-cache")
            headers.replaceOrAdd(name: "Connection", value: "close")
            try await outbound.write(.head(HTTPResponseHead(version: version, status: status, headers: headers)))
            let writer = SSEWriter(outbound: outbound)
            try await body(writer)
            try await outbound.write(.end(nil))
        }
    }
}