import Foundation

/// The canonical JSON form this protocol hashes.
///
/// This is a hand-written parser rather than `JSONSerialization`, and the reason is not preference:
/// the canonical form is defined over the *spelling* of the input, and every Foundation decoder
/// erases spelling before the caller sees anything. `1.0` and `1` both arrive as `NSNumber(1)`;
/// a lone surrogate is either rejected outright or silently repaired depending on the framework
/// version. Trying to recover the original text from a decoded tree is how two languages end up
/// hashing different bytes for the same body, which surfaces as an intent mismatch the phone
/// reports as "confirmation failed" with no way to tell which side was wrong.
///
/// The rules, which the Go agent and the Node reference implementation implement identically, are
/// documented in `agent/internal/wire/json.go`. The short version: keys sorted by UTF-8 bytes,
/// no insignificant whitespace, integers only in canonical spelling, strings escaped as `\" \\ \n
/// \r \t \b \f`, `\u00XX` for other control characters, literal UTF-8 for everything else, and a
/// lone surrogate escape is an error rather than a replacement character.
public enum CanonicalJSON {
    public enum Failure: Error, Equatable {
        case invalidJSON
        case trailingData
        case nonCanonicalNumber(String)
        case loneSurrogate(UInt32)
        case unsupportedValue
    }

    /// Canonicalises a JSON document. An empty or whitespace-only body canonicalises to no bytes at
    /// all, not to the four characters `null`, so "no body" and "empty body" cannot be confused.
    public static func canonicalize(_ text: String) throws -> Data {
        let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !trimmed.isEmpty else { return Data() }
        var parser = Parser(Array(trimmed.unicodeScalars))
        let value = try parser.parseValue()
        parser.skipWhitespace()
        guard parser.isAtEnd else { throw Failure.trailingData }
        var out = Data()
        try serialize(value, into: &out)
        return out
    }

    /// Canonicalises raw body bytes, which is what the intent digest needs.
    public static func canonicalize(_ data: Data) throws -> Data {
        guard let text = String(data: data, encoding: .utf8) else { throw Failure.invalidJSON }
        return try canonicalize(text)
    }

    // MARK: - Value tree

    indirect enum Value {
        case null
        case bool(Bool)
        case integer(String)
        case string(String)
        case array([Value])
        case object([(key: String, value: Value)])
    }

    // MARK: - Serialisation

    static func serialize(_ value: Value, into out: inout Data) throws {
        switch value {
        case .null:
            out.append(contentsOf: Array("null".utf8))
        case .bool(let flag):
            out.append(contentsOf: Array((flag ? "true" : "false").utf8))
        case .integer(let literal):
            out.append(contentsOf: Array(literal.utf8))
        case .string(let text):
            escape(text, into: &out)
        case .array(let items):
            out.append(UInt8(ascii: "["))
            for (index, item) in items.enumerated() {
                if index > 0 { out.append(UInt8(ascii: ",")) }
                try serialize(item, into: &out)
            }
            out.append(UInt8(ascii: "]"))
        case .object(let members):
            // UTF-8 byte order, which is what the Go agent's sort.Strings produces. Two shortcuts
            // are wrong here and both look right on ASCII: Swift's default String `<` uses Unicode
            // collation (so "é" sorts before "z"), and Foundation's `.sortedKeys` does the same.
            let sorted = members.sorted { Array($0.key.utf8).lexicographicallyPrecedes(Array($1.key.utf8)) }
            out.append(UInt8(ascii: "{"))
            for (index, member) in sorted.enumerated() {
                if index > 0 { out.append(UInt8(ascii: ",")) }
                escape(member.key, into: &out)
                out.append(UInt8(ascii: ":"))
                try serialize(member.value, into: &out)
            }
            out.append(UInt8(ascii: "}"))
        }
    }

