import CoreVideo
import Darwin
import Foundation

public struct StreamClient: StreamConnector {
    public let renderSize: PixelSize

    public init(renderSize: PixelSize) {
        self.renderSize = renderSize
    }

    public func connect(socketPath: String) async throws -> any StreamConnection {
        let renderSize = renderSize
        return try await withCheckedThrowingContinuation { continuation in
            DispatchQueue.global(qos: .userInitiated).async {
                continuation.resume(with: Result { try Self.connectBlocking(path: socketPath, renderSize: renderSize) })
            }
        }
    }

    static func connectBlocking(path: String, renderSize: PixelSize) throws -> StreamSession {
        let deadline = Date().addingTimeInterval(30)
        var connectedOnce = false
        while Date() < deadline {
            if let fd = openSocket(path) {
                connectedOnce = true
                setReadTimeout(fd, seconds: 2)
                var decoder = AgentMessageDecoder()
                switch readHello(fd, into: &decoder) {
                case .hello(let version, let size):
                    guard version == WireProtocol.version else {
                        Darwin.close(fd)
                        throw StreamError.versionMismatch(found: version, need: WireProtocol.version)
                    }
                    guard size == renderSize else {
                        Darwin.close(fd)
                        throw StreamError.protocolViolation("HELLO is \(size), asked for \(renderSize)")
                    }
                    setReadTimeout(fd, seconds: 0)
                    // A VM that stops reading input must not wedge the session's queue.
                    setSendTimeout(fd, seconds: 2)
                    return StreamSession(fd: fd, decoder: decoder, size: size)
                case .violation(let detail):
                    Darwin.close(fd)
                    throw StreamError.protocolViolation(detail)
                case .none:
                    Darwin.close(fd)
                }
            }
            usleep(100_000)
        }
        throw connectedOnce ? StreamError.timedOut : StreamError.refused
    }

    private enum HelloOutcome {
        case hello(version: Int, size: PixelSize)
        case violation(String)
        case none
    }

    private static func readHello(_ fd: Int32, into decoder: inout AgentMessageDecoder) -> HelloOutcome {
        var buffer = [UInt8](repeating: 0, count: 64 * 1024)
        while true {
            let n = buffer.withUnsafeMutableBytes { read(fd, $0.baseAddress!, $0.count) }
            if n < 0, errno == EINTR { continue }
            guard n > 0 else { return .none }
            decoder.append(Data(buffer[0..<n]))
            do {
                if case .hello(let version, let size)? = try decoder.next() { return .hello(version: version, size: size) }
            } catch {
                return .violation(error.description)
            }
        }
    }

    private static func openSocket(_ path: String) -> Int32? {
        let fd = socket(AF_UNIX, SOCK_STREAM, 0)
        guard fd >= 0 else { return nil }
        var address = sockaddr_un()
        address.sun_family = sa_family_t(AF_UNIX)
        let bytes = Array(path.utf8)
        guard bytes.count < MemoryLayout.size(ofValue: address.sun_path) else {
            Darwin.close(fd)
            return nil
        }
        withUnsafeMutableBytes(of: &address.sun_path) { $0.copyBytes(from: bytes) }
        let size = socklen_t(MemoryLayout<sockaddr_un>.size)
        let result = withUnsafePointer(to: &address) {
            $0.withMemoryRebound(to: sockaddr.self, capacity: 1) { Darwin.connect(fd, $0, size) }
        }
        guard result == 0 else {
            Darwin.close(fd)
            return nil
        }
        var on: Int32 = 1
        setsockopt(fd, SOL_SOCKET, SO_NOSIGPIPE, &on, socklen_t(MemoryLayout<Int32>.size))
        return fd
    }

    static func setReadTimeout(_ fd: Int32, seconds: Int) {
        var timeout = timeval.init(tv_sec: seconds, tv_usec: 0)
        setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, socklen_t(MemoryLayout<timeval>.size))
    }

    static func setSendTimeout(_ fd: Int32, seconds: Int) {
        var timeout = timeval.init(tv_sec: seconds, tv_usec: 0)
        setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &timeout, socklen_t(MemoryLayout<timeval>.size))
    }
}

