package wire

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"regexp"
	"sort"
)

// marshal is encoding/json with HTML escaping off.
//
// It is used for the envelope's own fields (ids, base64), never for payload strings: Go's encoder
// escapes U+2028 and U+2029 even with HTML escaping disabled, while Node's JSON.stringify and
// Swift's JSONSerialization both emit them literally. That difference is invisible in a hex dump of
// ASCII payloads and shows up as an intent mismatch the first time an SMS body contains a line
// separator, so payload strings go through escapeCanonicalString instead.
func marshal(value any) ([]byte, error) {
	buf := bytes.NewBuffer(nil)
	encoder := json.NewEncoder(buf)
	encoder.SetEscapeHTML(false)
	if err := encoder.Encode(value); err != nil {
		return nil, err
	}
	return bytes.TrimRight(buf.Bytes(), "\n"), nil
}

// escapeCanonicalString writes a JSON string literal following the documented canonical rules.
//
// It exists because Go's encoder is not a reference implementation for this protocol: it escapes
// U+2028/U+2029 (legal, unnecessary, and not what the other two languages do) and it silently
// rewrites lone surrogates to U+FFFD (which turns "reject this input" into "hash something else").
// Both are the kind of difference that produces a signature failure with no useful cause.
func escapeCanonicalString(buf *bytes.Buffer, value string) {
	const hexDigits = "0123456789abcdef"
	buf.WriteByte('"')
	for _, r := range value {
		switch r {
		case '"':
			buf.WriteString(`\"`)
		case '\\':
			buf.WriteString(`\\`)
		case '\n':
			buf.WriteString(`\n`)
		case '\r':
			buf.WriteString(`\r`)
		case '\t':
			buf.WriteString(`\t`)
		case '\b':
			buf.WriteString(`\b`)
		case '\f':
			buf.WriteString(`\f`)
		default:
			switch {
			case r < 0x20:
				// The short forms above cover the ones JSON has names for; everything else below
				// 0x20 is \u00XX, which is what both other languages emit.
				buf.WriteString(`\u00`)
				buf.WriteByte(hexDigits[(r>>4)&0xF])
				buf.WriteByte(hexDigits[r&0xF])
			default:
				// Literal UTF-8, including DEL, U+FFFD and U+2028/U+2029.
				buf.WriteRune(r)
			}
		}
	}
	buf.WriteByte('"')
}

func unmarshal(raw []byte, value any) error {
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.UseNumber()
	if err := decoder.Decode(value); err != nil {
		return err
	}
	// Trailing content would mean the caller is looking at a different document than the one that
	// was decoded, which for a signed field is a correctness bug rather than a style choice.
	var extra any
	if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) {
		return errors.New("wire: trailing data after JSON value")
	}
	return nil
}

// CanonicalJSON renders a JSON value in the one form this protocol hashes.
//
// The rules are deliberately small, and every implementation (Go agent, Node relay, Swift app)
// must follow exactly these:
//
//  1. Object keys are sorted by UTF-8 byte order and emitted without whitespace.
//  2. Arrays keep their order.
//  3. Strings use the same escapes as Go's encoder with HTML escaping disabled: \" \\ \n \r \t
//     \b \f, \u00XX for other control characters, and literal UTF-8 for everything else —
//     including U+2028 and U+2029, which Go's encoder escapes and the other two implementations
//     do not. Lone surrogates are rejected outright (rule 6).
//  4. Numbers must be integers in canonical form: an optional minus, no leading zero, no fraction,
//     no exponent, no plus. This is stricter than JSON allows and on purpose. Any wider grammar has
//     values that are the same number but different bytes (1 and 1.0, 1e2 and 100), and "which
//     spelling did the phone hash" is not a question this protocol is willing to leave open. The
//     APIs behind it take integers, strings, booleans, arrays and objects, so nothing is lost.
//  5. null, true and false are literal.
//  6. A lone surrogate escape (\uD800 without its pair, or an unpaired low surrogate) is an error.
//     Go's decoder rewrites it to U+FFFD, Node keeps the escape text, and Foundation refuses to
//     parse it at all: three behaviours, so the only way the three implementations can agree is to
//     accept none of them.
//
// An empty or absent body canonicalises to the empty byte string, not to the four characters
// "null", so "no body" and "empty body" cannot be confused.
func CanonicalJSON(raw []byte) ([]byte, error) {
	trimmed := bytes.TrimSpace(raw)
	if len(trimmed) == 0 {
		return nil, nil
	}
	if err := rejectLoneSurrogates(trimmed); err != nil {
		return nil, err
	}
	decoder := json.NewDecoder(bytes.NewReader(trimmed))
	decoder.UseNumber()
	var value any
	if err := decoder.Decode(&value); err != nil {
		return nil, fmt.Errorf("wire: body is not valid JSON: %w", err)
	}
	var extra any
	if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) {
		return nil, errors.New("wire: body has trailing data after the JSON value")
	}
	buf := bytes.NewBuffer(make([]byte, 0, len(trimmed)))
	if err := writeCanonical(buf, value); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// rejectLoneSurrogates scans the raw text for a \u escape in the surrogate range that is not
