import Foundation

#if canImport(FoundationNetworking)
import FoundationNetworking
#endif
#if canImport(Network)
import Network
#endif

/// A loopback token-injecting reverse proxy for Claude's API traffic.
///
/// The problem it solves: Nucleic keeps the rotating refresh token broker-only and hands a session
/// only a *static* `CLAUDE_CODE_OAUTH_TOKEN`, fixed for the life of the `claude` process. Access
/// tokens are short-lived (~1h), so a single long turn outlives its token and 401s mid-turn — and
/// `claude` can't self-refresh because it never held the refresh token. This proxy closes that gap
/// without loosening the security posture: `claude` is pointed at `http://127.0.0.1:<port>` via
/// `ANTHROPIC_BASE_URL`, and on **every request** the proxy overwrites the `Authorization` header
/// with a broker-fresh access token (``ClaudeCredentialBroker/freshAccessToken()``) before
/// forwarding to `https://api.anthropic.com`. The refresh token still never leaves the broker; a
/// mid-turn request simply picks up a freshly-rotated bearer.
///
/// The HTTP translation (parse request → rewrite auth → serialize the upstream response back, with
/// SSE streamed through untouched) is pure and unit-tested through an injected ``Upstream``; the
/// listeners (`AF_UNIX` for the container relay, loopback TCP for host runs) mirror
/// ``MCPApprovalServer``'s transport-agnostic accept loop.
public actor ClaudeTokenProxy {
    /// Supplies the current access token to inject, refreshing through Nucleic's own OAuth when near
    /// expiry. `nil` result ⇒ no subscription login (API-key mode or signed out); the proxy then
    /// forwards whatever `Authorization` the client sent, so an API key still works.
    public typealias TokenProvider = @Sendable () async -> String?

    /// One upstream round-trip, streamed. Injected so the HTTP translation is testable without the
    /// network; the default drives `URLSession` against `api.anthropic.com`.
    public typealias Upstream = @Sendable (_ request: UpstreamRequest) async throws -> UpstreamResponse

    /// The forwarded request as the proxy reconstructs it from the inbound connection.
    public struct UpstreamRequest: Sendable {
        public var method: String
        /// Path + query, e.g. `/v1/messages?beta=true`.
        public var path: String
        /// Header order preserved; `Authorization` already rewritten to the fresh bearer (or dropped
        /// when there is no bearer and the client sent none).
        public var headers: [(name: String, value: String)]
        public var body: Data
    }

    /// The upstream response, its body streamed so SSE (`text/event-stream`) flows to the client as
    /// it arrives rather than buffering a long-lived turn.
    public struct UpstreamResponse: Sendable {
        public var status: Int
        public var headers: [(name: String, value: String)]
        public var body: AsyncThrowingStream<Data, Error>
        public init(
            status: Int, headers: [(name: String, value: String)],
            body: AsyncThrowingStream<Data, Error>
        ) {
            self.status = status
            self.headers = headers
            self.body = body
        }
    }

    private let tokenProvider: TokenProvider
    private let upstream: Upstream

    public init(tokenProvider: @escaping TokenProvider, upstream: @escaping Upstream = ClaudeTokenProxy.liveUpstream) {
        self.tokenProvider = tokenProvider
        self.upstream = upstream
    }

    // MARK: - Request handling (pure; unit-tested over in-memory byte streams)

    /// Serve one client connection: read one HTTP/1.1 request, rewrite its auth, forward it, and
    /// stream the response back. One request per connection (`Connection: close`) — simplest correct
    /// framing for a mix of short JSON calls and long SSE streams, and `claude`'s pooled client just
    /// opens a fresh connection for the next request.
    func serve(reader: ByteReader, writer: ByteWriter) async {
        defer { Task { await writer.close() } }
        guard let request = await Self.readRequest(from: reader) else {
            await Self.writeSimpleResponse(400, "Bad Request", to: writer)
            return
        }
        // Overwrite Authorization with a broker-fresh bearer. When the broker has no token (API-key
        // mode / signed out) leave whatever the client sent, so an `ANTHROPIC_API_KEY` path is
        // untouched. `x-api-key` (API-key auth) is likewise left as the client set it.
        var headers = request.headers
        if let bearer = await tokenProvider(), !bearer.isEmpty {
            headers.removeAll { $0.name.caseInsensitiveCompare("authorization") == .orderedSame }
            headers.append((name: "Authorization", value: "Bearer \(bearer)"))
        }
        let forwarded = UpstreamRequest(
            method: request.method, path: request.path, headers: headers, body: request.body)

        let response: UpstreamResponse
        do {
            response = try await upstream(forwarded)
        } catch {
            claudeAuthLog.error("Claude token proxy upstream failed: \(String(describing: error), privacy: .public)")
            await Self.writeSimpleResponse(502, "Bad Gateway", to: writer)
            return
        }
        await Self.writeResponse(response, to: writer)
    }

    /// Parse one HTTP/1.1 request: request line, headers, then a `Content-Length` body (the only
    /// framing `claude`'s client sends — JSON bodies with an explicit length, or no body). Returns
    /// nil on a malformed head or a truncated body.
    static func readRequest(from reader: ByteReader) async -> ParsedRequest? {
        let head = await readUntilDoubleCRLF(reader)
        guard let (headText, leftover) = head else { return nil }
        var lines = headText.split(separator: "\r\n", omittingEmptySubsequences: false).map(String.init)
        guard !lines.isEmpty else { return nil }
        let requestLine = lines.removeFirst().split(separator: " ").map(String.init)
        guard requestLine.count >= 2 else { return nil }
        let method = requestLine[0]
        let path = requestLine[1]

        var headers: [(name: String, value: String)] = []
        var contentLength = 0
        for line in lines where !line.isEmpty {
            guard let colon = line.firstIndex(of: ":") else { continue }
            let name = String(line[..<colon]).trimmingCharacters(in: .whitespaces)
            let value = String(line[line.index(after: colon)...]).trimmingCharacters(in: .whitespaces)
            // Drop hop-by-hop / host-specific headers we re-derive; keep everything else (anthropic
            // beta, version, content-type, user-agent…) so the upstream sees an unchanged request.
            let lower = name.lowercased()
            if lower == "host" || lower == "connection" || lower == "proxy-connection" { continue }
            if lower == "content-length", let n = Int(value) { contentLength = n }
            headers.append((name: name, value: value))
        }

        var body = leftover
        while body.count < contentLength {
            guard let chunk = await reader.read(), !chunk.isEmpty else { return nil }
            body.append(chunk)
        }
        if body.count > contentLength { body = body.prefix(contentLength) }  // no pipelining
        return ParsedRequest(method: method, path: path, headers: headers, body: body)
    }

    /// A parsed inbound request (distinct from the auth-rewritten ``UpstreamRequest``).
    struct ParsedRequest: Sendable, Equatable {
        var method: String
        var path: String
        var headers: [(name: String, value: String)]
        var body: Data
        static func == (l: ParsedRequest, r: ParsedRequest) -> Bool {
            l.method == r.method && l.path == r.path && l.body == r.body
                && l.headers.count == r.headers.count
                && zip(l.headers, r.headers).allSatisfy { $0.name == $1.name && $0.value == $1.value }
        }
    }

    /// Serialize an upstream response back to the client: status line, headers (hop-by-hop stripped,
    /// `Connection: close` added, `Content-Length`/`Transfer-Encoding` dropped since we stream and
    /// close), then the body chunks as they arrive.
    static func writeResponse(_ response: UpstreamResponse, to writer: ByteWriter) async {
        var head = "HTTP/1.1 \(response.status) \(reasonPhrase(response.status))\r\n"
        for (name, value) in response.headers {
            let lower = name.lowercased()
            if lower == "content-length" || lower == "transfer-encoding" || lower == "connection" { continue }
            head += "\(name): \(value)\r\n"
        }
        head += "Connection: close\r\n\r\n"
        await writer.write(Data(head.utf8))
        do {
            for try await chunk in response.body where !chunk.isEmpty {
                await writer.write(chunk)
            }
        } catch {
            // The upstream stream faulted mid-body; we've already sent the head, so the best we can
            // do is close (below, via the caller's defer). The client sees a truncated stream.
            claudeAuthLog.notice("Claude token proxy response stream ended early: \(String(describing: error), privacy: .public)")
        }
    }

    static func writeSimpleResponse(_ status: Int, _ phrase: String, to writer: ByteWriter) async {
        let body = Data(phrase.utf8)
        let head = "HTTP/1.1 \(status) \(phrase)\r\nContent-Length: \(body.count)\r\nConnection: close\r\n\r\n"
        await writer.write(Data(head.utf8))
        await writer.write(body)
    }

    /// Read from `reader` until the header terminator (`\r\n\r\n`), returning the header text (without
    /// the terminator) and any bytes already read past it (the start of the body). Nil on EOF first.
    private static func readUntilDoubleCRLF(_ reader: ByteReader) async -> (head: String, leftover: Data)? {
        var buffer = Data()
        let terminator = Data("\r\n\r\n".utf8)
        // Bound the header size so a nonsense client can't grow the buffer without limit.
        let maxHead = 64 * 1024
        while buffer.count <= maxHead {
            if let range = buffer.range(of: terminator) {
                let head = String(decoding: buffer[..<range.lowerBound], as: UTF8.self)
                let leftover = Data(buffer[range.upperBound...])
                return (head, leftover)
            }
            guard let chunk = await reader.read(), !chunk.isEmpty else { return nil }
            buffer.append(chunk)
        }
        return nil
    }

    private static func reasonPhrase(_ status: Int) -> String {
        switch status {
        case 200: return "OK"
        case 201: return "Created"
        case 204: return "No Content"
        case 400: return "Bad Request"
        case 401: return "Unauthorized"
        case 403: return "Forbidden"
        case 429: return "Too Many Requests"
        case 500: return "Internal Server Error"
        case 502: return "Bad Gateway"
        case 503: return "Service Unavailable"
        default: return status < 400 ? "OK" : "Error"
        }
    }

    // MARK: - Live upstream (URLSession → api.anthropic.com)

    /// Where forwarded requests go. Overridable via `ANTHROPIC_PROXY_UPSTREAM` for local testing;
    /// defaults to the real API base.
    public static let upstreamBase: String =
        ProcessInfo.processInfo.environment["ANTHROPIC_PROXY_UPSTREAM"].flatMap { $0.isEmpty ? nil : $0 }
        ?? "https://api.anthropic.com"

    /// The default `Upstream`: forward to `api.anthropic.com` and stream the response body. Uses
    /// `URLSession.bytes(for:)` so an SSE stream is relayed incrementally rather than buffered.
    public static let liveUpstream: Upstream = { request in
        guard let url = URL(string: upstreamBase + request.path) else {
            throw ProxyError.badUpstreamURL(request.path)
        }
        var urlRequest = URLRequest(url: url)
        urlRequest.httpMethod = request.method
        for (name, value) in request.headers {
            urlRequest.setValue(value, forHTTPHeaderField: name)
        }
        if !request.body.isEmpty { urlRequest.httpBody = request.body }

        #if canImport(FoundationNetworking)
        // The Linux runner's URLSession lacks the async `bytes(for:)` overload; buffer instead. The
        // runner path barely exercises the proxy (its `claude` self-refreshes from the credentials
        // file), so a non-streamed body there is acceptable.
        let (data, response) = try await URLSession.shared.data(for: urlRequest)
        let http = response as? HTTPURLResponse
        let stream = AsyncThrowingStream<Data, Error> { continuation in
            if !data.isEmpty { continuation.yield(data) }
            continuation.finish()
        }
        return UpstreamResponse(
            status: http?.statusCode ?? 502,
            headers: ClaudeTokenProxy.headerPairs(http), body: stream)
        #else
        let (bytes, response) = try await URLSession.shared.bytes(for: urlRequest)
        let http = response as? HTTPURLResponse
        let stream = AsyncThrowingStream<Data, Error> { continuation in
            let task = Task {
                do {
                    var buffer = Data()
                    buffer.reserveCapacity(16 * 1024)
                    for try await byte in bytes {
                        buffer.append(byte)
                        // Flush on newline (SSE frames end in \n\n) or when the buffer fills, so a
                        // streamed event reaches the client promptly without a per-byte write storm.
                        if byte == 0x0A || buffer.count >= 16 * 1024 {
                            continuation.yield(buffer)
                            buffer.removeAll(keepingCapacity: true)
                        }
                    }
                    if !buffer.isEmpty { continuation.yield(buffer) }
                    continuation.finish()
                } catch {
                    continuation.finish(throwing: error)
                }
            }
            continuation.onTermination = { _ in task.cancel() }
        }
        return UpstreamResponse(
            status: http?.statusCode ?? 502,
            headers: ClaudeTokenProxy.headerPairs(http), body: stream)
        #endif
    }

    private static func headerPairs(_ http: HTTPURLResponse?) -> [(name: String, value: String)] {
        guard let http else { return [] }
        return http.allHeaderFields.compactMap { key, value in
            guard let name = key as? String else { return nil }
            return (name: name, value: String(describing: value))
        }
    }

    public enum ProxyError: Error, Equatable {
        case badUpstreamURL(String)
        case listenerFailed(String)
    }

    // MARK: - Listeners (Darwin host only)
    //
    // The proxy only ever runs on the macOS host — for a host `claude` (loopback TCP) or as the host
    // end of a container's relayed unix socket. The Linux runner's `claude` self-refreshes from its
    // own credentials file and never needs this, so the listeners are Apple-only; the pure `serve`
    // logic above stays cross-platform (and unit-tested on any host).

    #if canImport(Network)
    private var tcpListener: NWListener?
    private var unixListenFD: Int32?
    private var unixAcceptSource: DispatchSourceRead?
    private var boundUnixPath: String?

    /// Start a loopback TCP listener for a host run. Returns the bound ephemeral port; point `claude`
    /// at `http://127.0.0.1:<port>` via `ANTHROPIC_BASE_URL`.
    @discardableResult
    public func start(host: String = "127.0.0.1") async throws -> UInt16 {
        let params = NWParameters.tcp
        params.requiredLocalEndpoint = NWEndpoint.hostPort(host: NWEndpoint.Host(host), port: .any)
        let listener = try NWListener(using: params)
        listener.newConnectionHandler = { [weak self] connection in
            connection.start(queue: .global(qos: .userInitiated))
            let conn = NWSocketConn(connection)
            Task { await self?.serve(reader: conn, writer: conn) }
        }
        let resume = ProxyOnce()
        let port: UInt16 = try await withCheckedThrowingContinuation { cont in
            listener.stateUpdateHandler = { state in
                switch state {
                case .ready:
                    if let p = listener.port?.rawValue, resume.claim() { cont.resume(returning: p) }
                case .failed(let error):
                    if resume.claim() { cont.resume(throwing: error) }
                case .cancelled:
                    if resume.claim() { cont.resume(throwing: ProxyError.listenerFailed("cancelled")) }
                default:
                    break
                }
            }
            listener.start(queue: .global(qos: .userInitiated))
        }
        tcpListener = listener
        return port
    }

    /// Start an `AF_UNIX` listener at `path` (the host end of a container's relayed proxy socket).
    /// Mirrors ``MCPApprovalServer``'s real-socket bind so the containerization framework's host-side
    /// relay can connect to it.
    public func start(unixSocketPath path: String) async throws {
        let fm = FileManager.default
        try? fm.createDirectory(
            atPath: (path as NSString).deletingLastPathComponent, withIntermediateDirectories: true)
        try? fm.removeItem(atPath: path)

        let fd = socket(AF_UNIX, SOCK_STREAM, 0)
        guard fd >= 0 else {
            throw ProxyError.listenerFailed("socket(): \(String(cString: strerror(errno)))")
        }
        _ = fcntl(fd, F_SETFD, FD_CLOEXEC)

        var addr = sockaddr_un()
        addr.sun_family = sa_family_t(AF_UNIX)
        let pathBytes = Array(path.utf8)
        let capacity = MemoryLayout.size(ofValue: addr.sun_path)
        guard pathBytes.count < capacity else {
            close(fd)
            throw ProxyError.listenerFailed("socket path too long: \(path)")
        }
        withUnsafeMutablePointer(to: &addr.sun_path) { raw in
            raw.withMemoryRebound(to: CChar.self, capacity: capacity) { dst in
                for (i, byte) in pathBytes.enumerated() { dst[i] = CChar(bitPattern: byte) }
                dst[pathBytes.count] = 0
            }
        }
        let bindRC = withUnsafePointer(to: &addr) { ptr in
            ptr.withMemoryRebound(to: sockaddr.self, capacity: 1) {
                Darwin.bind(fd, $0, socklen_t(MemoryLayout<sockaddr_un>.size))
            }
        }
        guard bindRC == 0 else {
            let err = String(cString: strerror(errno))
            close(fd)
            throw ProxyError.listenerFailed("bind(\(path)): \(err)")
        }
        chmod(path, 0o600)
        guard listen(fd, SOMAXCONN) == 0 else {
            let err = String(cString: strerror(errno))
            close(fd)
            try? fm.removeItem(atPath: path)
            throw ProxyError.listenerFailed("listen(\(path)): \(err)")
        }

        let source = DispatchSource.makeReadSource(
            fileDescriptor: fd, queue: .global(qos: .userInitiated))
        source.setEventHandler { [weak self] in
            let client = accept(fd, nil, nil)
            guard client >= 0 else { return }
            _ = fcntl(client, F_SETFD, FD_CLOEXEC)
            let conn = UnixSocketConn(fd: client)
            Task { await self?.serve(reader: conn, writer: conn) }
        }
        source.setCancelHandler { close(fd) }
        source.resume()

        unixListenFD = fd
        unixAcceptSource = source
        boundUnixPath = path
    }

    /// Tear down every listener (session teardown). Idempotent.
    public func stop() {
        tcpListener?.cancel()
        tcpListener = nil
        unixAcceptSource?.cancel()  // cancel handler closes the fd
        unixAcceptSource = nil
        unixListenFD = nil
        if let path = boundUnixPath { try? FileManager.default.removeItem(atPath: path) }
        boundUnixPath = nil
    }
    #endif
}

