import Foundation
import VideoToolbox
import CoreImage
import IOSurface
import os.log

final class NativePortalStream: NSObject, URLSessionDataDelegate {
    private var session: URLSession!
    private var task: URLSessionDataTask?
    private var pending = Data()
    private let decoder = H264Decoder()
    private var dataCallbacks = 0
    var onFrame: ((CVPixelBuffer) -> Void)?
    var onStatus: ((String) -> Void)?

    func start(host: String, token: String? = PortalAuth.token) {
        stop()
        dataCallbacks = 0
        decoder.onFrame = { [weak self] b in self?.onFrame?(b) }
        let cfg = URLSessionConfiguration.default
        cfg.timeoutIntervalForRequest = .greatestFiniteMagnitude
        cfg.timeoutIntervalForResource = 7 * 24 * 60 * 60
        session = URLSession(configuration: cfg, delegate: self, delegateQueue: OperationQueue())
        guard let u = URL(string: "https://\(host)/video.h264") else {
            onStatus?("Invalid Portal address")
            return
        }
        onStatus?("Connecting native H.264 (HTTPS)…")
        var req = URLRequest(url: u)
        if let token {
            req.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
        }
        task = session.dataTask(with: req)
        task?.resume()
    }

    func stop() {
        task?.cancel()
        task = nil
        session?.invalidateAndCancel()
        session = nil
        pending.removeAll()
        decoder.stop()
    }

    func urlSession(_ session: URLSession, didReceive challenge: URLAuthenticationChallenge, completionHandler: @escaping (URLSession.AuthChallengeDisposition, URLCredential?) -> Void) {
        PortalTlsPinning.evaluate(challenge: challenge, completionHandler: completionHandler)
    }

    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive response: URLResponse, completionHandler: @escaping (URLSession.ResponseDisposition) -> Void) {
        os_log(.default, "PortalCam stream: response %{public}@", String(describing: response))
        completionHandler(.allow)
    }

    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
        dataCallbacks += 1
        if dataCallbacks <= 3 {
            os_log(.default, "PortalCam stream: received %{public}d bytes", data.count)
        }
        pending.append(data)
        consume()
    }

    func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
        if let error {
            onStatus?("Stream stopped: \(error.localizedDescription)")
        } else {
            onStatus?("Stream completed")
        }
    }

    private func consume() {
        while let first = startCode(in: pending, from: 0) {
            guard let next = startCode(in: pending, from: first.0 + first.1) else { break }
            let sc = next.0
            let nal = pending[(first.0 + first.1)..<sc]
            decoder.decode(nal: Data(nal))
            pending.removeSubrange(0..<sc)
        }
        if pending.count > 2 * 1024 * 1024 {
            pending = pending.suffix(256 * 1024)
        }
    }

    private func startCode(in d: Data, from: Int) -> (Int, Int)? {
        if d.count < 4 || from >= d.count { return nil }
        for i in from..<(d.count - 2) {
            if d[i] == 0 && d[i+1] == 0 && d[i+2] == 1 { return (i, 3) }
            if i + 3 < d.count && d[i] == 0 && d[i+1] == 0 && d[i+2] == 0 && d[i+3] == 1 { return (i, 4) }
        }
        return nil
    }
}

private final class H264Decoder {
    private var sps = Data(), pps = Data(), session: VTDecompressionSession?, format: CMVideoFormatDescription?
    var onFrame: ((CVPixelBuffer) -> Void)?
    private var decodedFrames = 0

    func stop() {
        if let session { VTDecompressionSessionInvalidate(session) }
        session = nil
        format = nil
    }