/// Its reader thread must never stop (the VM's watchdog needs a connected viewer); all else runs on `queue`.
final class StreamSession: StreamConnection, @unchecked Sendable {
    static let firstFrameTimeout: TimeInterval = 20
    static let decodeErrorLimit = 10

    let events: AsyncStream<StreamEvent>
    private let continuation: AsyncStream<StreamEvent>.Continuation
    private let queue = DispatchQueue(label: "aculix.knotchy.stream")

    private let fd: Int32
    private var decoder: AgentMessageDecoder
    private var parser = H264Parser()
    private var video: VideoDecoder?
    private var readiness = ContentReadiness()
    private var announcedReady = false
    private var enabled = true
    private var waitingForIDR = false
    private var frameHandler: (@Sendable (DecodedFrame) -> Void)?
    private var readinessTimer: DispatchSourceTimer?
    private var firstFrameTimer: DispatchSourceTimer?
    private var pendingPageURL: URL?
    private var pageURLTimer: DispatchSourceTimer?
    private var decodeErrors: [Date] = []
    private var ended = false

    init(fd: Int32, decoder: AgentMessageDecoder, size: PixelSize) {
        self.fd = fd
        self.decoder = decoder
        let (stream, continuation) = AsyncStream.makeStream(of: StreamEvent.self)
        events = stream
        self.continuation = continuation
        continuation.yield(.hello(version: WireProtocol.version, size: size))
        queue.async { self.begin() }
        let reader = Thread { [self] in readLoop() }
        reader.name = "knotchy stream reader"
        reader.start()
    }

    // MARK: StreamConnection

    var decodeEnabled: Bool {
        get { queue.sync { enabled } }
        set {
            queue.async {
                if newValue, !self.enabled {
                    self.waitingForIDR = true
                    self.write(.keyframe)
                }
                self.enabled = newValue
            }
        }
    }

    func setFrameHandler(_ handler: @escaping @Sendable (DecodedFrame) -> Void) {
        queue.async { self.frameHandler = handler }
    }

    func send(_ input: InputEvent) {
        queue.async { self.write(input) }
    }

    func close() {
        queue.async { self.end(nil) }
    }

    // MARK: Reading

    private func readLoop() {
        var buffer = [UInt8](repeating: 0, count: 64 * 1024)
        while true {
            let n = buffer.withUnsafeMutableBytes { read(fd, $0.baseAddress!, $0.count) }
            if n < 0, errno == EINTR { continue }
            guard n > 0 else { break }
            let chunk = Data(buffer[0..<n])
            queue.async { self.received(chunk) }
        }
        queue.async {
            self.end(.closed)
            Darwin.close(self.fd)
        }
    }

    private func begin() {
        readiness.hello(at: .now)
        scheduleReadinessTick()
        let timer = DispatchSource.makeTimerSource(queue: queue)
        timer.schedule(deadline: .now() + Self.firstFrameTimeout)
        timer.setEventHandler { [weak self] in self?.end(.noFirstFrame) }
        timer.resume()
        firstFrameTimer = timer
        received(Data())   // Messages that arrived together with HELLO.
    }

    private func received(_ chunk: Data) {
        guard !ended else { return }
        decoder.append(chunk)
        do {
            while let message = try decoder.next() { handle(message) }
        } catch {
            end(.protocolViolation(error.description))
        }
    }

    private func handle(_ message: AgentMessage) {
        switch message {
        case .navigated:
            readiness.navigated(at: .now)
            scheduleReadinessTick()
        case .progress(let percent):
            readiness.progress(percent)
            continuation.yield(.progress(percent))
            announceIfReady()
        case .video(_, let accessUnit):
            firstFrameTimer?.cancel()
            firstFrameTimer = nil
            let unit = parser.parse(accessUnit)
            if unit.newParameterSets != nil || video == nil, let sets = parser.parameterSets {
                video = VideoDecoder(parameterSets: sets) { [weak self] frame, status in
                    guard let self else { return }
                    self.queue.async { self.decoded(frame, status: status) }
                }
            }
            guard enabled, !unit.sample.isEmpty, let video else { return }
            if waitingForIDR {
                guard unit.isIDR else { return }
                waitingForIDR = false
            }
            video.decode(unit.sample)
        case .pageURL(let text):
            guard let url = URL(string: text) else { return }
            pendingPageURL = url
            passOnPageURL()
        case .hello, .unknown:
            break
        }
    }

