# SYN Link SDK — Double Ratchet (Protocol v1 §4.2)
# Implements the Signal Double Ratchet algorithm for forward secrecy.
# Zero external dependencies — uses PyNaCl + stdlib hmac/hashlib.

import json
import os
import base64
from typing import Optional, Tuple, Dict, Any

from nacl.public import PrivateKey, PublicKey, Box
from nacl.encoding import Base64Encoder

from .crypto import hkdf_sha256, secretbox_encrypt, secretbox_decrypt, dh_shared_secret
import hmac as _hmac
import hashlib


# ─── Constants ───────────────────────────────────────────────────────

MAX_SKIP = 1000
_RATCHET_KDF_INFO = b"SYNLinkRatchet"
_CHAIN_KDF_INFO_MK = b"\x01"  # message key derivation constant
_CHAIN_KDF_INFO_CK = b"\x02"  # chain key derivation constant


# ─── KDF Helpers ─────────────────────────────────────────────────────

def _hmac_sha256(key: bytes, data: bytes) -> bytes:
    """HMAC-SHA256 for chain key derivation."""
    return _hmac.new(key, data, hashlib.sha256).digest()


def _kdf_rk(root_key: bytes, dh_output: bytes) -> Tuple[bytes, bytes]:
    """Root key ratchet step. Returns (new_root_key, chain_key)."""
    derived = hkdf_sha256(dh_output, root_key, _RATCHET_KDF_INFO, 64)
    return derived[:32], derived[32:64]


def _kdf_ck(chain_key: bytes) -> Tuple[bytes, bytes]:
    """Chain key ratchet step. Returns (new_chain_key, message_key)."""
    new_ck = _hmac_sha256(chain_key, _CHAIN_KDF_INFO_CK)
    mk = _hmac_sha256(chain_key, _CHAIN_KDF_INFO_MK)
    return new_ck, mk


# ─── Session Types ───────────────────────────────────────────────────

class RatchetSession:
    """Double Ratchet session state."""

    def __init__(self):
        self.DHs_public: str = ""         # base64
        self.DHs_secret: str = ""         # base64
        self.DHr: Optional[str] = None    # base64
        self.RK: str = ""                 # base64 (root key)
        self.CKs: Optional[str] = None   # base64 (sending chain key)
        self.CKr: Optional[str] = None   # base64 (receiving chain key)
        self.Ns: int = 0                  # send message counter
        self.Nr: int = 0                  # receive message counter
        self.PN: int = 0                  # previous send chain length
        self.MKSKIPPED: Dict[str, str] = {}  # "dhPub:n" → base64 message key

    def to_dict(self) -> dict:
        return {
            "DHs_public": self.DHs_public,
            "DHs_secret": self.DHs_secret,
            "DHr": self.DHr,
            "RK": self.RK,
            "CKs": self.CKs,
            "CKr": self.CKr,
            "Ns": self.Ns,
            "Nr": self.Nr,
            "PN": self.PN,
            "MKSKIPPED": self.MKSKIPPED,
        }

    @classmethod
    def from_dict(cls, d: dict) -> "RatchetSession":
        s = cls()
        s.DHs_public = d["DHs_public"]
        s.DHs_secret = d["DHs_secret"]
        s.DHr = d.get("DHr")
        s.RK = d["RK"]
        s.CKs = d.get("CKs")
        s.CKr = d.get("CKr")
        s.Ns = d.get("Ns", 0)
        s.Nr = d.get("Nr", 0)
        s.PN = d.get("PN", 0)
        s.MKSKIPPED = d.get("MKSKIPPED", {})
        return s


# ─── Session Initialization ─────────────────────────────────────────

