import CoreFoundation
import Foundation
import IOKit

public struct ThermalSensorReading: Equatable, Sendable {
    public var name: String
    public var celsius: Double

    public init(name: String, celsius: Double) {
        self.name = name
        self.celsius = celsius
    }
}

public enum ThermalSensorDomain: Equatable, Sendable {
    case cpu
    case gpu
    case other
}

public struct ThermalSummary: Equatable, Sendable {
    public var cpuTemperatureCelsius: Double?
    public var gpuTemperatureCelsius: Double?

    public init(cpuTemperatureCelsius: Double?, gpuTemperatureCelsius: Double?) {
        self.cpuTemperatureCelsius = cpuTemperatureCelsius
        self.gpuTemperatureCelsius = gpuTemperatureCelsius
    }
}

public enum ThermalSensorClassifier {
    public static func domain(for name: String) -> ThermalSensorDomain {
        let normalized = name.lowercased()

        if normalized.contains("gpu") || normalized.contains("agx") {
            return .gpu
        }

        if normalized.hasPrefix("tg") {
            return .gpu
        }

        if normalized.contains("cpu")
            || normalized.contains("p-core")
            || normalized.contains("ecpu")
            || normalized.contains("pcpu")
            || normalized.contains("e-core")
            || normalized.hasPrefix("tf")
            || normalized.hasPrefix("te")
        {
            return .cpu
        }

        return .other
    }
}

public enum ThermalSensorAggregator {
    public static func aggregate(_ readings: [ThermalSensorReading]) -> ThermalSummary {
        var cpu: Double?
        var gpu: Double?

        for reading in readings {
            switch ThermalSensorClassifier.domain(for: reading.name) {
            case .cpu:
                cpu = max(cpu ?? reading.celsius, reading.celsius)
            case .gpu:
                gpu = max(gpu ?? reading.celsius, reading.celsius)
            case .other:
                continue
            }
        }

        return ThermalSummary(cpuTemperatureCelsius: cpu, gpuTemperatureCelsius: gpu)
    }
}

public struct ThermalSensorReader: Sendable {
    public init() {}

    public func readSensors() -> [ThermalSensorReading] {
        IOHIDThermalSensorReader().readSensors()
            + SMCThermalSensorReader().readSensors()
    }

    public func readSummary() -> ThermalSummary {
        ThermalSensorAggregator.aggregate(self.readSensors())
    }
}

private typealias IOHIDEventSystemClientRef = CFTypeRef
private typealias IOHIDServiceClientRef = CFTypeRef
private typealias IOHIDEventRef = CFTypeRef

@_silgen_name("IOHIDEventSystemClientCreate")
private func IOHIDEventSystemClientCreate(_ allocator: CFAllocator?) -> IOHIDEventSystemClientRef?

@_silgen_name("IOHIDEventSystemClientSetMatching")
private func IOHIDEventSystemClientSetMatching(
    _ client: IOHIDEventSystemClientRef,
    _ matching: CFDictionary)

@_silgen_name("IOHIDEventSystemClientCopyServices")
private func IOHIDEventSystemClientCopyServices(_ client: IOHIDEventSystemClientRef) -> CFArray?

@_silgen_name("IOHIDServiceClientCopyProperty")
private func IOHIDServiceClientCopyProperty(
    _ service: IOHIDServiceClientRef,
    _ key: CFString)
    -> CFTypeRef?

@_silgen_name("IOHIDServiceClientCopyEvent")
private func IOHIDServiceClientCopyEvent(
    _ service: IOHIDServiceClientRef,
    _ type: Int64,
    _ options: Int32,
    _ timeout: Int64)
    -> IOHIDEventRef?

@_silgen_name("IOHIDEventGetFloatValue")
private func IOHIDEventGetFloatValue(_ event: IOHIDEventRef, _ field: Int32) -> Double

private struct IOHIDThermalSensorReader {
    private static let temperatureEventType: Int64 = 15
    private static let temperatureField: Int32 = Int32(temperatureEventType << 16)

    func readSensors() -> [ThermalSensorReading] {
        guard let client = IOHIDEventSystemClientCreate(kCFAllocatorDefault) else { return [] }

        let matching = [
            "PrimaryUsagePage": NSNumber(value: 0xFF00),
            "PrimaryUsage": NSNumber(value: 5),
        ] as NSDictionary
        IOHIDEventSystemClientSetMatching(client, matching)

        guard let services = IOHIDEventSystemClientCopyServices(client) else {
            return []
        }

        let count = CFArrayGetCount(services)
        guard count > 0 else { return [] }

        return (0..<count).compactMap { index in
            guard let rawService = CFArrayGetValueAtIndex(services, index) else { return nil }
            let service = unsafeBitCast(rawService, to: IOHIDServiceClientRef.self)
            guard let event = IOHIDServiceClientCopyEvent(
                service,
                Self.temperatureEventType,
                0,
                0)
            else { return nil }

            let celsius = IOHIDEventGetFloatValue(event, Self.temperatureField)
            guard celsius.isFinite, celsius > 0 else { return nil }

            let name = Self.sensorName(for: service)
            return ThermalSensorReading(name: name, celsius: celsius)
        }
    }

    private static func sensorName(for service: IOHIDServiceClientRef) -> String {
        let keys = ["Product", "ProductName", "Name", "location", "IOHIDSensorName"]
        for key in keys {
            if let value = IOHIDServiceClientCopyProperty(service, key as CFString) {
                if let string = value as? String, !string.isEmpty {
                    return string
                }
                if let number = value as? NSNumber {
                    return "\(key)=\(number)"
                }
            }
        }
        return "unnamed thermal sensor"
    }
}