//
//  CamExtensionProvider.swift
//  CamExtension
//
//  Created by Vlad on 9/12/26.
//

import Foundation
import CoreMediaIO
import CoreVideo
import CoreGraphics
import CoreImage
import IOKit.audio
import IOSurface
import Security
import ObjectiveC
import os.log

let kFrameRate: Int = 30
let kPortalCamDeviceUUID = UUID(uuidString: "7B1E8C4B-8D2E-4A0D-9B95-8F5C1C7D9B11")!
let kPortalCamSourceStreamUUID = UUID(uuidString: "7B1E8C4B-8D2E-4A0D-9B95-8F5C1C7D9B12")!
let kPortalCamSinkStreamUUID = UUID(uuidString: "7B1E8C4B-8D2E-4A0D-9B95-8F5C1C7D9B13")!

let kStreamStartedNotification = "com.kovtash.portalcam.streamStarted"
let kStreamStoppedNotification = "com.kovtash.portalcam.streamStopped"

@_silgen_name("csops")
func csops(_ pid: pid_t, _ ops: UInt32, _ useraddr: UnsafeMutableRawPointer?, _ usersize: Int) -> Int32

// MARK: - Device Source

class CamExtensionDeviceSource: NSObject, CMIOExtensionDeviceSource {
	
	private(set) var device: CMIOExtensionDevice!
	
	private var _streamSource: CamExtensionStreamSource!
	private var _sinkStreamSource: CamExtensionSinkStreamSource!
	
	private var _streamingCounter: UInt32 = 0
	private var _videoDescription: CMFormatDescription!
	private var _bufferPool: CVPixelBufferPool!
	
	private var _placeholderPixelBuffer: CVPixelBuffer?
	private var _lastSinkFrameDate: Date?
	private let _lastSinkFrameLock = NSLock()
	
	private var _placeholderTimer: DispatchSourceTimer?
	private let _timerQueue = DispatchQueue(label: "com.kovtash.portalcam.cmio", qos: .userInteractive)
	private var _sentFrames = 0
	
	init(localizedName: String) {
		super.init()
		
		self.device = CMIOExtensionDevice(
			localizedName: localizedName,
			deviceID: kPortalCamDeviceUUID,
			legacyDeviceID: nil,
			source: self
		)
		
		let dims = CMVideoDimensions(width: 1280, height: 720)
		CMVideoFormatDescriptionCreate(
			allocator: kCFAllocatorDefault,
			codecType: kCVPixelFormatType_32BGRA,
			width: dims.width,
			height: dims.height,
			extensions: nil,
			formatDescriptionOut: &_videoDescription
		)
		
		let pixelBufferAttributes: NSDictionary = [
			kCVPixelBufferWidthKey: dims.width,
			kCVPixelBufferHeightKey: dims.height,
			kCVPixelBufferPixelFormatTypeKey: _videoDescription.mediaSubType,
			kCVPixelBufferIOSurfacePropertiesKey: [:] as NSDictionary
		]
		CVPixelBufferPoolCreate(kCFAllocatorDefault, nil, pixelBufferAttributes, &_bufferPool)
		
		let videoStreamFormat = CMIOExtensionStreamFormat(
			formatDescription: _videoDescription,
			maxFrameDuration: CMTime(value: 1, timescale: Int32(kFrameRate)),
			minFrameDuration: CMTime(value: 1, timescale: Int32(kFrameRate)),
			validFrameDurations: nil
		)
		
		_streamSource = CamExtensionStreamSource(
			localizedName: "PortalCam",
			streamID: kPortalCamSourceStreamUUID,
			streamFormat: videoStreamFormat,
			device: device
		)
		
		_sinkStreamSource = CamExtensionSinkStreamSource(
			localizedName: "PortalCam Sink",
			streamID: kPortalCamSinkStreamUUID,
			streamFormat: videoStreamFormat,
			device: device,
			deviceSource: self
		)
		
		do {
			try device.addStream(_streamSource.stream)
			try device.addStream(_sinkStreamSource.stream)
		} catch {
			fatalError("Failed to add streams: \(error.localizedDescription)")
		}
		
		createPlaceholderBuffer(width: Int(dims.width), height: Int(dims.height))
	}
	
