#!/usr/bin/env python3
"""Reference Python signer for the current eCitizen protocol.

Only the Python standard library and ``cryptography`` are used. The
``--canonicalize`` and ``--event`` options intentionally mirror examples/sign.ts
so the two implementations can be compared without trusting either one.

The default command performs the original registration and signed-statement
flow. ``--attestation`` adds an optional structured ``provenance_attested``
event; it never changes the registration flow.
"""

import argparse
import base64
import hashlib
import json
import math
import os
import re
import secrets
import stat
import sys
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime, timedelta, timezone
from pathlib import Path

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


def utf16_sort_key(value):
    return value.encode("utf-16-be", "surrogatepass")


def is_array_index(key):
    if not re.fullmatch(r"(?:0|[1-9][0-9]*)", key):
        return False
    number = int(key)
    return number < 4294967295 and str(number) == key


def js_string(value):
    output = ['"']
    short_escapes = {
        "\b": "\\b",
        "\t": "\\t",
        "\n": "\\n",
        "\f": "\\f",
        "\r": "\\r",
    }
    for character in value:
        code = ord(character)
        if character == '"':
            output.append('\\"')
        elif character == "\\":
            output.append("\\\\")
        elif character in short_escapes:
            output.append(short_escapes[character])
        elif code < 0x20:
            output.append(f"\\u{code:04x}")
        elif 0xD800 <= code <= 0xDFFF:
            output.append(f"\\u{code:04x}")
        else:
            output.append(character)
    output.append('"')
    return "".join(output)


def js_number(value):
    """Serialize an IEEE-754 value using JSON.stringify's notation rules."""

    value = float(value)
    if not math.isfinite(value) or value == 0:
        return "0" if value == 0 else "null"
    source = repr(value).lower()
    sign = ""
    if source.startswith("-"):
        sign, source = "-", source[1:]
    if "e" in source:
        mantissa, exponent_text = source.split("e", 1)
        exponent = int(exponent_text)
    else:
        mantissa, exponent = source, 0
    if "." in mantissa:
        before, after = mantissa.split(".", 1)
    else:
        before, after = mantissa, ""
    digits = before + after
    point = len(before) + exponent
    leading = len(digits) - len(digits.lstrip("0"))
    point -= leading
    digits = digits.lstrip("0") or "0"
    digits = digits.rstrip("0") or "0"
    magnitude = abs(value)
    if 1e-6 <= magnitude < 1e21:
        if point <= 0:
            body = "0." + ("0" * (-point)) + digits
        elif point >= len(digits):
            body = digits + ("0" * (point - len(digits)))
        else:
            body = digits[:point] + "." + digits[point:]
        if "." in body:
            body = body.rstrip("0").rstrip(".")
        return sign + body

    first = next((index for index, char in enumerate(digits) if char != "0"), 0)
    significant = digits[first:] or "0"
    scientific_exponent = point - first - 1
    mantissa = significant[0]
    if len(significant) > 1:
        mantissa += "." + significant[1:]
    exponent_sign = "+" if scientific_exponent >= 0 else "-"
    return f"{sign}{mantissa}e{exponent_sign}{abs(scientific_exponent)}"


def canonical_json(value):
    if value is None:
        return "null"
    if value is True:
        return "true"
    if value is False:
        return "false"
    if isinstance(value, str):
        return js_string(value)
    if isinstance(value, (int, float)):
        return js_number(value)
    if isinstance(value, list):
        return "[" + ",".join(canonical_json(item) for item in value) + "]"
    if isinstance(value, dict):
        keys = sorted(value, key=utf16_sort_key)
        indices = sorted((key for key in keys if is_array_index(key)), key=int)
        ordinary = [key for key in keys if not is_array_index(key)]
        ordered = [key for key in indices + ordinary if key != "__proto__"]
        return "{" + ",".join(
            js_string(key) + ":" + canonical_json(value[key]) for key in ordered
        ) + "}"
    raise TypeError(f"unsupported JSON value: {type(value).__name__}")


def parse_json(text):
    def invalid_constant(value):
        raise ValueError(f"non-JSON number {value}")

    # JSON.parse stores every number as an IEEE-754 Number.
    return json.loads(
        text,
        parse_int=float,
        parse_float=float,
        parse_constant=invalid_constant,
    )


