import Foundation

/// A decoded CBOR data item (RFC 8949), used as the intermediate tree between bytes and
/// `Codable`. We only need the subset the sync wire layer exercises: the encoder emits
/// definite-length items and float64s; the parser accepts that subset and rejects the rest.
///
/// Why CBOR (SYNC_PROTOCOL §5): the event-heavy transcript path is more compact than JSON,
/// and the same `Codable` types serialize over both LAN and the future relay unchanged.
enum CBORValue: Equatable {
    case unsigned(UInt64)        // major 0
    case negative(UInt64)        // major 1 — represents -1 - n
    case bytes([UInt8])          // major 2
    case text(String)            // major 3
    case array([CBORValue])      // major 4
    case map([(CBORValue, CBORValue)])  // major 5 (text keys in practice)
    case bool(Bool)              // major 7 / 20,21
    case null                    // major 7 / 22
    case undefined               // major 7 / 23
    case float(Double)           // major 7 / 27 (float64)

    static func == (lhs: CBORValue, rhs: CBORValue) -> Bool {
        switch (lhs, rhs) {
        case (.unsigned(let a), .unsigned(let b)): return a == b
        case (.negative(let a), .negative(let b)): return a == b
        case (.bytes(let a), .bytes(let b)): return a == b
        case (.text(let a), .text(let b)): return a == b
        case (.array(let a), .array(let b)): return a == b
        case (.map(let a), .map(let b)):
            guard a.count == b.count else { return false }
            return zip(a, b).allSatisfy { $0.0 == $1.0 && $0.1 == $1.1 }
        case (.bool(let a), .bool(let b)): return a == b
        case (.null, .null), (.undefined, .undefined): return true
        case (.float(let a), .float(let b)): return a.bitPattern == b.bitPattern
        default: return false
        }
    }
}

// MARK: - Serialization (deterministic: definite lengths, sorted map keys, float64)

extension CBORValue {
    func serialized() -> [UInt8] {
        var out: [UInt8] = []
        write(into: &out)
        return out
    }

    private func write(into out: inout [UInt8]) {
        switch self {
        case .unsigned(let n):
            CBORValue.writeHead(major: 0, value: n, into: &out)
        case .negative(let n):
            CBORValue.writeHead(major: 1, value: n, into: &out)
        case .bytes(let b):
            CBORValue.writeHead(major: 2, value: UInt64(b.count), into: &out)
            out.append(contentsOf: b)
        case .text(let s):
            let u = Array(s.utf8)
            CBORValue.writeHead(major: 3, value: UInt64(u.count), into: &out)
            out.append(contentsOf: u)
        case .array(let items):
            CBORValue.writeHead(major: 4, value: UInt64(items.count), into: &out)
            for item in items { item.write(into: &out) }
        case .map(let pairs):
            // Deterministic ordering: sort by the encoded key bytes. Keeps
            // `[String: JSONValue]` (nondeterministic dict iteration) stable on the wire.
            let encoded = pairs.map { (key, value) -> (k: [UInt8], v: CBORValue) in
                (key.serialized(), value)
            }.sorted { lexLess($0.k, $1.k) }
            CBORValue.writeHead(major: 5, value: UInt64(encoded.count), into: &out)
            for pair in encoded {
                out.append(contentsOf: pair.k)
                pair.v.write(into: &out)
            }
        case .bool(let b):
            out.append(b ? 0xf5 : 0xf4)
        case .null:
            out.append(0xf6)
        case .undefined:
            out.append(0xf7)
        case .float(let d):
            out.append(0xfb)
            var be = d.bitPattern.bigEndian
            withUnsafeBytes(of: &be) { out.append(contentsOf: $0) }
        }
    }

    private static func writeHead(major: UInt8, value: UInt64, into out: inout [UInt8]) {
        let high = major << 5
        switch value {
        case 0..<24:
            out.append(high | UInt8(value))
        case 24..<0x100:
            out.append(high | 24)
            out.append(UInt8(value))
        case 0x100..<0x1_0000:
            out.append(high | 25)
            out.append(UInt8(value >> 8))
            out.append(UInt8(value & 0xff))
        case 0x1_0000..<0x1_0000_0000:
            out.append(high | 26)
            for shift in stride(from: 24, through: 0, by: -8) {
                out.append(UInt8((value >> UInt64(shift)) & 0xff))
            }
        default:
            out.append(high | 27)
            for shift in stride(from: 56, through: 0, by: -8) {
                out.append(UInt8((value >> UInt64(shift)) & 0xff))
            }
        }
    }

    private func lexLess(_ a: [UInt8], _ b: [UInt8]) -> Bool {
        for (x, y) in zip(a, b) where x != y { return x < y }
        return a.count < b.count
    }
}

// MARK: - Parsing