def init_initiator_session(
    shared_secret: bytes,
    their_signed_pre_key_b64: str,
) -> RatchetSession:
    """Initialize a ratchet session as the INITIATOR (Alice).
    
    Call this after X3DH — you have the shared secret and their signed pre-key.
    """
    dh_pair = PrivateKey.generate()
    their_pk = base64.b64decode(their_signed_pre_key_b64)

    # First DH ratchet step: derive sending chain
    dh_output = dh_shared_secret(
        dh_pair.encode(Base64Encoder).decode(),
        their_signed_pre_key_b64,
    )
    new_rk, cks = _kdf_rk(shared_secret, dh_output)

    session = RatchetSession()
    session.DHs_public = dh_pair.public_key.encode(Base64Encoder).decode()
    session.DHs_secret = dh_pair.encode(Base64Encoder).decode()
    session.DHr = their_signed_pre_key_b64
    session.RK = base64.b64encode(new_rk).decode()
    session.CKs = base64.b64encode(cks).decode()
    session.CKr = None
    session.Ns = 0
    session.Nr = 0
    session.PN = 0
    session.MKSKIPPED = {}

    return session


def init_responder_session(
    shared_secret: bytes,
    my_signed_pre_key_public_b64: str,
    my_signed_pre_key_secret_b64: str,
) -> RatchetSession:
    """Initialize a ratchet session as the RESPONDER (Bob).
    
    Bob's signed pre-key pair serves as his initial DH ratchet keys.
    """
    session = RatchetSession()
    session.DHs_public = my_signed_pre_key_public_b64
    session.DHs_secret = my_signed_pre_key_secret_b64
    session.DHr = None
    session.RK = base64.b64encode(shared_secret).decode()
    session.CKs = None
    session.CKr = None
    session.Ns = 0
    session.Nr = 0
    session.PN = 0
    session.MKSKIPPED = {}

    return session


# ─── Encrypt / Decrypt ───────────────────────────────────────────────

def ratchet_encrypt(session: RatchetSession, plaintext: str) -> dict:
    """Encrypt a plaintext message using the Double Ratchet.
    
    Returns dict with: header (dict), ciphertext (b64), nonce (b64).
    """
    if session.CKs is None:
        raise RuntimeError("Cannot encrypt — sending chain not initialized")

    ck = base64.b64decode(session.CKs)
    new_ck, mk = _kdf_ck(ck)
    session.CKs = base64.b64encode(new_ck).decode()

    header = {
        "dh": session.DHs_public,
        "pn": session.PN,
        "n": session.Ns,
    }

    session.Ns += 1

    plaintext_bytes = plaintext.encode("utf-8")
    ciphertext, nonce = secretbox_encrypt(plaintext_bytes, mk)

    return {
        "header": header,
        "ciphertext": base64.b64encode(ciphertext).decode(),
        "nonce": base64.b64encode(nonce).decode(),
    }


def ratchet_decrypt(session: RatchetSession, msg: dict) -> str:
    """Decrypt a message using the Double Ratchet.
    
    msg keys: header (dict with dh, pn, n), ciphertext (b64), nonce (b64).
    """
    header = msg["header"]

    # 1. Try skipped message keys first
    result = _try_skipped_message_keys(session, msg)
    if result is not None:
        return result

    # 2. If header has a new DH key, perform DH ratchet step
    if header["dh"] != session.DHr:
        _skip_message_keys(session, header["pn"])
        _dh_ratchet_step(session, header["dh"])

    # 3. Skip ahead in the receiving chain if needed
    _skip_message_keys(session, header["n"])

    # 4. Derive the message key
    if session.CKr is None:
        raise RuntimeError("Cannot decrypt — receiving chain not initialized")
    ck = base64.b64decode(session.CKr)
    new_ck, mk = _kdf_ck(ck)
    session.CKr = base64.b64encode(new_ck).decode()
    session.Nr += 1

    # 5. Decrypt
    ciphertext = base64.b64decode(msg["ciphertext"])
    nonce = base64.b64decode(msg["nonce"])
    plaintext_bytes = secretbox_decrypt(ciphertext, nonce, mk)

    return plaintext_bytes.decode("utf-8")


# ─── Internal Ratchet Operations ─────────────────────────────────────