#if canImport(Network)
/// A resume-once guard for the `NWListener` state handler (`.ready` may be followed by `.failed`).
private final class ProxyOnce: @unchecked Sendable {
    private let lock = NSLock()
    private var claimed = false
    func claim() -> Bool {
        lock.lock(); defer { lock.unlock() }
        if claimed { return false }
        claimed = true
        return true
    }
}

/// ``ByteReader`` + ``ByteWriter`` over a Network.framework connection (host loopback TCP).
private final class NWSocketConn: ByteReader, ByteWriter, @unchecked Sendable {
    private let connection: NWConnection
    init(_ connection: NWConnection) { self.connection = connection }
    func read() async -> Data? {
        await withCheckedContinuation { cont in
            connection.receive(minimumIncompleteLength: 1, maximumLength: 64 * 1024) {
                data, _, isComplete, error in
                if let data, !data.isEmpty { cont.resume(returning: data) }
                else if isComplete || error != nil { cont.resume(returning: nil) }
                else { cont.resume(returning: Data()) }
            }
        }
    }
    func write(_ data: Data) async {
        await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
            connection.send(content: data, completion: .contentProcessed { _ in cont.resume() })
        }
    }
    func close() async { connection.cancel() }
}

/// ``ByteReader`` + ``ByteWriter`` over an accepted `AF_UNIX` stream socket (the relayed container
/// transport). fd-lifetime discipline mirrors ``MCPApprovalServer``'s `UnixSocketByteConn`: `close`
/// only `shutdown(2)`s; the descriptor is released once in `deinit`, after the last kernel call.
private final class UnixSocketConn: ByteReader, ByteWriter, @unchecked Sendable {
    private let fd: Int32
    private let lock = NSLock()
    private var closed = false
    init(fd: Int32) {
        self.fd = fd
        var one: Int32 = 1
        setsockopt(fd, SOL_SOCKET, SO_NOSIGPIPE, &one, socklen_t(MemoryLayout<Int32>.size))
    }
    deinit { Darwin.close(fd) }
    private var isClosed: Bool { lock.lock(); defer { lock.unlock() }; return closed }
    func read() async -> Data? {
        guard !isClosed else { return nil }
        return await withCheckedContinuation { cont in
            DispatchQueue.global(qos: .userInitiated).async {
                var buffer = [UInt8](repeating: 0, count: 64 * 1024)
                while true {
                    let n = buffer.withUnsafeMutableBytes { Darwin.read(self.fd, $0.baseAddress, 64 * 1024) }
                    if n > 0 { cont.resume(returning: Data(buffer[0..<n])) }
                    else if n == 0 { cont.resume(returning: nil) }
                    else if errno == EINTR { continue }
                    else { cont.resume(returning: nil) }
                    return
                }
            }
        }
    }
    func write(_ data: Data) async {
        guard !isClosed else { return }
        await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
            DispatchQueue.global(qos: .userInitiated).async {
                data.withUnsafeBytes { raw in
                    guard let base = raw.baseAddress else { return }
                    var offset = 0
                    while offset < raw.count {
                        let n = Darwin.send(self.fd, base + offset, raw.count - offset, Int32(MSG_NOSIGNAL_COMPAT))
                        if n > 0 { offset += n }
                        else if n < 0 && errno == EINTR { continue }
                        else { break }
                    }
                }
                cont.resume()
            }
        }
    }
    func close() async { shutdownOnce() }

    /// Synchronous so it can hold the lock (NSLock is unavailable from an async context).
    private func shutdownOnce() {
        lock.lock(); defer { lock.unlock() }
        guard !closed else { return }
        closed = true
        shutdown(fd, Int32(SHUT_RDWR))
    }
}

// macOS has no MSG_NOSIGNAL; SO_NOSIGPIPE (set on the fd above) covers it, so send flags are 0.
private let MSG_NOSIGNAL_COMPAT: Int32 = 0
#endif

// MARK: - Byte-stream seams (in-memory in tests; socket-backed in the listeners)

/// A one-directional byte source. `read()` returns the next available bytes, or nil at EOF.
protocol ByteReader: Sendable {
    func read() async -> Data?
}

/// A one-directional byte sink.
protocol ByteWriter: Sendable {
    func write(_ data: Data) async
    func close() async
}