    /// Writes a JSON string literal in the canonical spelling.
    ///
    /// The two rules worth stating: control characters below 0x20 use the short forms JSON names
    /// where they exist and `\u00XX` otherwise (matching Go's encoder with HTML escaping off), and
    /// everything else — including U+2028, U+2029, DEL and astral characters — is emitted as
    /// literal UTF-8. Go's encoder escapes U+2028/U+2029 even with HTML escaping disabled, which is
    /// why the Go side had to grow its own writer too; the vectors pin the literal form.
    static func escape(_ text: String, into out: inout Data) {
        out.append(UInt8(ascii: "\""))
        for scalar in text.unicodeScalars {
            switch scalar {
            case "\"":
                out.append(contentsOf: Array("\\\"".utf8))
            case "\\":
                out.append(contentsOf: Array("\\\\".utf8))
            case "\n":
                out.append(contentsOf: Array("\\n".utf8))
            case "\r":
                out.append(contentsOf: Array("\\r".utf8))
            case "\t":
                out.append(contentsOf: Array("\\t".utf8))
            case Unicode.Scalar(0x08):
                out.append(contentsOf: Array("\\b".utf8))
            case Unicode.Scalar(0x0C):
                out.append(contentsOf: Array("\\f".utf8))
            default:
                if scalar.value < 0x20 {
                    let hex = String(format: "%04x", scalar.value)
                    out.append(contentsOf: Array("\\u\(hex)".utf8))
                } else {
                    out.append(contentsOf: Array(String(scalar).utf8))
                }
            }
        }
        out.append(UInt8(ascii: "\""))
    }

    // MARK: - Parser

    struct Parser {
        private let scalars: [Unicode.Scalar]
        private var index = 0

        init(_ scalars: [Unicode.Scalar]) {
            self.scalars = scalars
        }

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

        mutating func skipWhitespace() {
            while index < scalars.count, isSpace(scalars[index]) { index += 1 }
        }

        private func isSpace(_ scalar: Unicode.Scalar) -> Bool {
            scalar == " " || scalar == "\t" || scalar == "\n" || scalar == "\r"
        }

        private func peek() -> Unicode.Scalar? {
            index < scalars.count ? scalars[index] : nil
        }

        private mutating func expect(_ scalar: Unicode.Scalar) throws {
            guard peek() == scalar else { throw Failure.invalidJSON }
            index += 1
        }

        mutating func parseValue() throws -> Value {
            skipWhitespace()
            guard let scalar = peek() else { throw Failure.invalidJSON }
            switch scalar {
            case "{": return try parseObject()
            case "[": return try parseArray()
            case "\"": return .string(try parseString())
            case "t":
                try expectWord("true")
                return .bool(true)
            case "f":
                try expectWord("false")
                return .bool(false)
            case "n":
                try expectWord("null")
                return .null
            default:
                return .integer(try parseNumber())
            }
        }

        private mutating func expectWord(_ word: String) throws {
            for scalar in word.unicodeScalars {
                try expect(scalar)
            }
        }

        private mutating func parseObject() throws -> Value {
            try expect("{")
            var members: [(key: String, value: Value)] = []
            skipWhitespace()
            if peek() == "}" {
                index += 1
                return .object(members)
            }
            while true {
                skipWhitespace()
                let key = try parseString()
                skipWhitespace()
                try expect(":")
                let value = try parseValue()
                // A repeated key takes the last occurrence, which is what Go's decoder and
                // JSON.parse both do. Refusing here would make the phone stricter than the agent
                // that recomputes the intent digest, and the two disagreeing about a body is the
                // failure this whole file exists to prevent.
                if let existing = members.firstIndex(where: { $0.key == key }) {
                    members[existing] = (key, value)
                } else {
                    members.append((key, value))
                }
                skipWhitespace()
                if peek() == "," {
                    index += 1
                    continue
                }
                try expect("}")
                return .object(members)
            }
        }

        private mutating func parseArray() throws -> Value {
            try expect("[")
            var items: [Value] = []
            skipWhitespace()
            if peek() == "]" {
                index += 1
                return .array(items)
            }
            while true {
                items.append(try parseValue())
                skipWhitespace()
                if peek() == "," {
                    index += 1
                    continue
                }
                try expect("]")
                return .array(items)
            }
        }

