// NTPDissector.swift
//
// Parses an NTPv4 packet per RFC 5905 §7.3 (header layout).
// 48-byte fixed header; optional extension fields and KeyID/MAC are
// surfaced as a single "extension" data blob without further parsing.

import Foundation
import CPAPluginSDK

enum NTPMode: UInt8, Sendable {
    case reserved = 0
    case symmetricActive = 1
    case symmetricPassive = 2
    case client = 3
    case server = 4
    case broadcast = 5
    case controlMessage = 6
    case privateUse = 7

    var displayName: String {
        switch self {
        case .reserved:
            return "Reserved"
        case .symmetricActive:
            return "Symmetric Active"
        case .symmetricPassive:
            return "Symmetric Passive"
        case .client:
            return "Client"
        case .server:
            return "Server"
        case .broadcast:
            return "Broadcast"
        case .controlMessage:
            return "NTP Control Message"
        case .privateUse:
            return "Reserved (Private Use)"
        }
    }
}

enum NTPLeapIndicator: UInt8, Sendable {
    case noWarning = 0
    case lastMinute61 = 1
    case lastMinute59 = 2
    case alarm = 3

    var displayName: String {
        switch self {
        case .noWarning:
            return "No Warning"
        case .lastMinute61:
            return "Last minute has 61 seconds"
        case .lastMinute59:
            return "Last minute has 59 seconds"
        case .alarm:
            return "Clock not synchronized"
        }
    }
}

struct NTPMessage: Sendable {
    let leapIndicator: NTPLeapIndicator
    let versionNumber: UInt8
    let mode: NTPMode
    let stratum: UInt8
    let poll: Int8
    let precision: Int8
    let rootDelaySeconds: Double
    let rootDispersionSeconds: Double
    let referenceID: ReferenceID
    let referenceTimestamp: NTPTimestamp
    let originTimestamp: NTPTimestamp
    let receiveTimestamp: NTPTimestamp
    let transmitTimestamp: NTPTimestamp

    enum ReferenceID: Sendable {
        case kissCode(String)        // stratum 0
        case identifier(String)      // stratum 1
        case ipv4(String)            // stratum 2+
        case rawHex(String)          // fallback

        var displayValue: String {
            switch self {
            case .kissCode(let code):
                return "KoD: \(code)"
            case .identifier(let identifier):
                return identifier
            case .ipv4(let address):
                return address
            case .rawHex(let hex):
                return "0x\(hex)"
            }
        }
    }
}

/// 64-bit NTP timestamp: seconds since 1900-01-01 + 32-bit fraction.
struct NTPTimestamp: Sendable {
    let secondsSince1900: UInt32
    let fraction: UInt32

    /// Seconds offset from 1900 to 1970 (Unix epoch). RFC 5905 §6.
    static let unixEpochOffset: UInt32 = 2_208_988_800

    /// Returns the timestamp as a `Date`. Returns nil if the timestamp is
    /// zero (commonly seen in client-only requests).
    var date: Date? {
        guard secondsSince1900 != 0 else { return nil }
        let unix = Int64(secondsSince1900) - Int64(NTPTimestamp.unixEpochOffset)
        let fractional = Double(fraction) / Double(UInt32.max)
        return Date(timeIntervalSince1970: TimeInterval(unix) + fractional)
    }

    /// Compact "1900-seconds.fraction" form when no Unix-mappable date.
    var rawValue: String {
        "\(secondsSince1900).\(fraction)"
    }
}

final class NTPDissector: Sendable {
    static let headerSize: Int = 48

    enum ProtocolAttribute: String {
        case ntpSize
        case ntpMode
        case ntpStratum
        case ntpReference
        case ntpProtocolInfo
    }

