"""
pqc_sign.py — hybrid EdDSA + ML-DSA-65 JWT-like token. MIT.

Install:
    pip install pqcrypto pynacl
"""
from __future__ import annotations
import base64
import json
import time
from typing import Optional

from pqcrypto.sign import ml_dsa_65
from nacl.signing import SigningKey, VerifyKey


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


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


class HybridSigner:
    def __init__(self, ed_seed: bytes, pq_pub: bytes, pq_priv: bytes) -> None:
        self.ed = SigningKey(ed_seed)
        self.pq_pub = pq_pub
        self.pq_priv = pq_priv

    @classmethod
    def fresh(cls) -> "HybridSigner":
        pq_pub, pq_priv = ml_dsa_65.generate_keypair()
        from os import urandom
        return cls(urandom(32), pq_pub, pq_priv)

    @property
    def public_bundle(self) -> dict:
        return {
            "ed": b64(bytes(self.ed.verify_key)),
            "pq": b64(self.pq_pub),
        }

    def issue(self, claims: dict) -> str:
        header = {"alg": "EdDSA+ML-DSA-65", "typ": "JWT-HYBRID"}
        payload = {**claims, "iat": int(time.time())}
        hb = b64(json.dumps(header, separators=(",", ":")).encode())
        pb = b64(json.dumps(payload, separators=(",", ":")).encode())
        msg = f"{hb}.{pb}".encode()
        sed = self.ed.sign(msg).signature
        spq = ml_dsa_65.sign(self.pq_priv, msg)
        return f"{hb}.{pb}.{b64(sed)}.{b64(spq)}"

    def verify(self, token: str) -> Optional[dict]:
        try:
            hb, pb, s_ed, s_pq = token.split(".")
            msg = f"{hb}.{pb}".encode()
            VerifyKey(bytes(self.ed.verify_key)).verify(msg, ub64(s_ed))
            ml_dsa_65.verify(self.pq_pub, msg, ub64(s_pq))
            return json.loads(ub64(pb))
        except Exception:
            return None


if __name__ == "__main__":
    s = HybridSigner.fresh()
    tok = s.issue({"sub": "alice", "scope": "read:all"})
    print("token bytes:", len(tok))
    print("verifies   :", s.verify(tok))