    private func decoded(_ frame: DecodedFrame?, status: OSStatus) {
        guard !ended else { return }
        guard let frame else { return decodeFailed() }
        if !readiness.isReady {
            readiness.frame(isUniform: Self.isUniform(frame.pixelBuffer))
            announceIfReady()
        }
        frameHandler?(frame)
    }

    private func decodeFailed() {
        let now = Date()
        decodeErrors = decodeErrors.filter { now.timeIntervalSince($0) < 60 } + [now]
        if decodeErrors.count > Self.decodeErrorLimit { return end(.decodeErrors) }
        waitingForIDR = true
        write(.keyframe)
    }

    private static func isUniform(_ buffer: CVPixelBuffer) -> Bool {
        CVPixelBufferLockBaseAddress(buffer, .readOnly)
        defer { CVPixelBufferUnlockBaseAddress(buffer, .readOnly) }
        guard let base = CVPixelBufferGetBaseAddress(buffer) else { return true }
        let bytes = UnsafeRawBufferPointer(start: base, count: CVPixelBufferGetDataSize(buffer))
        return FrameSampler.isUniform(bytes, width: CVPixelBufferGetWidth(buffer),
                                      height: CVPixelBufferGetHeight(buffer),
                                      bytesPerRow: CVPixelBufferGetBytesPerRow(buffer))
    }

    // MARK: Readiness

    private func announceIfReady() {
        guard readiness.isReady, !announcedReady else { return }
        announcedReady = true
        readinessTimer?.cancel()
        readinessTimer = nil
        continuation.yield(.contentReady)
    }

    private func scheduleReadinessTick() {
        readinessTimer?.cancel()
        readinessTimer = nil
        guard let deadline = readiness.nextDeadline else { return }
        let wait = deadline - .now
        let seconds = max(0, Double(wait.components.seconds) + Double(wait.components.attoseconds) / 1e18)
        let timer = DispatchSource.makeTimerSource(queue: queue)
        timer.schedule(deadline: .now() + seconds)
        timer.setEventHandler { [weak self] in
            guard let self else { return }
            self.readiness.tick(now: .now)
            self.announceIfReady()
        }
        timer.resume()
        readinessTimer = timer
    }

    /// At most 10 a second, always including the last: a page can change its address in a loop.
    private func passOnPageURL() {
        guard pageURLTimer == nil, let url = pendingPageURL else { return }
        pendingPageURL = nil
        continuation.yield(.pageURL(url))
        let timer = DispatchSource.makeTimerSource(queue: queue)
        timer.schedule(deadline: .now() + 0.1)
        timer.setEventHandler { [weak self] in
            guard let self, !self.ended else { return }
            self.pageURLTimer = nil
            self.passOnPageURL()
        }
        timer.resume()
        pageURLTimer = timer
    }

    // MARK: Writing and ending

    private func write(_ input: InputEvent) {
        guard !ended, let data = input.encoded() else { return }
        let stalled = data.withUnsafeBytes { raw -> Bool in
            var offset = 0
            while offset < raw.count {
                let n = Darwin.write(fd, raw.baseAddress! + offset, raw.count - offset)
                if n > 0 { offset += n } else if n < 0, errno == EINTR { continue } else {
                    return n < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)
                }
            }
            return false
        }
        if stalled { end(.socketError("The VM stopped taking input.")) }
    }

    private func end(_ reason: EndReason?) {
        guard !ended else { return }
        ended = true
        shutdown(fd, SHUT_RDWR)
        readinessTimer?.cancel()
        firstFrameTimer?.cancel()
        pageURLTimer?.cancel()
        video = nil
        frameHandler = nil
        if let reason { continuation.yield(.ended(reason)) }
        continuation.finish()
    }
}