import Foundation

/// A JSON value that preserves object key order.
///
/// The protocol's plaintexts and envelopes are transmitted as `JSON.stringify`
/// output: insertion-ordered keys, no whitespace, JavaScript escaping rules.
/// Foundation's `JSONSerialization` neither preserves order nor formats numbers
/// the same way, so the SDK carries its own small serializer and parser.
public indirect enum JSONValue: Sendable, Hashable {
    case null
    case bool(Bool)
    case number(Double)
    case string(String)
    case array([JSONValue])
    /// Ordered key/value pairs. Equality is order-insensitive (JSON object semantics);
    /// serialization is order-preserving (wire semantics).
    case object(JSONObject)

    public static func int(_ value: Int) -> JSONValue { .number(Double(value)) }

    public subscript(key: String) -> JSONValue? {
        if case .object(let o) = self { return o[key] }
        return nil
    }

    public var stringValue: String? { if case .string(let s) = self { return s }; return nil }
    public var doubleValue: Double? { if case .number(let n) = self { return n }; return nil }
    public var intValue: Int? {
        guard case .number(let n) = self, n.rounded() == n, abs(n) < 9.007_199_254_740_992e15 else { return nil }
        return Int(n)
    }
    public var boolValue: Bool? { if case .bool(let b) = self { return b }; return nil }
    public var arrayValue: [JSONValue]? { if case .array(let a) = self { return a }; return nil }
    public var objectValue: JSONObject? { if case .object(let o) = self { return o }; return nil }
    public var isNull: Bool { if case .null = self { return true }; return false }
}

/// An insertion-ordered JSON object.
public struct JSONObject: Sendable, Hashable, ExpressibleByDictionaryLiteral {
    public private(set) var keys: [String] = []
    private var storage: [String: JSONValue] = [:]

    public init() {}

    public init(_ pairs: [(String, JSONValue)]) {
        for (k, v) in pairs { self[k] = v }
    }

    public init(dictionaryLiteral elements: (String, JSONValue)...) {
        for (k, v) in elements { self[k] = v }
    }

    public subscript(key: String) -> JSONValue? {
        get { storage[key] }
        set {
            if let newValue {
                if storage[key] == nil { keys.append(key) }
                storage[key] = newValue
            } else if storage.removeValue(forKey: key) != nil {
                keys.removeAll { $0 == key }
            }
        }
    }

    public var pairs: [(String, JSONValue)] { keys.map { ($0, storage[$0]!) } }
    public var count: Int { keys.count }

    public static func == (lhs: JSONObject, rhs: JSONObject) -> Bool { lhs.storage == rhs.storage }
    public func hash(into hasher: inout Hasher) { hasher.combine(storage) }
}

// MARK: - Serialization (JSON.stringify compatible)

extension JSONValue {
    /// Compact serialization matching JavaScript's `JSON.stringify(value)`.
    public var serialized: String {
        var out = ""
        write(into: &out)
        return out
    }

    /// UTF-8 bytes of `serialized`.
    public var serializedData: Data { Data(serialized.utf8) }

    private func write(into out: inout String) {
        switch self {
        case .null: out += "null"
        case .bool(let b): out += b ? "true" : "false"
        case .number(let n): out += JSONValue.formatNumber(n)
        case .string(let s): JSONValue.writeString(s, into: &out)
        case .array(let a):
            out += "["
            for (i, v) in a.enumerated() {
                if i > 0 { out += "," }
                v.write(into: &out)
            }
            out += "]"
        case .object(let o):
            out += "{"
            for (i, (k, v)) in o.pairs.enumerated() {
                if i > 0 { out += "," }
                JSONValue.writeString(k, into: &out)
                out += ":"
                v.write(into: &out)
            }
            out += "}"
        }
    }

    /// JavaScript number formatting: integers without a fraction, otherwise the
    /// shortest round-trip form with a bare exponent (`1e-7`, `1e+21`), and `null`
    /// for non-finite values.
    static func formatNumber(_ n: Double) -> String {
        guard n.isFinite else { return "null" }
        if n == 0 { return "0" } // JSON.stringify(-0) === "0"
        if n.rounded() == n, abs(n) < 1e21 {
            return String(format: "%.0f", n)
        }
        var s = "\(n)"
        // Swift prints "1e-07" / "1e+21"; JavaScript prints "1e-7" / "1e+21".
        if let e = s.firstIndex(where: { $0 == "e" || $0 == "E" }) {
            let mantissa = s[..<e]
            var exp = String(s[s.index(after: e)...])
            var sign = "+"
            if exp.hasPrefix("-") { sign = "-"; exp.removeFirst() } else if exp.hasPrefix("+") { exp.removeFirst() }
            while exp.count > 1, exp.hasPrefix("0") { exp.removeFirst() }
            s = "\(mantissa)e\(sign)\(exp)"
        }
        return s
    }

