import ByreAPI
import Foundation
import Network
import Observation

typealias Preview = Components.Schemas.Preview

@MainActor
@Observable
final class Forwards {
    private struct Key: Hashable {
        let taskId: String
        let name: String
        let port: Int
    }

    private let host: String
    private var localPorts: [Key: Int] = [:]
    @ObservationIgnored private var forwards: [Key: Forward] = [:]

    init(host: String) {
        self.host = host
    }

    func url(_ taskId: String, _ port: Int) -> URL? {
        guard let (key, local) = localPorts.first(where: { $0.key.taskId == taskId && $0.key.port == port }) else {
            return nil
        }
        var c = URLComponents()
        c.scheme = "http"
        c.host = "\(slug(key.name).isEmpty ? "preview" : slug(key.name))-\(taskId.prefix(6)).localhost"
        c.port = local
        c.path = "/"
        return c.url
    }

    func load(_ api: API, _ taskId: String) async {
        do {
            set(taskId, try await api.client.listPreviews(path: .init(taskId: taskId)).ok.body.json)
        } catch {
            log.error("list previews: \(errorText(error), privacy: .public)")
        }
    }

    func set(_ taskId: String, _ previews: [Preview]) {
        let wanted = Set(previews.map { Key(taskId: taskId, name: $0.name, port: $0.port) })
        for key in forwards.keys where key.taskId == taskId && !wanted.contains(key) {
            drop(key)
        }
        for key in wanted where forwards[key] == nil {
            forwards[key] = Forward(remote: "\(host):\(key.port)", port: key.port) { [weak self] local in
                Task { @MainActor in
                    if self?.forwards[key] != nil {
                        self?.localPorts[key] = local
                    }
                }
            }
        }
    }

    func sync(_ api: API, _ taskIds: Set<String>) {
        for key in forwards.keys where !taskIds.contains(key.taskId) {
            drop(key)
        }
        for taskId in taskIds {
            Task {
                await load(api, taskId)
            }
        }
    }

    func close(_ taskId: String) {
        for key in forwards.keys where key.taskId == taskId {
            drop(key)
        }
    }

    func closeAll() {
        for key in forwards.keys {
            drop(key)
        }
    }

    private func drop(_ key: Key) {
        forwards.removeValue(forKey: key)?.close()
        localPorts[key] = nil
    }
}

private final class Forward: @unchecked Sendable {
    private let remote: String
    private let ready: @Sendable (Int) -> Void
    private let queue = DispatchQueue(label: "run.fc.byre.forward")
    private var listener: NWListener?
    private var splices: [ObjectIdentifier: Splice] = [:]
    private var closed = false

    init(remote: String, port: Int, ready: @escaping @Sendable (Int) -> Void) {
        self.remote = remote
        self.ready = ready
        queue.async {
            self.listen(NWEndpoint.Port(rawValue: UInt16(port)) ?? .any)
        }
    }

    func close() {
        queue.async {
            self.closed = true
            self.listener?.newConnectionHandler = nil
            self.listener?.cancel()
            self.listener = nil
            for splice in self.splices.values {
                splice.close()
            }
            self.splices = [:]
        }
    }

    private func listen(_ port: NWEndpoint.Port) {
        guard !closed else {
            return
        }
        let params = NWParameters.tcp
        params.requiredLocalEndpoint = .hostPort(host: .ipv4(.loopback), port: port)
        let listener: NWListener
        do {
            listener = try NWListener(using: params)
        } catch {
            log.error("preview listener for \(self.remote, privacy: .public): \(String(describing: error), privacy: .public)")
            return
        }
        self.listener = listener
        listener.stateUpdateHandler = { [weak self, weak listener] state in
            guard let self, let listener else {
                return
            }
            switch state {
            case .ready:
                if let port = listener.port {
                    self.ready(Int(port.rawValue))
                }
            case .failed(let error):
                listener.cancel()
                if port == .any {
                    log.error("preview listener for \(self.remote, privacy: .public): \(String(describing: error), privacy: .public)")
                    return
                }
                self.listen(.any)
            default:
                break
            }
        }
        listener.newConnectionHandler = { [remote, queue] conn in
            conn.start(queue: queue)
            DispatchQueue.global().async {
                do {
                    let fd = try tailnetNode.dial(remote)
                    queue.async {
                        let splice = Splice(conn, fd, queue)
                        if self.closed {
                            splice.close()
                            return
                        }
                        let id = ObjectIdentifier(splice)
                        self.splices[id] = splice
                        splice.start {
                            self.splices[id] = nil
                        }
                    }
                } catch {
                    log.error("preview: \(errorText(error), privacy: .public)")
                    conn.cancel()
                }
            }
        }
        listener.start(queue: queue)
    }
}

private final class Splice: @unchecked Sendable {
    private let conn: NWConnection
    private let fd: Int32
    private let io: DispatchIO
    private let queue: DispatchQueue
    private static let chunk = 1 << 16
    private var open = 2
    private var ended: (() -> Void)?
    private var read = 0
    private var sending = 0
    private var readDone = false

    init(_ conn: NWConnection, _ fd: Int32, _ queue: DispatchQueue) {
        self.conn = conn
        self.fd = fd
        self.queue = queue
        io = DispatchIO(type: .stream, fileDescriptor: fd, queue: queue) { _ in
            Darwin.close(fd)
        }
        io.setLimit(lowWater: 1)
    }

    func close() {
        guard open > 0 else {
            return
        }
        open = 0
        conn.cancel()
        io.close(flags: .stop)
        ended?()
        ended = nil
    }

    func start(_ ended: @escaping () -> Void) {
        self.ended = ended
        readChunk()
        receive()
    }

    // One chunk at a time: the next read starts only once everything read so far has been sent.
    private func readChunk() {
        read = 0
        readDone = false
        io.read(offset: 0, length: Self.chunk, queue: queue) { done, data, error in
            guard self.open > 0 else {
                return
            }
            if let data, !data.isEmpty {
                self.read += data.count
                self.sending += 1
                self.conn.send(content: Data(data), completion: .contentProcessed { error in
                    self.sending -= 1
                    if error != nil {
                        self.close()
                        return
                    }
                    self.nextChunk()
                })
            }
            guard done else {
                return
            }
            if error != 0 || self.read < Self.chunk {
                self.conn.send(content: nil, contentContext: .finalMessage, isComplete: true, completion: .idempotent)
                self.half()
                return
            }
            self.readDone = true
            self.nextChunk()
        }
    }

    private func nextChunk() {
        if open > 0, readDone, sending == 0 {
            readChunk()
        }
    }

    private func receive() {
        conn.receive(minimumIncompleteLength: 1, maximumLength: Self.chunk) { data, _, complete, error in
            guard self.open > 0 else {
                return
            }
            let finished = complete || error != nil
            guard let data, !data.isEmpty else {
                finished ? self.endWrites() : self.receive()
                return
            }
            let chunk = unsafe data.withUnsafeBytes { unsafe DispatchData(bytes: $0) }
            self.io.write(offset: 0, data: chunk, queue: self.queue) { done, _, error in
                if error != 0 {
                    self.close()
                    return
                }
                guard done, self.open > 0 else {
                    return
                }
                finished ? self.endWrites() : self.receive()
            }
        }
    }

    private func endWrites() {
        io.barrier {
            Darwin.shutdown(self.fd, SHUT_WR)
        }
        half()
    }

    private func half() {
        guard open > 0 else {
            return
        }
        open -= 1
        if open == 0 {
            conn.cancel()
            io.close()
            ended?()
            ended = nil
        }
    }
}