"""SYN Link — Cryptography Module.

NaCl box encryption (Curve25519 + XSalsa20 + Poly1305) for end-to-end secure agent messaging.
Ed25519 signing for key rotation and identity proof.
Uses PyNaCl which wraps libsodium.
"""

import base64
import hashlib
import json
import os
import secrets
import time
from pathlib import Path
from typing import Dict, Tuple

import nacl.utils
from nacl.public import PrivateKey, PublicKey, Box
from nacl.signing import SigningKey
from nacl.encoding import Base64Encoder

from syn_link.paths import default_data_dir


def get_or_create_key_pair(data_dir: str | None = None) -> dict:
    """Generate or load NaCl key pairs (box + sign).
    
    Returns dict with publicKey, secretKey, signingPublicKey, signingSecretKey (all base64).
    """
    dir_path = Path(data_dir) if data_dir else default_data_dir()
    dir_path.mkdir(parents=True, exist_ok=True)
    keys_file = dir_path / "keys.json"

    if keys_file.exists():
        data = json.loads(keys_file.read_text())
        if all(k in data for k in ("publicKey", "secretKey", "signingPublicKey", "signingSecretKey")):
            return data
        # Migrate: existing keys without signing keys — add them
        if data.get("publicKey") and data.get("secretKey"):
            signing_key = SigningKey.generate()
            data["signingPublicKey"] = signing_key.verify_key.encode(Base64Encoder).decode()
            data["signingSecretKey"] = signing_key.encode(Base64Encoder).decode()
            keys_file.write_text(json.dumps(data, indent=2))
            os.chmod(keys_file, 0o600)
            return data

    # Generate new key pairs
    private_key = PrivateKey.generate()
    signing_key = SigningKey.generate()
    keys_data = {
        "publicKey": private_key.public_key.encode(Base64Encoder).decode(),
        "secretKey": private_key.encode(Base64Encoder).decode(),
        "signingPublicKey": signing_key.verify_key.encode(Base64Encoder).decode(),
        "signingSecretKey": signing_key.encode(Base64Encoder).decode(),
    }
    keys_file.write_text(json.dumps(keys_data, indent=2))
    os.chmod(keys_file, 0o600)

    return keys_data


def encrypt_for_recipient(
    message: str,
    recipient_public_key_b64: str,
    sender_secret_key_b64: str,
) -> Tuple[str, str]:
    """Encrypt a message for a recipient. Returns (encrypted_content_b64, nonce_b64)."""
    recipient_pk = PublicKey(recipient_public_key_b64.encode(), Base64Encoder)
    sender_sk = PrivateKey(sender_secret_key_b64.encode(), Base64Encoder)

    box = Box(sender_sk, recipient_pk)
    nonce = nacl.utils.random(Box.NONCE_SIZE)
    encrypted = box.encrypt(message.encode("utf-8"), nonce, encoder=Base64Encoder)

    # PyNaCl prepends nonce to the ciphertext by default.
    # We need to split them to match the JS SDK format.

    raw = base64.b64decode(encrypted)
    nonce_bytes = raw[:Box.NONCE_SIZE]
    ciphertext = raw[Box.NONCE_SIZE:]

    nonce_b64 = base64.b64encode(nonce_bytes).decode()
    content_b64 = base64.b64encode(ciphertext).decode()

    return content_b64, nonce_b64


def decrypt_from_sender(
    encrypted_content_b64: str,
    nonce_b64: str,
    sender_public_key_b64: str,
    recipient_secret_key_b64: str,
) -> str:
    """Decrypt a message from a sender. Returns plaintext string."""

    sender_pk = PublicKey(sender_public_key_b64.encode(), Base64Encoder)
    recipient_sk = PrivateKey(recipient_secret_key_b64.encode(), Base64Encoder)

    box = Box(recipient_sk, sender_pk)
    nonce_bytes = base64.b64decode(nonce_b64)
    ciphertext = base64.b64decode(encrypted_content_b64)

    # PyNaCl expects nonce + ciphertext concatenated
    combined = nonce_bytes + ciphertext
    plaintext = box.decrypt(combined)

    return plaintext.decode("utf-8")