    static func writeString(_ s: String, into out: inout String) {
        out += "\""
        for scalar in s.unicodeScalars {
            switch scalar {
            case "\"": out += "\\\""
            case "\\": out += "\\\\"
            case "\u{08}": out += "\\b"
            case "\u{0C}": out += "\\f"
            case "\n": out += "\\n"
            case "\r": out += "\\r"
            case "\t": out += "\\t"
            case _ where scalar.value < 0x20:
                out += String(format: "\\u%04x", scalar.value)
            default:
                out.unicodeScalars.append(scalar)
            }
        }
        out += "\""
    }
}

// MARK: - Parsing

extension JSONValue {
    public struct ParseError: Error, Sendable, CustomStringConvertible {
        public let offset: Int
        public let reason: String
        public var description: String { "JSON parse error at byte \(offset): \(reason)" }
    }

    /// Parses standard JSON (RFC 8259) into an order-preserving value.
    public static func parse(_ data: Data) throws -> JSONValue {
        var parser = Parser(bytes: Array(data))
        parser.skipWhitespace()
        let value = try parser.parseValue(depth: 0)
        parser.skipWhitespace()
        guard parser.atEnd else { throw parser.error("trailing characters") }
        return value
    }

    public static func parse(_ string: String) throws -> JSONValue { try parse(Data(string.utf8)) }

    private struct Parser {
        let bytes: [UInt8]
        var i = 0
        static let maxDepth = 128

        var atEnd: Bool { i >= bytes.count }
        func error(_ reason: String) -> ParseError { ParseError(offset: i, reason: reason) }

        mutating func skipWhitespace() {
            while i < bytes.count, [0x20, 0x09, 0x0A, 0x0D].contains(bytes[i]) { i += 1 }
        }

        mutating func parseValue(depth: Int) throws -> JSONValue {
            guard depth < Parser.maxDepth else { throw error("nesting too deep") }
            guard i < bytes.count else { throw error("unexpected end") }
            switch bytes[i] {
            case UInt8(ascii: "{"): return try parseObject(depth: depth)
            case UInt8(ascii: "["): return try parseArray(depth: depth)
            case UInt8(ascii: "\""): return .string(try parseString())
            case UInt8(ascii: "t"): try expect("true"); return .bool(true)
            case UInt8(ascii: "f"): try expect("false"); return .bool(false)
            case UInt8(ascii: "n"): try expect("null"); return .null
            default: return try parseNumber()
            }
        }

        mutating func expect(_ word: String) throws {
            let w = Array(word.utf8)
            guard i + w.count <= bytes.count, Array(bytes[i..<i + w.count]) == w else { throw error("unexpected token") }
            i += w.count
        }

        mutating func parseObject(depth: Int) throws -> JSONValue {
            i += 1
            var obj = JSONObject()
            skipWhitespace()
            if i < bytes.count, bytes[i] == UInt8(ascii: "}") { i += 1; return .object(obj) }
            while true {
                skipWhitespace()
                guard i < bytes.count, bytes[i] == UInt8(ascii: "\"") else { throw error("expected object key") }
                let key = try parseString()
                skipWhitespace()
                guard i < bytes.count, bytes[i] == UInt8(ascii: ":") else { throw error("expected ':'") }
                i += 1
                skipWhitespace()
                obj[key] = try parseValue(depth: depth + 1)
                skipWhitespace()
                guard i < bytes.count else { throw error("unterminated object") }
                if bytes[i] == UInt8(ascii: ",") { i += 1; continue }
                if bytes[i] == UInt8(ascii: "}") { i += 1; return .object(obj) }
                throw error("expected ',' or '}'")
            }
        }

        mutating func parseArray(depth: Int) throws -> JSONValue {
            i += 1
            var arr: [JSONValue] = []
            skipWhitespace()
            if i < bytes.count, bytes[i] == UInt8(ascii: "]") { i += 1; return .array(arr) }
            while true {
                skipWhitespace()
                arr.append(try parseValue(depth: depth + 1))
                skipWhitespace()
                guard i < bytes.count else { throw error("unterminated array") }
                if bytes[i] == UInt8(ascii: ",") { i += 1; continue }
                if bytes[i] == UInt8(ascii: "]") { i += 1; return .array(arr) }
                throw error("expected ',' or ']'")
            }
        }