	private func createPlaceholderBuffer(width: Int, height: Int) {
		let attrs: [CFString: Any] = [
			kCVPixelBufferCGImageCompatibilityKey: true,
			kCVPixelBufferCGBitmapContextCompatibilityKey: true,
			kCVPixelBufferIOSurfacePropertiesKey: [:] as NSDictionary
		]
		let status = CVPixelBufferCreate(
			kCFAllocatorDefault,
			width,
			height,
			kCVPixelFormatType_32BGRA,
			attrs as CFDictionary,
			&_placeholderPixelBuffer
		)
		guard status == kCVReturnSuccess, let buffer = _placeholderPixelBuffer else { return }
		
		CVPixelBufferLockBaseAddress(buffer, [])
		defer { CVPixelBufferUnlockBaseAddress(buffer, []) }
		guard let base = CVPixelBufferGetBaseAddress(buffer) else { return }
		let bytesPerRow = CVPixelBufferGetBytesPerRow(buffer)
		let colorSpace = CGColorSpaceCreateDeviceRGB()
		let bitmapInfo = CGBitmapInfo.byteOrder32Little.rawValue | CGImageAlphaInfo.premultipliedFirst.rawValue
		guard let ctx = CGContext(
			data: base,
			width: width,
			height: height,
			bitsPerComponent: 8,
			bytesPerRow: bytesPerRow,
			space: colorSpace,
			bitmapInfo: bitmapInfo
		) else { return }
		
		// Fill dark slate background
		ctx.setFillColor(red: 0.10, green: 0.12, blue: 0.16, alpha: 1.0)
		ctx.fill(CGRect(x: 0, y: 0, width: width, height: height))
		
		// Rounded border
		let insetRect = CGRect(x: 40, y: 40, width: width - 80, height: height - 80)
		let path = CGPath(roundedRect: insetRect, cornerWidth: 20, cornerHeight: 20, transform: nil)
		ctx.addPath(path)
		ctx.setStrokeColor(red: 0.28, green: 0.45, blue: 0.75, alpha: 0.8)
		ctx.setLineWidth(4)
		ctx.strokePath()
	}
	
	var availableProperties: Set<CMIOExtensionProperty> {
		return [.deviceTransportType, .deviceModel]
	}
	
	func deviceProperties(forProperties properties: Set<CMIOExtensionProperty>) throws -> CMIOExtensionDeviceProperties {
		let deviceProperties = CMIOExtensionDeviceProperties(dictionary: [:])
		if properties.contains(.deviceTransportType) {
			deviceProperties.transportType = kIOAudioDeviceTransportTypeVirtual
		}
		if properties.contains(.deviceModel) {
			deviceProperties.model = "PortalCam Camera"
		}
		return deviceProperties
	}
	
	func setDeviceProperties(_ deviceProperties: CMIOExtensionDeviceProperties) throws {
		// Settable properties
	}
	
	func startStreaming() {
		_streamingCounter += 1
		os_log(.default, "PortalCam extension: startStreaming (counter=%{public}u)", _streamingCounter)
		
		if _streamingCounter == 1 {
			// Notify main app that a client has opened the camera
			if let center = CFNotificationCenterGetDarwinNotifyCenter() {
				CFNotificationCenterPostNotification(
					center,
					CFNotificationName(kStreamStartedNotification as CFString),
					nil,
					nil,
					true
				)
			}
			
			// Schedule fallback/standby timer
			_placeholderTimer = DispatchSource.makeTimerSource(queue: _timerQueue)
			_placeholderTimer?.schedule(deadline: .now(), repeating: 1.0 / Double(kFrameRate), leeway: .milliseconds(2))
			_placeholderTimer?.setEventHandler { [weak self] in
				self?.sendPlaceholderIfNeeded()
			}
			_placeholderTimer?.resume()
		}
	}
	
