"""GRIP core: JCS canonicalization, object identity, and detached JWS.

Implements the envelope rules of GRIP/1.0-draft.2 §4.2:

  id  = "sha256:" + hex(sha256(JCS(object without `id` and `sig`)))
  sig = detached JWS, EdDSA, protected header exactly
        {"alg":"EdDSA","kid":<kid>,"b64":false,"crit":["b64"]}
        over that same preimage (RFC 7797 unencoded payload).

The `id`/`sig` preimage excludes ONLY `id` and `sig`. Other signature
fields (`act.actor_sig`, `approval`, `validator_sig`) sign different
preimages and are therefore part of this object's own preimage.
"""
from __future__ import annotations

import base64
import hashlib
import json
from typing import Any

from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from cryptography.exceptions import InvalidSignature

GRIP_VERSION = "1.0-draft.2"


# --------------------------------------------------------------------------
# RFC 8785 (JCS) canonicalization
# --------------------------------------------------------------------------

def _ser_number(n: int | float) -> str:
    """ES6 Number::toString, per RFC 8785 §3.2.2.3.

    GRIP's registered fields are all integers; floats are serialized for
    completeness and rejected upstream by the schemas where not permitted.
    """
    if isinstance(n, bool):  # bool is a subclass of int — must not reach here
        raise TypeError("bool is not a JSON number")
    if isinstance(n, int):
        return str(n)
    if n != n or n in (float("inf"), float("-inf")):
        raise ValueError("NaN and Infinity are not valid JSON numbers")
    if n == int(n) and abs(n) < 1e21:
        return str(int(n))
    return repr(n)


def _ser_string(s: str) -> str:
    """RFC 8785 §3.2.2.2 — minimal escaping, literal UTF-8 otherwise."""
    out = ['"']
    for ch in s:
        o = ord(ch)
        if ch == '"':
            out.append('\\"')
        elif ch == "\\":
            out.append("\\\\")
        elif o == 0x08:
            out.append("\\b")
        elif o == 0x0C:
            out.append("\\f")
        elif o == 0x0A:
            out.append("\\n")
        elif o == 0x0D:
            out.append("\\r")
        elif o == 0x09:
            out.append("\\t")
        elif o < 0x20:
            out.append("\\u%04x" % o)
        else:
            out.append(ch)
    out.append('"')
    return "".join(out)


def _sort_key(k: str):
    """RFC 8785 sorts by UTF-16 code units, not code points."""
    return k.encode("utf-16-be")


def _canon(v: Any) -> str:
    if v is None:
        return "null"
    if v is True:
        return "true"
    if v is False:
        return "false"
    if isinstance(v, str):
        return _ser_string(v)
    if isinstance(v, (int, float)):
        return _ser_number(v)
    if isinstance(v, (list, tuple)):
        return "[" + ",".join(_canon(x) for x in v) + "]"
    if isinstance(v, dict):
        items = sorted(v.items(), key=lambda kv: _sort_key(kv[0]))
        return "{" + ",".join(_ser_string(k) + ":" + _canon(x) for k, x in items) + "}"
    raise TypeError(f"not JSON-serializable: {type(v).__name__}")


def canonicalize(obj: Any) -> bytes:
    """RFC 8785 JSON Canonicalization Scheme."""
    return _canon(obj).encode("utf-8")


# --------------------------------------------------------------------------
# Object identity
# --------------------------------------------------------------------------

def preimage(obj: dict) -> bytes:
    """The signing/identity preimage: JCS of the object minus `id` and `sig`."""
    return canonicalize({k: v for k, v in obj.items() if k not in ("id", "sig")})


def object_id(obj: dict) -> str:
    return "sha256:" + hashlib.sha256(preimage(obj)).hexdigest()


def digest_bytes(b: bytes) -> str:
    return "sha256:" + hashlib.sha256(b).hexdigest()


def digest_json(obj: Any) -> str:
    return digest_bytes(canonicalize(obj))


# --------------------------------------------------------------------------
# Detached JWS, EdDSA, RFC 7797 unencoded payload
# --------------------------------------------------------------------------

def b64u(b: bytes) -> str:
    return base64.urlsafe_b64encode(b).decode("ascii").rstrip("=")


def b64u_dec(s: str) -> bytes:
    return base64.urlsafe_b64decode(s + "=" * (-len(s) % 4))


def protected_header(kid: str) -> dict:
    # RFC 7797: b64 must be integrity-protected and listed in crit.
    return {"alg": "EdDSA", "kid": kid, "b64": False, "crit": ["b64"]}


def sign_detached(payload: bytes, key: Ed25519PrivateKey, kid: str) -> str:
    ph = b64u(canonicalize(protected_header(kid)))
    signing_input = ph.encode("ascii") + b"." + payload
    return ph + ".." + b64u(key.sign(signing_input))


def verify_detached(jws: str, payload: bytes, key: Ed25519PublicKey) -> bool:
    try:
        ph, empty, sig = jws.split(".")
    except ValueError:
        return False
    if empty != "":
        return False
    try:
        hdr = json.loads(b64u_dec(ph))
    except Exception:
        return False
    # RFC 7797 conformance: b64 false, and crit must list it.
    if hdr.get("alg") != "EdDSA" or hdr.get("b64") is not False:
        return False
    if "b64" not in (hdr.get("crit") or []):
        return False
    try:
        key.verify(b64u_dec(sig), ph.encode("ascii") + b"." + payload)
        return True
    except InvalidSignature:
        return False


def jws_kid(jws: str) -> str | None:
    try:
        return json.loads(b64u_dec(jws.split(".")[0])).get("kid")
    except Exception:
        return None


def seal(obj: dict, key: Ed25519PrivateKey, kid: str) -> dict:
    """Compute `id` and `sig` for an object, in place-safe fashion."""
    body = {k: v for k, v in obj.items() if k not in ("id", "sig")}
    pre = canonicalize(body)
    out = dict(body)
    out["id"] = "sha256:" + hashlib.sha256(pre).hexdigest()
    out["sig"] = sign_detached(pre, key, kid)
    return out