import Foundation
import MarkBase
import Metal
import RDMAKit

// ─────────────────────────────────────────────────────────────
// RDMA Distribution Service — Pipeline Parallelism
// ─────────────────────────────────────────────────────────────

public actor RDMADistributionService {
    public enum Role: Sendable {
        case primary(splitLayer: Int)
        case secondary
    }

    public struct Config: Sendable {
        public let role: Role
        public let peerAddress: String?
        public let peerPort: UInt16

        public init(role: Role, peerAddress: String? = nil, peerPort: UInt16 = 0) {
            self.role = role
            self.peerAddress = peerAddress
            self.peerPort = peerPort
        }
    }

    // RDMA state
    private var context: RDMAContext?
    private var pd: RDMAProtectionDomain?
    private var sendCQ: RDMACompletionQueue?
    private var recvCQ: RDMACompletionQueue?
    private var qp: RDMAQueuePair?

    // Registered GPU buffers
    private var registeredHiddenState: RegisteredMemory?

    public let config: Config
    public let discovery: RDMADiscovery

    public init(config: Config) {
        self.config = config
        self.discovery = RDMADiscovery()
    }

    /// Initialize RDMA device and resources. Returns false if no RDMA device is available.
    @discardableResult
    public func initialize() throws -> Bool {
        let devices = discovery.listDevices()
        guard let device = devices.first else {
            print("  RDMA: No devices found — running in local-only mode")
            return false
        }

        let tbInfo = device.isThunderbolt ? " (Thunderbolt)" : ""
        print("  RDMA: Found device '\(device.name)'\(tbInfo)")
        context = try discovery.openDevice(device)
        let attrs = try context!.queryDeviceAttributes()
        print("  RDMA: Firmware \(attrs.fwVer), max MR size \(attrs.maxMRSize) bytes")

        pd = try RDMAProtectionDomain(context: context!)
        sendCQ = try RDMACompletionQueue(context: context!, cqe: 128)
        recvCQ = try RDMACompletionQueue(context: context!, cqe: 128)
        qp = try RDMAQueuePair(pd: pd!, sendCQ: sendCQ!, recvCQ: recvCQ!)
        try qp!.modifyToInit()
        print("  RDMA: QP initialized in RC mode")

        if case .primary = config.role {
            try qp!.modifyToRTR(destQPN: 0)
            try qp!.modifyToRTS()
        }

        print("  RDMA: Ready")
        return true
    }

    /// Register a Metal buffer for remote RDMA access
    public func registerBuffer(_ buffer: MTLBuffer) throws -> RegisteredMemory? {
        guard let pd else { return nil }
        let reg = try pd.registerMTLBuffer(
            UnsafeMutableRawPointer(buffer.contents()),
            length: buffer.length
        )
        print("  RDMA: Registered buffer \(buffer.length) bytes (lkey=\(reg.lkey), rkey=\(reg.rkey))")
        return reg
    }

    /// Register the model's hidden state buffer for pipeline-split transfer
    public func registerHiddenState(_ model: E4BModel) throws {
        let buf = model.temps.io
        let reg = try registerBuffer(buf)
        registeredHiddenState = reg
        print("  RDMA: Hidden state buffer registered (\(buf.length) bytes)")
    }

    /// Split point for pipeline parallelism
    public var splitLayer: Int? {
        if case .primary(let layer) = config.role { return layer }
        return nil
    }
}