	func stopStreaming() {
		if _streamingCounter > 1 {
			_streamingCounter -= 1
		} else {
			_streamingCounter = 0
			os_log(.default, "PortalCam extension: stopStreaming - posting stop notification")
			
			_placeholderTimer?.cancel()
			_placeholderTimer = nil
			
			_lastSinkFrameLock.lock()
			_lastSinkFrameDate = nil
			_lastSinkFrameLock.unlock()
			
			if let center = CFNotificationCenterGetDarwinNotifyCenter() {
				CFNotificationCenterPostNotification(
					center,
					CFNotificationName(kStreamStoppedNotification as CFString),
					nil,
					nil,
					true
				)
			}
		}
	}
	
	private func sendPlaceholderIfNeeded() {
		guard _streamingCounter > 0 else { return }
		
		_lastSinkFrameLock.lock()
		let lastDate = _lastSinkFrameDate
		_lastSinkFrameLock.unlock()
		
		// If live sink frames are actively arriving within the last 500ms, don't send placeholder
		if let lastDate = lastDate, Date().timeIntervalSince(lastDate) < 0.5 {
			return
		}
		
		guard let placeholder = _placeholderPixelBuffer else { return }
		
		let hostTime = CMClockGetTime(CMClockGetHostTimeClock())
		var timing = CMSampleTimingInfo(
			duration: CMTime(value: 1, timescale: Int32(kFrameRate)),
			presentationTimeStamp: hostTime,
			decodeTimeStamp: .invalid
		)
		
		var sample: CMSampleBuffer?
		let status = CMSampleBufferCreateForImageBuffer(
			allocator: kCFAllocatorDefault,
			imageBuffer: placeholder,
			dataReady: true,
			makeDataReadyCallback: nil,
			refcon: nil,
			formatDescription: _videoDescription,
			sampleTiming: &timing,
			sampleBufferOut: &sample
		)
		
		if status == noErr, let sample = sample {
			let nano = UInt64(hostTime.seconds * Double(NSEC_PER_SEC))
			_streamSource.stream.send(sample, discontinuity: [], hostTimeInNanoseconds: nano)
		}
	}
	
	func deliverSinkSampleBuffer(_ sampleBuffer: CMSampleBuffer) {
		_lastSinkFrameLock.lock()
		_lastSinkFrameDate = Date()
		_lastSinkFrameLock.unlock()
		
		guard _streamingCounter > 0 else { return }
		
		let hostTime = CMClockGetTime(CMClockGetHostTimeClock())
		let nano = UInt64(hostTime.seconds * Double(NSEC_PER_SEC))
		
		_sentFrames += 1
		if _sentFrames <= 3 {
			os_log(.default, "PortalCam extension: forwarded sink frame %{public}d to source stream", self._sentFrames)
		}
		
		_streamSource.stream.send(sampleBuffer, discontinuity: [], hostTimeInNanoseconds: nano)
	}
}

// MARK: - Source Stream (FaceTime, Zoom, Photo Booth)

class CamExtensionStreamSource: NSObject, CMIOExtensionStreamSource {
	
	private(set) var stream: CMIOExtensionStream!
	let device: CMIOExtensionDevice
	private let _streamFormat: CMIOExtensionStreamFormat
	
	init(localizedName: String, streamID: UUID, streamFormat: CMIOExtensionStreamFormat, device: CMIOExtensionDevice) {
		self.device = device
		self._streamFormat = streamFormat
		super.init()
		self.stream = CMIOExtensionStream(
			localizedName: localizedName,
			streamID: streamID,
			direction: .source,
			clockType: .hostTime,
			source: self
		)
	}
	
	var formats: [CMIOExtensionStreamFormat] {
		return [_streamFormat]
	}
	
	var activeFormatIndex: Int = 0 {
		didSet {
			if activeFormatIndex >= 1 {
				os_log(.error, "Invalid index: %{public}d", activeFormatIndex)
			}
		}
	}
	