        mutating func parseString() throws -> String {
            i += 1
            var scalars = String.UnicodeScalarView()
            var pendingHigh: UInt32? = nil
            func flushPending() throws {
                if pendingHigh != nil { throw ParseError(offset: 0, reason: "lone surrogate") }
            }
            while true {
                guard i < bytes.count else { throw error("unterminated string") }
                let b = bytes[i]
                if b == UInt8(ascii: "\"") {
                    try flushPending()
                    i += 1
                    return String(scalars)
                }
                if b == UInt8(ascii: "\\") {
                    i += 1
                    guard i < bytes.count else { throw error("bad escape") }
                    let e = bytes[i]
                    i += 1
                    var scalar: Unicode.Scalar?
                    switch e {
                    case UInt8(ascii: "\""): scalar = "\""
                    case UInt8(ascii: "\\"): scalar = "\\"
                    case UInt8(ascii: "/"): scalar = "/"
                    case UInt8(ascii: "b"): scalar = "\u{08}"
                    case UInt8(ascii: "f"): scalar = "\u{0C}"
                    case UInt8(ascii: "n"): scalar = "\n"
                    case UInt8(ascii: "r"): scalar = "\r"
                    case UInt8(ascii: "t"): scalar = "\t"
                    case UInt8(ascii: "u"):
                        guard i + 4 <= bytes.count, let code = UInt32(String(decoding: bytes[i..<i + 4], as: UTF8.self), radix: 16) else { throw error("bad \\u escape") }
                        i += 4
                        if (0xD800...0xDBFF).contains(code) {
                            try flushPending()
                            pendingHigh = code
                            continue
                        }
                        if (0xDC00...0xDFFF).contains(code) {
                            guard let hi = pendingHigh else { throw error("lone low surrogate") }
                            pendingHigh = nil
                            let combined = 0x10000 + ((hi - 0xD800) << 10) + (code - 0xDC00)
                            scalar = Unicode.Scalar(combined)
                        } else {
                            try flushPending()
                            scalar = Unicode.Scalar(code)
                        }
                    default: throw error("bad escape")
                    }
                    guard let s = scalar else { throw error("bad escape") }
                    scalars.append(s)
                    continue
                }
                try flushPending()
                if b < 0x20 { throw error("control character in string") }
                // decode one UTF-8 scalar
                let len: Int = b < 0x80 ? 1 : (b >> 5 == 0b110 ? 2 : (b >> 4 == 0b1110 ? 3 : (b >> 3 == 0b11110 ? 4 : 0)))
                guard len > 0, i + len <= bytes.count else { throw error("invalid UTF-8") }
                guard let str = String(bytes: bytes[i..<i + len], encoding: .utf8), let sc = str.unicodeScalars.first, str.unicodeScalars.count == 1 else { throw error("invalid UTF-8") }
                scalars.append(sc)
                i += len
            }
        }

        mutating func parseNumber() throws -> JSONValue {
            let start = i
            if i < bytes.count, bytes[i] == UInt8(ascii: "-") { i += 1 }
            guard i < bytes.count, (0x30...0x39).contains(bytes[i]) else { throw error("expected value") }
            if bytes[i] == 0x30 { i += 1 } else { while i < bytes.count, (0x30...0x39).contains(bytes[i]) { i += 1 } }
            if i < bytes.count, bytes[i] == UInt8(ascii: ".") {
                i += 1
                guard i < bytes.count, (0x30...0x39).contains(bytes[i]) else { throw error("bad fraction") }
                while i < bytes.count, (0x30...0x39).contains(bytes[i]) { i += 1 }
            }
            if i < bytes.count, bytes[i] == UInt8(ascii: "e") || bytes[i] == UInt8(ascii: "E") {
                i += 1
                if i < bytes.count, bytes[i] == UInt8(ascii: "+") || bytes[i] == UInt8(ascii: "-") { i += 1 }
                guard i < bytes.count, (0x30...0x39).contains(bytes[i]) else { throw error("bad exponent") }
                while i < bytes.count, (0x30...0x39).contains(bytes[i]) { i += 1 }
            }
            guard let d = Double(String(decoding: bytes[start..<i], as: UTF8.self)) else { throw error("bad number") }
            return .number(d)
        }
    }
}

// MARK: - Codable bridging

extension JSONValue {
    /// Encodes any `Encodable` into a `JSONValue` (key order follows the encoder's output order).
    public init<T: Encodable>(encoding value: T) throws {
        let encoder = JSONEncoder()
        encoder.outputFormatting = [.sortedKeys]
        self = try JSONValue.parse(try encoder.encode(value))
    }

    /// Decodes this value into a `Decodable` type.
    public func decode<T: Decodable>(_ type: T.Type) throws -> T {
        try JSONDecoder().decode(type, from: serializedData)
    }
}