// correctly paired.
//
// It works on the encoded bytes rather than on the decoded string because decoding is exactly what
// destroys the evidence: Go replaces the lone surrogate with U+FFFD, after which "the caller sent
// an invalid escape" and "the caller sent U+FFFD" are the same bytes and only one of them is a
// legitimate thing to hash.
//
// The scan is escape-aware rather than a byte search: in the JSON text `"\\ud800"` the backslash is
// escaped, so the four characters after it are literal text and not a surrogate escape. Walking
// escapes one at a time is what distinguishes the two.
func rejectLoneSurrogates(raw []byte) error {
	inString := false
	for i := 0; i < len(raw); i++ {
		switch raw[i] {
		case '"':
			inString = !inString
		case '\\':
			if !inString {
				continue
			}
			if i+1 >= len(raw) {
				return errors.New("wire: JSON string ends with a dangling backslash")
			}
			if raw[i+1] != 'u' {
				i++
				continue
			}
			code, ok := hex4(raw, i+2)
			if !ok {
				return errors.New("wire: malformed \\u escape in JSON string")
			}
			switch {
			case code >= 0xDC00 && code <= 0xDFFF:
				return fmt.Errorf("wire: unpaired low surrogate U+%04X in JSON string", code)
			case code >= 0xD800 && code <= 0xDBFF:
				// A high surrogate must be followed immediately by its low half.
				if i+11 >= len(raw) || raw[i+6] != '\\' || raw[i+7] != 'u' {
					return fmt.Errorf("wire: unpaired high surrogate U+%04X in JSON string", code)
				}
				low, ok := hex4(raw, i+8)
				if !ok || low < 0xDC00 || low > 0xDFFF {
					return fmt.Errorf("wire: high surrogate U+%04X is not followed by a low surrogate", code)
				}
				i += 11
				continue
			}
			i += 5
		}
	}
	return nil
}

// hex4 reads the four hex digits starting at offset, reporting whether all four were hex.
func hex4(raw []byte, offset int) (uint32, bool) {
	if offset+4 > len(raw) {
		return 0, false
	}
	var value uint32
	for _, c := range raw[offset : offset+4] {
		value <<= 4
		switch {
		case c >= '0' && c <= '9':
			value |= uint32(c - '0')
		case c >= 'a' && c <= 'f':
			value |= uint32(c-'a') + 10
		case c >= 'A' && c <= 'F':
			value |= uint32(c-'A') + 10
		default:
			return 0, false
		}
	}
	return value, true
}

// canonicalNumber is the only number spelling this protocol will hash. See rule 4 above.
//
// Negative zero is excluded even though it is an integer: no implementation can be relied on to
// round-trip "-0" as text rather than as the number 0, and nothing in these payloads needs it.
var canonicalNumber = regexp.MustCompile(`^-?[1-9][0-9]*$|^-?[1-9][0-9]*$|^0$`)

func writeCanonical(buf *bytes.Buffer, value any) error {
	switch typed := value.(type) {
	case nil:
		buf.WriteString("null")
	case bool:
		if typed {
			buf.WriteString("true")
		} else {
			buf.WriteString("false")
		}
	case json.Number:
		literal := typed.String()
		if !canonicalNumber.MatchString(literal) || literal == "-0" {
			return fmt.Errorf("wire: number %q is not a canonical integer", literal)
		}
		buf.WriteString(literal)
	case string:
		escapeCanonicalString(buf, typed)
	case []any:
		buf.WriteByte('[')
		for index, item := range typed {
			if index > 0 {
				buf.WriteByte(',')
			}
			if err := writeCanonical(buf, item); err != nil {
				return err
			}
		}
		buf.WriteByte(']')
	case map[string]any:
		keys := make([]string, 0, len(typed))
		for key := range typed {
			keys = append(keys, key)
		}
		sort.Strings(keys)
		buf.WriteByte('{')
		for index, key := range keys {
			if index > 0 {
				buf.WriteByte(',')
			}
			escapeCanonicalString(buf, key)
			buf.WriteByte(':')
			if err := writeCanonical(buf, typed[key]); err != nil {
				return err
			}
		}
		buf.WriteByte('}')
	default:
		return fmt.Errorf("wire: cannot canonicalise %T", value)
	}
	return nil
}