	var availableProperties: Set<CMIOExtensionProperty> {
		return [.streamActiveFormatIndex, .streamFrameDuration]
	}
	
	func streamProperties(forProperties properties: Set<CMIOExtensionProperty>) throws -> CMIOExtensionStreamProperties {
		let streamProperties = CMIOExtensionStreamProperties(dictionary: [:])
		if properties.contains(.streamActiveFormatIndex) {
			streamProperties.activeFormatIndex = 0
		}
		if properties.contains(.streamFrameDuration) {
			streamProperties.frameDuration = CMTime(value: 1, timescale: Int32(kFrameRate))
		}
		return streamProperties
	}
	
	func setStreamProperties(_ streamProperties: CMIOExtensionStreamProperties) throws {
		if let activeFormatIndex = streamProperties.activeFormatIndex {
			self.activeFormatIndex = activeFormatIndex
		}
	}
	
	func authorizedToStartStream(for client: CMIOExtensionClient) -> Bool {
		return true
	}
	
	func startStream() throws {
		guard let deviceSource = device.source as? CamExtensionDeviceSource else {
			fatalError("Unexpected source type \(String(describing: device.source))")
		}
		deviceSource.startStreaming()
	}
	
	func stopStream() throws {
		guard let deviceSource = device.source as? CamExtensionDeviceSource else {
			fatalError("Unexpected source type \(String(describing: device.source))")
		}
		deviceSource.stopStreaming()
	}
}

// MARK: - Sink Stream (fed by PortalCam.app)

class CamExtensionSinkStreamSource: NSObject, CMIOExtensionStreamSource {
	
	private(set) var stream: CMIOExtensionStream!
	let device: CMIOExtensionDevice
	weak var deviceSource: CamExtensionDeviceSource?
	private let _streamFormat: CMIOExtensionStreamFormat
	
	private let _consumeQueue = DispatchQueue(label: "com.kovtash.portalcam.sink.consume", qos: .userInteractive)
	private var _isStreaming: Bool = false
	private var _activeClient: CMIOExtensionClient?
	
	init(
		localizedName: String,
		streamID: UUID,
		streamFormat: CMIOExtensionStreamFormat,
		device: CMIOExtensionDevice,
		deviceSource: CamExtensionDeviceSource
	) {
		self.device = device
		self.deviceSource = deviceSource
		self._streamFormat = streamFormat
		super.init()
		self.stream = CMIOExtensionStream(
			localizedName: localizedName,
			streamID: streamID,
			direction: .sink,
			clockType: .hostTime,
			source: self
		)
	}
	
	var formats: [CMIOExtensionStreamFormat] {
		return [_streamFormat]
	}
	
	var activeFormatIndex: Int = 0 {
		didSet {
			if activeFormatIndex >= 1 {
				os_log(.error, "Invalid sink format index: %{public}d", activeFormatIndex)
			}
		}
	}
	
	var availableProperties: Set<CMIOExtensionProperty> {
		return [
			.streamActiveFormatIndex,
			.streamFrameDuration,
			.streamSinkBufferQueueSize,
			.streamSinkBuffersRequiredForStartup
		]
	}
	
	func streamProperties(forProperties properties: Set<CMIOExtensionProperty>) throws -> CMIOExtensionStreamProperties {
		let streamProperties = CMIOExtensionStreamProperties(dictionary: [:])
		if properties.contains(.streamActiveFormatIndex) {
			streamProperties.activeFormatIndex = 0
		}
		if properties.contains(.streamFrameDuration) {
			streamProperties.frameDuration = CMTime(value: 1, timescale: Int32(kFrameRate))
		}
		if properties.contains(.streamSinkBufferQueueSize) {
			streamProperties.sinkBufferQueueSize = 8
		}
		if properties.contains(.streamSinkBuffersRequiredForStartup) {
			streamProperties.sinkBuffersRequiredForStartup = 1
		}
		return streamProperties
	}
	
