//
//  CBOR.swift
//  AT-IRC
//
//  Minimal CBOR decoder for atirc's Event Stream v0 framing
//  (two concatenated CBOR values: header object + body object).
//

import Foundation

indirect enum CBORValue {
    case unsigned(UInt64)
    case negative(Int64)
    case bytes(Data)
    case text(String)
    case array([CBORValue])
    case map([(CBORValue, CBORValue)])
    case bool(Bool)
    case null
    case float(Double)
    case tag(UInt64, CBORValue)
    case undefined

    var string: String? {
        if case .text(let value) = self { return value }
        return nil
    }

    var int: Int? {
        switch self {
        case .unsigned(let value): return Int(value)
        case .negative(let value): return Int(value)
        default: return nil
        }
    }

    subscript(key: String) -> CBORValue? {
        guard case .map(let pairs) = self else { return nil }
        for (candidate, value) in pairs where candidate.string == key {
            return value
        }
        return nil
    }

    /// Convert into JSON-compatible Foundation objects.
    func jsonObject() -> Any {
        switch self {
        case .unsigned(let value): return value
        case .negative(let value): return value
        case .bytes(let data): return Base64URL.encode(data)
        case .text(let value): return value
        case .array(let values): return values.map { $0.jsonObject() }
        case .map(let pairs):
            var result: [String: Any] = [:]
            for (key, value) in pairs {
                if let name = key.string {
                    result[name] = value.jsonObject()
                }
            }
            return result
        case .bool(let value): return value
        case .null, .undefined: return NSNull()
        case .float(let value): return value
        case .tag(_, let value): return value.jsonObject()
        }
    }
}

struct CBORDecodingError: Error, CustomStringConvertible {
    let description: String
}

struct CBORDecoder {
    private let bytes: [UInt8]
    private var index = 0

    init(_ data: Data) {
        bytes = [UInt8](data)
    }

    var isAtEnd: Bool { index >= bytes.count }

    mutating func decode() throws -> CBORValue {
        let initial = try readByte()
        let major = initial >> 5
        let additional = initial & 0x1F

        switch major {
        case 0:
            return .unsigned(try readLength(additional))
        case 1:
            let magnitude = try readLength(additional)
            return .negative(-1 - Int64(magnitude))
        case 2:
            return .bytes(try readBytes(Int(try readLength(additional))))
        case 3:
            let data = try readBytes(Int(try readLength(additional)))
            return .text(String(decoding: data, as: UTF8.self))
        case 4:
            let count = try readLength(additional)
            var values: [CBORValue] = []
            values.reserveCapacity(Int(count))
            for _ in 0..<count { values.append(try decode()) }
            return .array(values)
        case 5:
            let count = try readLength(additional)
            var pairs: [(CBORValue, CBORValue)] = []
            pairs.reserveCapacity(Int(count))
            for _ in 0..<count {
                let key = try decode()
                let value = try decode()
                pairs.append((key, value))
            }
            return .map(pairs)
        case 6:
            let tag = try readLength(additional)
            return .tag(tag, try decode())
        case 7:
            switch additional {
            case 20: return .bool(false)
            case 21: return .bool(true)
            case 22: return .null
            case 23: return .undefined
            case 25: return .float(Double(try readUInt16()))
            case 26:
                let bits = try readUInt32()
                return .float(Double(Float(bitPattern: bits)))
            case 27:
                let bits = try readUInt64()
                return .float(Double(bitPattern: bits))
            default:
                if additional < 24 { return .undefined }
                throw CBORDecodingError(description: "unsupported simple value \(additional)")
            }
        default:
            throw CBORDecodingError(description: "unsupported major type \(major)")
        }
    }

    private mutating func readByte() throws -> UInt8 {
        guard index < bytes.count else {
            throw CBORDecodingError(description: "unexpected end of CBOR input")
        }
        defer { index += 1 }
        return bytes[index]
    }

    private mutating func readBytes(_ count: Int) throws -> Data {
        guard count >= 0, index + count <= bytes.count else {
            throw CBORDecodingError(description: "CBOR string length out of bounds")
        }
        defer { index += count }
        return Data(bytes[index..<(index + count)])
    }

    private mutating func readUInt16() throws -> UInt16 {
        let data = try readBytes(2)
        return data.reduce(0) { ($0 << 8) | UInt16($1) }
    }

    private mutating func readUInt32() throws -> UInt32 {
        let data = try readBytes(4)
        return data.reduce(0) { ($0 << 8) | UInt32($1) }
    }

    private mutating func readUInt64() throws -> UInt64 {
        let data = try readBytes(8)
        return data.reduce(0) { ($0 << 8) | UInt64($1) }
    }

    private mutating func readLength(_ additional: UInt8) throws -> UInt64 {
        switch additional {
        case 0...23: return UInt64(additional)
        case 24: return UInt64(try readByte())
        case 25: return UInt64(try readUInt16())
        case 26: return UInt64(try readUInt32())
        case 27: return try readUInt64()
        case 31: throw CBORDecodingError(description: "indefinite-length CBOR is not supported")
        default: throw CBORDecodingError(description: "invalid additional info \(additional)")
        }
    }
}

/// One atirc WebSocket frame: Event Stream v0 header + body.
struct AtircFrame {
    let op: Int?
    let type: String?
    let error: String?
    let body: [String: Any]

    static func decode(_ data: Data) throws -> AtircFrame {
        var decoder = CBORDecoder(data)
        let header = try decoder.decode()
        let body = try decoder.decode()
        let headerObject = header.jsonObject() as? [String: Any] ?? [:]
        let bodyObject = body.jsonObject() as? [String: Any] ?? [:]
        return AtircFrame(
            op: headerObject["op"] as? Int,
            type: headerObject["t"] as? String,
            error: headerObject["error"] as? String,
            body: bodyObject
        )
    }
}