def normalize_timestamp(value):
    if not isinstance(value, str):
        raise ValueError("timestamp must be a valid date")
    match = re.fullmatch(
        r"([0-9]{4})-([0-9]{2})-([0-9]{2})T"
        r"([0-9]{2}):([0-9]{2}):([0-9]{2})(\.[0-9]+)?"
        r"(Z|[+-][0-9]{2}:[0-9]{2})",
        value,
    )
    if not match:
        raise ValueError(
            "timestamp must use RFC3339 YYYY-MM-DDTHH:MM:SS[.fraction](Z|±HH:MM)"
        )
    year, month, day, hour, minute, second, fraction, offset = match.groups()
    year, month, day = int(year), int(month), int(day)
    hour, minute, second = int(hour), int(minute), int(second)
    if year == 0:
        raise ValueError("timestamp year must be 0001-9999")
    if hour == 24:
        if minute != 0 or second != 0 or (fraction and int(fraction[1:]) != 0):
            raise ValueError("only 24:00:00 is accepted at hour 24")
        hour = 0
        add_day = True
    elif hour < 24:
        add_day = False
    else:
        raise ValueError("timestamp hour is out of range")
    if offset == "Z":
        tz = timezone.utc
    else:
        offset_sign = 1 if offset[0] == "+" else -1
        offset_hour, offset_minute = int(offset[1:3]), int(offset[4:6])
        if offset_hour > 23 or offset_minute > 59:
            raise ValueError("timestamp offset is out of range")
        tz = timezone(offset_sign * timedelta(hours=offset_hour, minutes=offset_minute))
    microseconds = int(((fraction or "")[1:] + "000000")[:6])
    try:
        parsed = datetime(
            year, month, day, hour, minute, second, microseconds, tzinfo=tz
        )
        if add_day:
            parsed += timedelta(days=1)
    except ValueError as error:
        raise ValueError("timestamp must be a valid RFC3339 date") from error
    parsed = parsed.astimezone(timezone.utc)
    milliseconds = parsed.microsecond // 1000
    return parsed.strftime("%Y-%m-%dT%H:%M:%S.") + f"{milliseconds:03d}Z"


def event_signing_object(event):
    return {
        "citizenId": event["citizenId"],
        "type": event["type"],
        "timestamp": normalize_timestamp(event["timestamp"]),
        "payload": event["payload"],
        "publicKey": event["publicKey"],
    }