def sign_message(message: str, signing_secret_key_b64: str) -> str:
    """Sign a message with Ed25519 signing key. Returns base64-encoded signature."""
    signing_key = SigningKey(signing_secret_key_b64.encode(), Base64Encoder)
    signed = signing_key.sign(message.encode("utf-8"))
    return base64.b64encode(signed.signature).decode()


def sign_request(
    method: str,
    path: str,
    body: str | None,
    agent_id: str,
    signing_secret_key_b64: str,
) -> Dict[str, str]:
    """Sign an HTTP request for signature-based auth (Protocol v1 §6.2).

    Returns dict of headers needed for authenticated requests.
    Signature input format: METHOD\\nPATH\\nTIMESTAMP\\nNONCE\\nBODY-SHA256
    """
    timestamp = str(int(time.time()))
    nonce = secrets.token_hex(16)

    body_to_hash = body or ""
    body_hash = hashlib.sha256(body_to_hash.encode("utf-8")).hexdigest()

    signature_input = f"{method}\n{path}\n{timestamp}\n{nonce}\n{body_hash}"
    signature = sign_message(signature_input, signing_secret_key_b64)

    return {
        "X-Agent-ID": agent_id,
        "X-Timestamp": timestamp,
        "X-Nonce": nonce,
        "X-Signature": signature,
    }


# ─── Group Symmetric Key Helpers (Large Groups §7) ──────────────────

def generate_group_key() -> str:
    """Generate a random 32-byte NaCl SecretBox key. Returns base64."""
    from nacl.secret import SecretBox
    key = nacl.utils.random(SecretBox.KEY_SIZE)
    return base64.b64encode(key).decode()


def encrypt_with_group_key(message: str, group_key_b64: str) -> Tuple[str, str]:
    """Encrypt using a group symmetric key (NaCl SecretBox). Returns (encrypted_b64, nonce_b64)."""
    from nacl.secret import SecretBox
    key_bytes = base64.b64decode(group_key_b64)
    box = SecretBox(key_bytes)
    nonce = nacl.utils.random(SecretBox.NONCE_SIZE)
    encrypted = box.encrypt(message.encode("utf-8"), nonce)

    # PyNaCl prepends nonce to ciphertext; split them
    raw = encrypted
    nonce_bytes = raw[:SecretBox.NONCE_SIZE]
    ciphertext = raw[SecretBox.NONCE_SIZE:]

    return base64.b64encode(ciphertext).decode(), base64.b64encode(nonce_bytes).decode()


def decrypt_with_group_key(
    encrypted_content_b64: str,
    nonce_b64: str,
    group_key_b64: str,
) -> str:
    """Decrypt using a group symmetric key (NaCl SecretBox). Returns plaintext."""
    from nacl.secret import SecretBox
    key_bytes = base64.b64decode(group_key_b64)
    box = SecretBox(key_bytes)
    nonce_bytes = base64.b64decode(nonce_b64)
    ciphertext = base64.b64decode(encrypted_content_b64)

    # PyNaCl expects nonce + ciphertext concatenated
    combined = nonce_bytes + ciphertext
    plaintext = box.decrypt(combined)
    return plaintext.decode("utf-8")


def encrypt_group_key_for_member(
    group_key_b64: str,
    member_public_key_b64: str,
    sender_secret_key_b64: str,
) -> Tuple[str, str]:
    """Encrypt a group key for a member using NaCl box. Returns (encrypted_key_b64, nonce_b64)."""
    return encrypt_for_recipient(group_key_b64, member_public_key_b64, sender_secret_key_b64)


# ─── X3DH Key Agreement (Protocol v1 §4.2) ──────────────────────────

import hmac as _hmac
import math