        /// Parses a string, resolving escapes.
        ///
        /// A `\uXXXX` escape in the surrogate range must be correctly paired. A lone surrogate is
        /// an error here for the same reason it is in Go and Node: the three frameworks disagree
        /// about what to do with one, so the only rule all three can share is to accept none.
        private mutating func parseString() throws -> String {
            try expect("\"")
            var output = String.UnicodeScalarView()
            while true {
                guard let scalar = peek() else { throw Failure.invalidJSON }
                index += 1
                switch scalar {
                case "\"":
                    return String(output)
                case "\\":
                    output.append(try parseEscape())
                default:
                    if scalar.value < 0x20 { throw Failure.invalidJSON }
                    output.append(scalar)
                }
            }
        }

        private mutating func parseEscape() throws -> Unicode.Scalar {
            guard let scalar = peek() else { throw Failure.invalidJSON }
            index += 1
            switch scalar {
            case "\"": return "\""
            case "\\": return "\\"
            case "/": return "/"
            case "b": return Unicode.Scalar(0x08)
            case "f": return Unicode.Scalar(0x0C)
            case "n": return "\n"
            case "r": return "\r"
            case "t": return "\t"
            case "u":
                let high = try parseHex4()
                if high >= 0xD800, high <= 0xDBFF {
                    guard peek() == "\\" else { throw Failure.loneSurrogate(high) }
                    index += 1
                    guard peek() == "u" else { throw Failure.loneSurrogate(high) }
                    index += 1
                    let low = try parseHex4()
                    guard low >= 0xDC00, low <= 0xDFFF else { throw Failure.loneSurrogate(high) }
                    let combined = 0x10000 + ((high - 0xD800) << 10) + (low - 0xDC00)
                    guard let scalar = Unicode.Scalar(combined) else { throw Failure.loneSurrogate(high) }
                    return scalar
                }
                guard high < 0xDC00 || high > 0xDFFF else { throw Failure.loneSurrogate(high) }
                guard let scalar = Unicode.Scalar(high) else { throw Failure.loneSurrogate(high) }
                return scalar
            default:
                throw Failure.invalidJSON
            }
        }

        private mutating func parseHex4() throws -> UInt32 {
            var value: UInt32 = 0
            for _ in 0..<4 {
                guard let scalar = peek(), let digit = hexDigit(scalar) else { throw Failure.invalidJSON }
                index += 1
                value = value << 4 | digit
            }
            return value
        }

        private func hexDigit(_ scalar: Unicode.Scalar) -> UInt32? {
            switch scalar {
            case "0"..."9": return scalar.value - 0x30
            case "a"..."f": return scalar.value - 0x61 + 10
            case "A"..."F": return scalar.value - 0x41 + 10
            default: return nil
            }
        }

        /// Parses an integer in the one spelling this protocol hashes.
        ///
        /// The literal is kept as text rather than converted to a number, so what gets hashed is
        /// what the phone sent. `1.0`, `1e3`, `+1`, `.5`, `01` and `-0` are all rejected: each of
        /// them is the same number as some accepted spelling, and "which spelling did the phone
        /// hash" is not a question this protocol leaves open.
        private mutating func parseNumber() throws -> String {
            var literal = ""
            if peek() == "-" {
                literal.append("-")
                index += 1
            }
            guard let first = peek(), first >= "0", first <= "9" else {
                throw Failure.invalidJSON
            }
            if first == "0" {
                literal.append("0")
                index += 1
                if let next = peek(), next >= "0", next <= "9" {
                    throw Failure.nonCanonicalNumber(literal + String(next))
                }
            } else {
                while let next = peek(), next >= "0", next <= "9" {
                    literal.append(Character(next))
                    index += 1
                }
            }
            if let next = peek(), next == "." || next == "e" || next == "E" || next == "+" {
                throw Failure.nonCanonicalNumber(literal + String(next))
            }
            if literal == "-0" { throw Failure.nonCanonicalNumber(literal) }
            return literal
        }
    }
}