    func parseMessage(_ data: Data) -> NTPMessage? {
        guard data.count >= NTPDissector.headerSize else { return nil }

        let firstByte: UInt8 = data[data.startIndex]
        let liRaw: UInt8 = (firstByte >> 6) & 0x03
        let vnRaw: UInt8 = (firstByte >> 3) & 0x07
        let modeRaw: UInt8 = firstByte & 0x07

        let stratum: UInt8 = data[data.startIndex + 1]
        let pollRaw: UInt8 = data[data.startIndex + 2]
        let poll = Int8(bitPattern: pollRaw)
        let precisionRaw: UInt8 = data[data.startIndex + 3]
        let precision = Int8(bitPattern: precisionRaw)

        let rootDelay = readFixedPoint(data, offset: 4)
        let rootDispersion = readFixedPoint(data, offset: 8)

        let referenceID = decodeReferenceID(data: data, stratum: stratum)

        let referenceTimestamp = readTimestamp(data, offset: 16)
        let originTimestamp = readTimestamp(data, offset: 24)
        let receiveTimestamp = readTimestamp(data, offset: 32)
        let transmitTimestamp = readTimestamp(data, offset: 40)

        return NTPMessage(
            leapIndicator: NTPLeapIndicator(rawValue: liRaw) ?? .alarm,
            versionNumber: vnRaw,
            mode: NTPMode(rawValue: modeRaw) ?? .reserved,
            stratum: stratum,
            poll: poll,
            precision: precision,
            rootDelaySeconds: rootDelay,
            rootDispersionSeconds: rootDispersion,
            referenceID: referenceID,
            referenceTimestamp: referenceTimestamp,
            originTimestamp: originTimestamp,
            receiveTimestamp: receiveTimestamp,
            transmitTimestamp: transmitTimestamp
        )
    }

    func protocolInfo(_ data: Data) -> String? {
        guard let message = parseMessage(data) else { return nil }
        return "\(message.mode.displayName), v\(message.versionNumber), stratum \(message.stratum)"
    }

    // MARK: - Attribute lookup

    var protocolProperty: [String: @Sendable (Data) -> Any?] {
        [
            ProtocolAttribute.ntpSize.rawValue: { data in
                NSNumber(value: data.count)
            },
            ProtocolAttribute.ntpMode.rawValue: { [weak self] data in
                self?.parseMessage(data)?.mode.displayName
            },
            ProtocolAttribute.ntpStratum.rawValue: { [weak self] data in
                guard let stratum = self?.parseMessage(data)?.stratum else {
                    return nil
                }
                return NSNumber(value: stratum)
            },
            ProtocolAttribute.ntpReference.rawValue: { [weak self] data in
                self?.parseMessage(data)?.referenceID.displayValue
            },
            ProtocolAttribute.ntpProtocolInfo.rawValue: { [weak self] data in
                self?.protocolInfo(data)
            }
        ]
    }

    // MARK: - Internal

    /// Read a 32-bit signed fixed-point value (16.16) and return seconds.
    /// RFC 5905 expresses Root Delay / Root Dispersion this way.
    private func readFixedPoint(_ data: Data, offset: Int) -> Double {
        let raw: UInt32 = readUInt32(data, offset: offset)
        let signed = Int32(bitPattern: raw)
        return Double(signed) / 65536.0
    }

    private func readTimestamp(_ data: Data, offset: Int) -> NTPTimestamp {
        let seconds = readUInt32(data, offset: offset)
        let fraction = readUInt32(data, offset: offset + 4)
        return NTPTimestamp(secondsSince1900: seconds, fraction: fraction)
    }

    private func readUInt32(_ data: Data, offset: Int) -> UInt32 {
        let absoluteOffset = data.startIndex + offset
        guard absoluteOffset + 4 <= data.endIndex else { return 0 }
        let byte0 = UInt32(data[absoluteOffset])
        let byte1 = UInt32(data[absoluteOffset + 1])
        let byte2 = UInt32(data[absoluteOffset + 2])
        let byte3 = UInt32(data[absoluteOffset + 3])
        return (byte0 << 24) | (byte1 << 16) | (byte2 << 8) | byte3
    }

    private func decodeReferenceID(data: Data, stratum: UInt8) -> NTPMessage.ReferenceID {
        let absoluteOffset = data.startIndex + 12
        guard absoluteOffset + 4 <= data.endIndex else {
            return .rawHex("00000000")
        }
        let bytes: [UInt8] = [
            data[absoluteOffset],
            data[absoluteOffset + 1],
            data[absoluteOffset + 2],
            data[absoluteOffset + 3]
        ]

        switch stratum {
        case 0:
            let ascii = String(bytes: bytes.filter { $0 != 0 }, encoding: .ascii) ?? ""
            return .kissCode(ascii)
        case 1:
            let ascii = String(bytes: bytes.filter { $0 != 0 }, encoding: .ascii) ?? ""
            return .identifier(ascii)
        case 2...:
            return .ipv4("\(bytes[0]).\(bytes[1]).\(bytes[2]).\(bytes[3])")
        default:
            let hex = bytes.map { String(format: "%02x", $0) }.joined()
            return .rawHex(hex)
        }
    }
}