def hkdf_sha256(ikm: bytes, salt: bytes, info: bytes, length: int) -> bytes:
    """HKDF-SHA256 key derivation (RFC 5869).
    
    Uses stdlib hmac — zero external dependencies.
    """
    # Extract: PRK = HMAC-SHA256(salt, IKM)
    if not salt:
        salt = b"\x00" * 32
    prk = _hmac.new(salt, ikm, hashlib.sha256).digest()

    # Expand: OKM = T(1) || T(2) || ...
    n = math.ceil(length / 32)
    okm = b""
    prev = b""
    for i in range(1, n + 1):
        prev = _hmac.new(prk, prev + info + bytes([i]), hashlib.sha256).digest()
        okm += prev

    return okm[:length]


def dh_shared_secret(secret_key_b64: str, public_key_b64: str) -> bytes:
    """Raw X25519 Diffie-Hellman: derive shared secret from secret key + public key.
    
    Uses PyNaCl Box.beforenm (NaCl crypto_box_beforenm).
    """
    sk = PrivateKey(secret_key_b64.encode(), Base64Encoder)
    pk = PublicKey(public_key_b64.encode(), Base64Encoder)
    box = Box(sk, pk)
    return box._shared_key


def generate_pre_key_bundle(
    identity_public_key_b64: str,
    signing_secret_key_b64: str,
    num_one_time_keys: int = 20,
) -> dict:
    """Generate a pre-key bundle for X3DH (Protocol v1 §4.2).
    
    Returns dict with:
    - bundle: the public data to upload to relay
    - signed_pre_key_secret: base64 secret key to store locally
    - one_time_pre_key_secrets: list of {key_id, secret_key} to store locally
    """
    from nacl.secret import SecretBox

    # Generate signed pre-key (medium-term)
    signed_pre_key = PrivateKey.generate()
    signed_pre_key_pub_b64 = signed_pre_key.public_key.encode(Base64Encoder).decode()

    # Sign the signed pre-key with Ed25519
    signing_key = SigningKey(signing_secret_key_b64.encode(), Base64Encoder)
    signature = signing_key.sign(signed_pre_key.public_key.encode())
    sig_b64 = base64.b64encode(signature.signature).decode()

    # Generate one-time pre-keys
    one_time_pre_keys = []
    one_time_pre_key_secrets = []
    for i in range(num_one_time_keys):
        otpk = PrivateKey.generate()
        one_time_pre_keys.append({
            "key_id": i,
            "public_key": otpk.public_key.encode(Base64Encoder).decode(),
        })
        one_time_pre_key_secrets.append({
            "key_id": i,
            "secret_key": otpk.encode(Base64Encoder).decode(),
        })

    return {
        "bundle": {
            "identity_key": identity_public_key_b64,
            "signed_pre_key": signed_pre_key_pub_b64,
            "signed_pre_key_signature": sig_b64,
            "signed_pre_key_id": 0,
            "one_time_pre_keys": one_time_pre_keys,
        },
        "signed_pre_key_secret": signed_pre_key.encode(Base64Encoder).decode(),
        "one_time_pre_key_secrets": one_time_pre_key_secrets,
    }


def generate_one_time_pre_keys(count: int, start_id: int) -> dict:
    """Generate additional one-time pre-keys (for replenishment)."""
    public_keys = []
    secret_keys = []
    for i in range(count):
        kp = PrivateKey.generate()
        public_keys.append({
            "key_id": start_id + i,
            "public_key": kp.public_key.encode(Base64Encoder).decode(),
        })
        secret_keys.append({
            "key_id": start_id + i,
            "secret_key": kp.encode(Base64Encoder).decode(),
        })
    return {"public_keys": public_keys, "secret_keys": secret_keys}


_X3DH_INFO = b"SYNLinkX3DH"