	func setStreamProperties(_ streamProperties: CMIOExtensionStreamProperties) throws {
		if let activeFormatIndex = streamProperties.activeFormatIndex {
			self.activeFormatIndex = activeFormatIndex
		}
	}
	
	func authorizedToStartStream(for client: CMIOExtensionClient) -> Bool {
		guard isValidPortalCamClient(client: client) else {
			os_log(.error, "PortalCam extension: rejected unauthorized sink client PID %d, signingID %{public}@",
				   client.pid, client.signingID ?? "unknown")
			return false
		}
		_activeClient = client
		os_log(.default, "PortalCam extension: authorized verified sink client PID %d, signingID %{public}@",
			   client.pid, client.signingID ?? "unknown")
		return true
	}
	
	private func isValidPortalCamClient(client: CMIOExtensionClient) -> Bool {
		let pid = client.pid
		
		// 1. Check code signing status via kernel csops
		// CS_OPS_STATUS = 0, CS_VALID = 0x00000001
		var status: UInt32 = 0
		let statusRet = csops(pid, 0, &status, MemoryLayout<UInt32>.size)
		guard statusRet == 0 else {
			os_log(.error, "PortalCam extension: csops CS_OPS_STATUS failed for PID %d: errno %d", pid, errno)
			return false
		}
		
		guard (status & 0x00000001) != 0 else {
			os_log(.error, "PortalCam extension: process %d code signature is not valid (status 0x%{public}x)", pid, status)
			return false
		}
		
		// 2. Check code signing identity (Bundle Identifier)
		struct CSHeader {
			var magic: UInt32 = 0
			var length: UInt32 = 0
		}
		
		let CS_OPS_IDENTITY: UInt32 = 11
		var idHeader = CSHeader()
		let idHeaderRet = csops(pid, CS_OPS_IDENTITY, &idHeader, MemoryLayout<CSHeader>.size)
		guard idHeaderRet == 0 || errno == ERANGE else {
			os_log(.error, "PortalCam extension: csops CS_OPS_IDENTITY header failed for PID %d: errno %d", pid, errno)
			return false
		}
		let idLen = Int(UInt32(bigEndian: idHeader.length))
		guard idLen >= MemoryLayout<CSHeader>.size else {
			os_log(.error, "PortalCam extension: invalid identity length %d for PID %d", idLen, pid)
			return false
		}
		var idBuf = [UInt8](repeating: 0, count: idLen)
		guard csops(pid, CS_OPS_IDENTITY, &idBuf, idLen) == 0 else {
			os_log(.error, "PortalCam extension: csops CS_OPS_IDENTITY failed for PID %d", pid)
			return false
		}
		let identity = String(cString: Array(idBuf[MemoryLayout<CSHeader>.size...]))
		
		// 3. Check Team Identifier
		let CS_OPS_TEAMID: UInt32 = 14
		var teamHeader = CSHeader()
		let teamHeaderRet = csops(pid, CS_OPS_TEAMID, &teamHeader, MemoryLayout<CSHeader>.size)
		guard teamHeaderRet == 0 || errno == ERANGE else {
			os_log(.error, "PortalCam extension: csops CS_OPS_TEAMID header failed for PID %d: errno %d", pid, errno)
			return false
		}
		let teamLen = Int(UInt32(bigEndian: teamHeader.length))
		guard teamLen >= MemoryLayout<CSHeader>.size else {
			os_log(.error, "PortalCam extension: invalid team length %d for PID %d", teamLen, pid)
			return false
		}
		var teamBuf = [UInt8](repeating: 0, count: teamLen)
		guard csops(pid, CS_OPS_TEAMID, &teamBuf, teamLen) == 0 else {
			os_log(.error, "PortalCam extension: csops CS_OPS_TEAMID failed for PID %d", pid)
			return false
		}
		let teamID = String(cString: Array(teamBuf[MemoryLayout<CSHeader>.size...]))
		
		os_log(.default, "PortalCam extension: verified client PID %d: identity='%{public}@', teamID='%{public}@', flags=0x%{public}x",
			   pid, identity, teamID, status)
		
		guard identity == "com.kovtash.portalcam", teamID == "ENT9X9U544" else {
			os_log(.error, "PortalCam extension: rejected client PID %d: identity mismatch ('%{public}@' != 'com.kovtash.portalcam') or team mismatch ('%{public}@' != 'ENT9X9U544')",
				   pid, identity, teamID)
			return false
		}
		
		return true
	}
	