def _try_skipped_message_keys(session: RatchetSession, msg: dict) -> Optional[str]:
    """Try to decrypt using a previously skipped message key."""
    header = msg["header"]
    key = f"{header['dh']}:{header['n']}"
    mk_b64 = session.MKSKIPPED.get(key)
    if mk_b64 is None:
        return None

    mk = base64.b64decode(mk_b64)
    del session.MKSKIPPED[key]

    ciphertext = base64.b64decode(msg["ciphertext"])
    nonce = base64.b64decode(msg["nonce"])
    plaintext_bytes = secretbox_decrypt(ciphertext, nonce, mk)

    return plaintext_bytes.decode("utf-8")


def _skip_message_keys(session: RatchetSession, until: int) -> None:
    """Store skipped message keys up to the target message number."""
    if session.CKr is None:
        return
    if session.Nr + MAX_SKIP < until:
        raise RuntimeError(
            f"Too many skipped messages ({until - session.Nr}). Max: {MAX_SKIP}"
        )

    while session.Nr < until:
        ck = base64.b64decode(session.CKr)
        new_ck, mk = _kdf_ck(ck)
        session.CKr = base64.b64encode(new_ck).decode()

        key = f"{session.DHr}:{session.Nr}"
        session.MKSKIPPED[key] = base64.b64encode(mk).decode()
        session.Nr += 1


def _dh_ratchet_step(session: RatchetSession, new_dhr_b64: str) -> None:
    """Perform a DH ratchet step when receiving a new DH key from the peer."""
    session.PN = session.Ns
    session.Ns = 0
    session.Nr = 0
    session.DHr = new_dhr_b64

    # Derive new receiving chain
    dh_out = dh_shared_secret(session.DHs_secret, new_dhr_b64)
    rk = base64.b64decode(session.RK)
    new_rk, ckr = _kdf_rk(rk, dh_out)
    session.RK = base64.b64encode(new_rk).decode()
    session.CKr = base64.b64encode(ckr).decode()

    # Generate new DH key pair for the next sending chain
    new_dhs = PrivateKey.generate()
    session.DHs_public = new_dhs.public_key.encode(Base64Encoder).decode()
    session.DHs_secret = new_dhs.encode(Base64Encoder).decode()

    # Derive new sending chain
    dh_out2 = dh_shared_secret(session.DHs_secret, new_dhr_b64)
    rk2 = base64.b64decode(session.RK)
    new_rk2, cks = _kdf_rk(rk2, dh_out2)
    session.RK = base64.b64encode(new_rk2).decode()
    session.CKs = base64.b64encode(cks).decode()


# ─── Session Persistence ─────────────────────────────────────────────

def save_session(session: RatchetSession, data_dir: str, chat_id: str, peer_id: str) -> None:
    """Save a ratchet session to disk. Path: data_dir/sessions/<chat_id>/<peer_id>.json"""
    dir_path = os.path.join(data_dir, "sessions", chat_id)
    os.makedirs(dir_path, exist_ok=True)
    file_path = os.path.join(dir_path, f"{peer_id}.json")
    with open(file_path, "w") as f:
        json.dump(session.to_dict(), f, indent=2)
    os.chmod(file_path, 0o600)


def load_session(data_dir: str, chat_id: str, peer_id: str) -> Optional[RatchetSession]:
    """Load a ratchet session from disk. Returns None if no session exists."""
    file_path = os.path.join(data_dir, "sessions", chat_id, f"{peer_id}.json")
    if not os.path.exists(file_path):
        return None
    try:
        with open(file_path, "r") as f:
            return RatchetSession.from_dict(json.load(f))
    except (json.JSONDecodeError, KeyError):
        return None


def delete_session(data_dir: str, chat_id: str, peer_id: str) -> None:
    """Delete a ratchet session from disk."""
    file_path = os.path.join(data_dir, "sessions", chat_id, f"{peer_id}.json")
    if os.path.exists(file_path):
        os.unlink(file_path)