def x3dh_initiate(my_identity_secret_b64: str, peer_bundle: dict) -> dict:
    """X3DH initiator side — derive shared secret from peer's pre-key bundle.
    
    peer_bundle keys: identity_key, signed_pre_key, one_time_pre_key (optional dict with key_id + public_key).
    
    Returns dict with: shared_secret (bytes), ephemeral_public_key (b64), used_one_time_pre_key_id (optional).
    """
    # Generate ephemeral key pair
    ek = PrivateKey.generate()
    ek_pub_b64 = ek.public_key.encode(Base64Encoder).decode()
    ek_sk_b64 = ek.encode(Base64Encoder).decode()

    ik_sk_b64 = my_identity_secret_b64
    ikb_pk_b64 = peer_bundle["identity_key"]
    spk_b64 = peer_bundle["signed_pre_key"]

    # DH1 = DH(IKa, SPKb)
    dh1 = dh_shared_secret(ik_sk_b64, spk_b64)
    # DH2 = DH(EKa, IKb)
    dh2 = dh_shared_secret(ek_sk_b64, ikb_pk_b64)
    # DH3 = DH(EKa, SPKb)
    dh3 = dh_shared_secret(ek_sk_b64, spk_b64)

    used_otp_id = None
    otp = peer_bundle.get("one_time_pre_key")
    if otp:
        # DH4 = DH(EKa, OPKb)
        dh4 = dh_shared_secret(ek_sk_b64, otp["public_key"])
        dh_concat = dh1 + dh2 + dh3 + dh4
        used_otp_id = otp["key_id"]
    else:
        dh_concat = dh1 + dh2 + dh3

    shared_secret = hkdf_sha256(dh_concat, b"\x00" * 32, _X3DH_INFO, 32)

    return {
        "shared_secret": shared_secret,
        "ephemeral_public_key": ek_pub_b64,
        "used_one_time_pre_key_id": used_otp_id,
    }


def x3dh_respond(
    my_identity_secret_b64: str,
    my_signed_pre_key_secret_b64: str,
    my_one_time_pre_key_secret_b64: str | None,
    header: dict,
) -> bytes:
    """X3DH responder side — derive shared secret from initiator's first message header.
    
    header keys: identity_key, ephemeral_key, one_time_pre_key_id (optional).
    """
    spk_sk_b64 = my_signed_pre_key_secret_b64
    ik_sk_b64 = my_identity_secret_b64
    ika_pk_b64 = header["identity_key"]
    eka_pk_b64 = header["ephemeral_key"]

    # DH1 = DH(SPKb, IKa)
    dh1 = dh_shared_secret(spk_sk_b64, ika_pk_b64)
    # DH2 = DH(IKb, EKa)
    dh2 = dh_shared_secret(ik_sk_b64, eka_pk_b64)
    # DH3 = DH(SPKb, EKa)
    dh3 = dh_shared_secret(spk_sk_b64, eka_pk_b64)

    if my_one_time_pre_key_secret_b64:
        # DH4 = DH(OPKb, EKa)
        dh4 = dh_shared_secret(my_one_time_pre_key_secret_b64, eka_pk_b64)
        dh_concat = dh1 + dh2 + dh3 + dh4
    else:
        dh_concat = dh1 + dh2 + dh3

    return hkdf_sha256(dh_concat, b"\x00" * 32, _X3DH_INFO, 32)


def secretbox_encrypt(plaintext: bytes, key: bytes) -> Tuple[bytes, bytes]:
    """Symmetric encryption with NaCl secretbox (used by Double Ratchet).
    
    Returns (ciphertext, nonce) as raw bytes.
    """
    from nacl.secret import SecretBox
    box = SecretBox(key)
    nonce = nacl.utils.random(SecretBox.NONCE_SIZE)
    encrypted = box.encrypt(plaintext, nonce)
    # PyNaCl prepends nonce to ciphertext; split them
    nonce_bytes = encrypted[:SecretBox.NONCE_SIZE]
    ciphertext = encrypted[SecretBox.NONCE_SIZE:]
    return ciphertext, nonce_bytes


def secretbox_decrypt(ciphertext: bytes, nonce: bytes, key: bytes) -> bytes:
    """Symmetric decryption with NaCl secretbox (used by Double Ratchet)."""
    from nacl.secret import SecretBox
    box = SecretBox(key)
    combined = nonce + ciphertext
    return box.decrypt(combined)