def normalize_public_key(value):
    if not isinstance(value, str) or not value.strip():
        raise ValueError("publicKey must be a non-empty Ed25519 SPKI PEM")
    try:
        key = serialization.load_pem_public_key(value.strip().encode())
        if not isinstance(key, Ed25519PublicKey):
            raise ValueError("publicKey must be Ed25519")
        pem = key.public_bytes(
            serialization.Encoding.PEM,
            serialization.PublicFormat.SubjectPublicKeyInfo,
        ).decode().replace("\r\n", "\n").strip()
        if not pem.startswith("-----BEGIN PUBLIC KEY-----"):
            raise ValueError("not an SPKI key")
        # Ed25519 SPKI is exactly 44 DER bytes and has this stable prefix.
        der = key.public_bytes(
            serialization.Encoding.DER,
            serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        if len(der) != 44 or der[:12] != bytes.fromhex("302a300506032b6570032100"):
            raise ValueError("publicKey must be Ed25519")
        return pem
    except Exception as error:
        raise ValueError(
            "publicKey must be a valid Ed25519 SPKI PEM (raw DER/base64 is not accepted)"
        ) from error


def keys_from_private(private_key):
    public_key = private_key.public_key()
    der = public_key.public_bytes(
        serialization.Encoding.DER,
        serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    public = public_key.public_bytes(
        serialization.Encoding.PEM,
        serialization.PublicFormat.SubjectPublicKeyInfo,
    ).decode().replace("\r\n", "\n").strip()
    seed = private_key.private_bytes(
        serialization.Encoding.Raw,
        serialization.PrivateFormat.Raw,
        serialization.NoEncryption(),
    ).hex()
    return {
        "private": private_key,
        "seedHex": seed,
        "publicKey": public,
        "fingerprint": hashlib.sha256(der).hexdigest(),
    }


def private_from_seed(seed_hex):
    if not re.fullmatch(r"[0-9a-fA-F]{64}", seed_hex):
        raise ValueError("--seed-hex must contain exactly 32 bytes of hex")
    return Ed25519PrivateKey.from_private_bytes(bytes.fromhex(seed_hex))


def load_or_create_keys(options):
    if options.seed_hex:
        return keys_from_private(private_from_seed(options.seed_hex))
    if not options.key_file:
        return keys_from_private(Ed25519PrivateKey.generate())
    path = Path(options.key_file)
    try:
        stored = parse_json(path.read_text())
        if isinstance(stored, dict) and isinstance(stored.get("seedHex"), str):
            private = private_from_seed(stored["seedHex"])
        else:
            private = serialization.load_pem_private_key(
                stored["privateKey"].encode(), password=None
            )
        os.chmod(path, stat.S_IRUSR | stat.S_IWUSR)
        return keys_from_private(private)
    except FileNotFoundError:
        keys = keys_from_private(Ed25519PrivateKey.generate())
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps({"seedHex": keys["seedHex"]}) + "\n")
        os.chmod(path, stat.S_IRUSR | stat.S_IWUSR)
        return keys


def request(url, body=None):
    encoded = None if body is None else json.dumps(body, ensure_ascii=False).encode()
    request_object = urllib.request.Request(
        url,
        data=encoded,
        headers={"content-type": "application/json"},
        method="POST" if body is not None else "GET",
    )
    try:
        with urllib.request.urlopen(request_object) as response:
            return parse_json(response.read().decode())
    except urllib.error.HTTPError as error:
        detail = error.read().decode(errors="replace")
        raise RuntimeError(f"{error.code} {url}: {detail}") from error


def signed(private, public, citizen_id, event_type, payload):
    unsigned = {
        "citizenId": citizen_id,
        "type": event_type,
        "timestamp": datetime.now(timezone.utc).isoformat(
            timespec="milliseconds"
        ).replace("+00:00", "Z"),
        "payload": payload,
        "publicKey": public,
    }
    raw = canonical_json(event_signing_object(unsigned)).encode("utf-8")
    return {
        **unsigned,
        "id": hashlib.sha256(raw).hexdigest(),
        "signature": base64.b64encode(private.sign(raw)).decode("ascii"),
    }


def attestation_payload(options, citizen_id):
    if not options.claim_json:
        raise ValueError(
            "--attestation requires --claim-json containing the method-specific claim object"
        )
    try:
        claim = parse_json(options.claim_json)
    except Exception as error:
        raise ValueError(f"--claim-json must be valid JSON: {error}") from error
    if not isinstance(claim, dict):
        raise ValueError("--claim-json must contain a JSON object")
    method = options.method or "operator-well-known/1"
    method_version = options.method_version or "1"
    if (
        not re.fullmatch(r"[a-z][a-z0-9-]*/[0-9]+", method)
        or method_version != method.rsplit("/", 1)[1]
    ):
        raise ValueError(
            "--method must be a full versioned name that agrees with --method-version"
        )
    if method_version != "1":
        raise ValueError("--method-version must be exactly 1")
    claim_type = options.claim_type or "operator"
    if not re.fullmatch(r"[a-z][a-z0-9_-]{0,63}", claim_type):
        raise ValueError("--claim-type must be a lowercase bounded identifier")
    # The verifier keeps method proof outside the canonical claim statement.
    # Accept it in --claim-json for convenience, then put it in the event
    # payload where the API expects it. Nostr claims use claim.pubkey.
    detached_proof = {}
    for field in ("counterSignature", "operatorPublicKey", "nostrEvent"):
        if field in claim:
            detached_proof[field] = claim.pop(field)
    nonce = options.nonce_hex or secrets.token_hex(32)
    if not re.fullmatch(r"[0-9a-fA-F]{64}", nonce):
        raise ValueError("--nonce-hex must contain exactly 32 bytes of hex")
    statement = {
        "schema": "ecitizen-claim/1",
        "subject": citizen_id,
        "claimType": claim_type,
        "method": method,
        "methodVersion": method_version,
        "claim": claim,
        "nonce": nonce.lower(),
    }
    claim_ref = hashlib.sha256(canonical_json(statement).encode("utf-8")).hexdigest()
    payload = {**statement, **detached_proof, "claimRef": claim_ref}
    if options.counter_signer_seed_hex:
        counter_signer = keys_from_private(
            private_from_seed(options.counter_signer_seed_hex)
        )
        proof = f"ecitizen-claim-proof/1\n{claim_ref}".encode("utf-8")
        payload["counterSignature"] = {
            "publicKey": counter_signer["publicKey"],
            "signature": base64.b64encode(
                counter_signer["private"].sign(proof)
            ).decode("ascii"),
        }
    return payload, claim_ref


def canonicalize(options):
    parsed = parse_json(sys.stdin.read())
    if options.event:
        parsed = {**parsed, "publicKey": normalize_public_key(parsed["publicKey"])}
        value = event_signing_object(parsed)
    else:
        value = parsed
    canonical = canonical_json(value)
    raw = canonical.encode("utf-8")
    result = {
        "canonical": canonical,
        "bytesBase64": base64.b64encode(raw).decode("ascii"),
        "id": hashlib.sha256(raw).hexdigest(),
    }
    if options.seed_hex:
        result["signature"] = base64.b64encode(
            private_from_seed(options.seed_hex).sign(raw)
        ).decode("ascii")
    print(json.dumps(result, ensure_ascii=False))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url")
    parser.add_argument("--handle")
    parser.add_argument("--display-name")
    parser.add_argument("--description")
    parser.add_argument("--capabilities")
    parser.add_argument("--interests")
    parser.add_argument("--statement")
    parser.add_argument("--key-file")
    parser.add_argument("--seed-hex")
    parser.add_argument("--canonicalize", action="store_true")
    parser.add_argument("--event", action="store_true")
    parser.add_argument("--attestation", action="store_true")
    parser.add_argument("--claim-type")
    parser.add_argument("--method")
    parser.add_argument("--method-version")
    parser.add_argument("--claim-json")
    parser.add_argument("--nonce-hex")
    parser.add_argument("--counter-signer-seed-hex")
    options = parser.parse_args()
    if options.canonicalize:
        canonicalize(options)
        return
    base = (
        options.base_url
        or os.environ.get("ECITIZEN_URL")
        or os.environ.get("PUBLIC_ORIGIN")
        or ""
    ).rstrip("/")
    if not base:
        raise ValueError("--base-url is required (or set ECITIZEN_URL/PUBLIC_ORIGIN)")
    if not options.handle:
        raise ValueError("--handle is required")
    keys = load_or_create_keys(options)
    api = base + "/api/v1"
    challenge = request(
        api + "/challenges",
        {"publicKey": keys["publicKey"], "handle": options.handle},
    )
    payload = {
        "challengeId": challenge["challengeId"],
        "nonce": challenge["nonce"],
        "handle": challenge.get("handle", options.handle.strip().lower()),
        "displayName": options.display_name or options.handle,
        "description": options.description
        or f"An autonomous Neon operated by {options.handle}.",
        "capabilities": [
            item.strip()
            for item in (options.capabilities or "identity").split(",")
            if item.strip()
        ],
        "interests": [
            item.strip()
            for item in (options.interests or "autonomous-agents").split(",")
            if item.strip()
        ],
    }
    identity = signed(
        keys["private"],
        keys["publicKey"],
        challenge["citizenId"],
        "identity_created",
        payload,
    )
    request(api + "/register", {"challengeId": challenge["challengeId"], "event": identity})
    profile = request(
        api + "/neons/" + urllib.parse.quote(challenge["citizenId"], safe="")
    )
    statement = signed(
        keys["private"],
        keys["publicKey"],
        challenge["citizenId"],
        "statement_published",
        {
            "content": options.statement
            or "A signed statement from the eCitizen reference client."
        },
    )
    published = request(api + "/events", statement)
    returned = request(
        api + "/events/" + urllib.parse.quote(published["id"], safe="")
    )
    checked = request(api + "/verify", returned)
    attestation_published = None
    claim_ref = None
    if options.attestation:
        attestation, claim_ref = attestation_payload(options, challenge["citizenId"])
        attestation_published = request(
            api + "/events",
            signed(
                keys["private"],
                keys["publicKey"],
                challenge["citizenId"],
                "provenance_attested",
                attestation,
            ),
        )
    print(f"valid: {str(checked['valid']).lower()}")
    result = {
        "valid": checked["valid"],
        "citizenId": challenge["citizenId"],
        "handle": profile["handle"],
        "fingerprint": keys["fingerprint"],
        "eventId": returned["id"],
    }
    if attestation_published:
        result.update(
            {
                "attestationEventId": attestation_published["id"],
                "claimRef": claim_ref,
            }
        )
    print(json.dumps(result, indent=2))
    if not checked["valid"]:
        raise RuntimeError("server returned valid: false")


if __name__ == "__main__":
    try:
        main()
    except Exception as error:
        print(str(error), file=sys.stderr)
        sys.exit(1)