	func startStream() throws {
		os_log(.default, "PortalCam extension: sink startStream called")
		_isStreaming = true
		_consumeQueue.async { [weak self] in
			self?.consumeNext()
		}
	}
	
	func stopStream() throws {
		os_log(.default, "PortalCam extension: sink stopStream called")
		_isStreaming = false
		_activeClient = nil
	}
	
	private func consumeNext() {
		guard _isStreaming, let stream = self.stream else { return }
		
		guard let client = _activeClient ?? stream.streamingClients.first else {
			// No client connected yet; wait briefly and check again
			_consumeQueue.asyncAfter(deadline: .now() + .milliseconds(50)) { [weak self] in
				self?.consumeNext()
			}
			return
		}
		
		stream.consumeSampleBuffer(from: client) { [weak self] sampleBuffer, sequenceNumber, discontinuity, hasMore, error in
			guard let self = self, self._isStreaming else { return }
			
			if let error = error {
				os_log(.default, "PortalCam extension: sink consume error: %{public}@", error.localizedDescription)
				self._consumeQueue.asyncAfter(deadline: .now() + .milliseconds(100)) { [weak self] in
					self?.consumeNext()
				}
				return
			}
			
			if let sampleBuffer = sampleBuffer {
				self.deviceSource?.deliverSinkSampleBuffer(sampleBuffer)
				
				let pts = sampleBuffer.presentationTimeStamp
				let nano = pts.isValid && pts.seconds > 0 ? UInt64(pts.seconds * Double(NSEC_PER_SEC)) : mach_absolute_time()
				let output = CMIOExtensionScheduledOutput(sequenceNumber: sequenceNumber, hostTimeInNanoseconds: nano)
				stream.notifyScheduledOutputChanged(output)
				
				// Pull next buffer immediately
				self._consumeQueue.async { [weak self] in
					self?.consumeNext()
				}
			} else {
				// Queue is empty, wait half a frame duration (~16ms)
				self._consumeQueue.asyncAfter(deadline: .now() + .milliseconds(16)) { [weak self] in
					self?.consumeNext()
				}
			}
		}
	}
}

// MARK: - Provider Source

class CamExtensionProviderSource: NSObject, CMIOExtensionProviderSource {
	
	private(set) var provider: CMIOExtensionProvider!
	private var deviceSource: CamExtensionDeviceSource!
	
	init(clientQueue: DispatchQueue?) {
		super.init()
		
		provider = CMIOExtensionProvider(source: self, clientQueue: clientQueue)
		deviceSource = CamExtensionDeviceSource(localizedName: "PortalCam")
		
		do {
			try provider.addDevice(deviceSource.device)
		} catch {
			fatalError("Failed to add device: \(error.localizedDescription)")
		}
	}
	
	func connect(to client: CMIOExtensionClient) throws {
		// Client connected
	}
	
	func disconnect(from client: CMIOExtensionClient) {
		// Client disconnected
	}
	
	var availableProperties: Set<CMIOExtensionProperty> {
		return [.providerManufacturer]
	}
	
	func providerProperties(forProperties properties: Set<CMIOExtensionProperty>) throws -> CMIOExtensionProviderProperties {
		let providerProperties = CMIOExtensionProviderProperties(dictionary: [:])
		if properties.contains(.providerManufacturer) {
			providerProperties.manufacturer = "PortalCam"
		}
		return providerProperties
	}
	
	func setProviderProperties(_ providerProperties: CMIOExtensionProviderProperties) throws {
		// Settable properties
	}
}