enum CBORParseError: Error, Equatable {
    case truncated
    case unsupportedAdditionalInfo(major: UInt8, info: UInt8)
    case indefiniteLengthUnsupported
    case invalidUTF8
    case trailingBytes
}

extension CBORValue {
    /// Parse exactly one top-level item; throws if trailing bytes remain.
    static func parse(_ bytes: [UInt8]) throws -> CBORValue {
        var cursor = 0
        let value = try parseItem(bytes, &cursor)
        guard cursor == bytes.count else { throw CBORParseError.trailingBytes }
        return value
    }

    private static func parseItem(_ bytes: [UInt8], _ cursor: inout Int) throws -> CBORValue {
        guard cursor < bytes.count else { throw CBORParseError.truncated }
        let initial = bytes[cursor]; cursor += 1
        let major = initial >> 5
        let info = initial & 0x1f

        switch major {
        case 0:
            return .unsigned(try readLength(bytes, &cursor, info: info, major: major))
        case 1:
            return .negative(try readLength(bytes, &cursor, info: info, major: major))
        case 2:
            let count = Int(try readLength(bytes, &cursor, info: info, major: major))
            return .bytes(try take(bytes, &cursor, count))
        case 3:
            let count = Int(try readLength(bytes, &cursor, info: info, major: major))
            let raw = try take(bytes, &cursor, count)
            guard let s = String(bytes: raw, encoding: .utf8) else { throw CBORParseError.invalidUTF8 }
            return .text(s)
        case 4:
            let count = Int(try readLength(bytes, &cursor, info: info, major: major))
            var items: [CBORValue] = []
            items.reserveCapacity(count)
            for _ in 0..<count { items.append(try parseItem(bytes, &cursor)) }
            return .array(items)
        case 5:
            let count = Int(try readLength(bytes, &cursor, info: info, major: major))
            var pairs: [(CBORValue, CBORValue)] = []
            pairs.reserveCapacity(count)
            for _ in 0..<count {
                let k = try parseItem(bytes, &cursor)
                let v = try parseItem(bytes, &cursor)
                pairs.append((k, v))
            }
            return .map(pairs)
        case 7:
            switch info {
            case 20: return .bool(false)
            case 21: return .bool(true)
            case 22: return .null
            case 23: return .undefined
            case 25:  // float16
                let raw = try take(bytes, &cursor, 2)
                return .float(decodeFloat16(raw))
            case 26:  // float32
                let raw = try take(bytes, &cursor, 4)
                let bits = raw.reduce(UInt32(0)) { ($0 << 8) | UInt32($1) }
                return .float(Double(Float(bitPattern: bits)))
            case 27:  // float64
                let raw = try take(bytes, &cursor, 8)
                let bits = raw.reduce(UInt64(0)) { ($0 << 8) | UInt64($1) }
                return .float(Double(bitPattern: bits))
            default:
                throw CBORParseError.unsupportedAdditionalInfo(major: major, info: info)
            }
        default:
            throw CBORParseError.unsupportedAdditionalInfo(major: major, info: info)
        }
    }

    private static func readLength(
        _ bytes: [UInt8], _ cursor: inout Int, info: UInt8, major: UInt8
    ) throws -> UInt64 {
        switch info {
        case 0..<24:
            return UInt64(info)
        case 24:
            return UInt64(try take(bytes, &cursor, 1)[0])
        case 25:
            return try take(bytes, &cursor, 2).reduce(UInt64(0)) { ($0 << 8) | UInt64($1) }
        case 26:
            return try take(bytes, &cursor, 4).reduce(UInt64(0)) { ($0 << 8) | UInt64($1) }
        case 27:
            return try take(bytes, &cursor, 8).reduce(UInt64(0)) { ($0 << 8) | UInt64($1) }
        case 31:
            throw CBORParseError.indefiniteLengthUnsupported
        default:
            throw CBORParseError.unsupportedAdditionalInfo(major: major, info: info)
        }
    }

    private static func take(_ bytes: [UInt8], _ cursor: inout Int, _ count: Int) throws -> [UInt8] {
        guard count >= 0, cursor + count <= bytes.count else { throw CBORParseError.truncated }
        defer { cursor += count }
        return Array(bytes[cursor..<cursor + count])
    }

    private static func decodeFloat16(_ raw: [UInt8]) -> Double {
        let bits = UInt16(raw[0]) << 8 | UInt16(raw[1])
        let sign = (bits >> 15) & 0x1
        let exp = (bits >> 10) & 0x1f
        let frac = bits & 0x3ff
        let value: Double
        if exp == 0 {
            value = Double(frac) * pow(2.0, -24)
        } else if exp == 0x1f {
            value = frac == 0 ? .infinity : .nan
        } else {
            value = (1.0 + Double(frac) / 1024.0) * pow(2.0, Double(Int(exp) - 15))
        }
        return sign == 1 ? -value : value
    }
}