    func decode(nal: Data) {
        guard let type = nal.first.map({ $0 & 0x1f }) else { return }
        if type == 7 {
            os_log(.default, "PortalCam decoder: SPS %{public}d bytes", nal.count)
            sps = nal
            rebuild()
            return
        }
        if type == 8 {
            os_log(.default, "PortalCam decoder: PPS %{public}d bytes", nal.count)
            pps = nal
            rebuild()
            return
        }
        guard (type == 1 || type == 5), let format else { return }
        var size = UInt32(nal.count).bigEndian
        var avcc = Data(bytes: &size, count: 4)
        avcc.append(nal)
        var block: CMBlockBuffer?
        avcc.withUnsafeBytes { raw in
            _ = CMBlockBufferCreateWithMemoryBlock(
                allocator: kCFAllocatorDefault,
                memoryBlock: UnsafeMutableRawPointer(mutating: raw.baseAddress!),
                blockLength: raw.count,
                blockAllocator: kCFAllocatorNull,
                customBlockSource: nil,
                offsetToData: 0,
                dataLength: raw.count,
                flags: 0,
                blockBufferOut: &block
            )
        }
        guard let block else { return }
        var sample: CMSampleBuffer?
        var timing = CMSampleTimingInfo(
            duration: CMTime(value: 1, timescale: 30),
            presentationTimeStamp: CMClockGetTime(CMClockGetHostTimeClock()),
            decodeTimeStamp: .invalid
        )
        var sizes = [avcc.count]
        _ = CMSampleBufferCreateReady(
            allocator: kCFAllocatorDefault,
            dataBuffer: block,
            formatDescription: format,
            sampleCount: 1,
            sampleTimingEntryCount: 1,
            sampleTimingArray: &timing,
            sampleSizeEntryCount: 1,
            sampleSizeArray: &sizes,
            sampleBufferOut: &sample
        )
        if let sample {
            let rc = VTDecompressionSessionDecodeFrame(
                session!,
                sampleBuffer: sample,
                flags: [],
                frameRefcon: Unmanaged.passUnretained(self).toOpaque(),
                infoFlagsOut: nil
            )
            if rc != noErr {
                os_log(.error, "PortalCam decoder: decode returned %{public}d", rc)
            }
        }
    }

    private func rebuild() {
        guard !sps.isEmpty, !pps.isEmpty else { return }
        stop()
        var f: CMVideoFormatDescription?
        let rc: OSStatus = sps.withUnsafeBytes { sr in
            pps.withUnsafeBytes { pr in
                var ps = [sr.bindMemory(to: UInt8.self).baseAddress!, pr.bindMemory(to: UInt8.self).baseAddress!]
                var lens = [sps.count, pps.count]
                return CMVideoFormatDescriptionCreateFromH264ParameterSets(
                    allocator: kCFAllocatorDefault,
                    parameterSetCount: 2,
                    parameterSetPointers: &ps,
                    parameterSetSizes: &lens,
                    nalUnitHeaderLength: 4,
                    formatDescriptionOut: &f
                )
            }
        }
        guard rc == noErr, let f else {
            os_log(.error, "PortalCam decoder: format creation failed %{public}d", rc)
            return
        }
        os_log(.default, "PortalCam decoder: format ready")
        format = f
        let callback: VTDecompressionOutputCallback = { refCon, _, status, _, image, _, _ in
            if status == noErr, let image, let refCon {
                let decoder = Unmanaged<H264Decoder>.fromOpaque(refCon).takeUnretainedValue()
                decoder.decodedFrames += 1
                if decoder.decodedFrames <= 3 {
                    os_log(.default, "PortalCam decoder: frame decoded (#%{public}d)", decoder.decodedFrames)
                }
                decoder.onFrame?(image)
            }
        }
        var cb = VTDecompressionOutputCallbackRecord(
            decompressionOutputCallback: callback,
            decompressionOutputRefCon: Unmanaged.passUnretained(self).toOpaque()
        )
        var attrs: CFDictionary = [
            kCVPixelBufferPixelFormatTypeKey: kCVPixelFormatType_32BGRA,
            kCVPixelBufferIOSurfacePropertiesKey: [:]
        ] as CFDictionary
        var ds: VTDecompressionSession?
        let rc2 = VTDecompressionSessionCreate(
            allocator: kCFAllocatorDefault,
            formatDescription: f,
            decoderSpecification: nil,
            imageBufferAttributes: attrs,
            outputCallback: &cb,
            decompressionSessionOut: &ds
        )
        os_log(.default, "PortalCam decoder: session create %{public}d", rc2)
        session = ds
    }
}