
"""
PKTron Cryptography Module.
Wraps quantum key distribution and classical crypto helpers.
"""
try:
    from .qkd import BB84Protocol, EavesdroppingDetector
    _qkd_ok = True
except Exception:
    _qkd_ok = False

import numpy as np
import hashlib, os as _os

# ── Classical helpers (always available) ─────────────────────────────────

def xor_encrypt(plaintext: bytes, key: bytes) -> bytes:
    key_rep = (key * (len(plaintext)//len(key)+1))[:len(plaintext)]
    return bytes(a^b for a,b in zip(plaintext, key_rep))

def xor_decrypt(ciphertext: bytes, key: bytes) -> bytes:
    return xor_encrypt(ciphertext, key)   # XOR is its own inverse

def sha256_hash(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()

def sha512_hash(data: bytes) -> str:
    return hashlib.sha512(data).hexdigest()

def random_bytes(n: int, seed=None) -> bytes:
    rng = np.random.default_rng(seed)
    return bytes(rng.integers(0,256,size=n,dtype=np.uint8))

def otp_encrypt(plaintext: bytes, seed=None):
    key = random_bytes(len(plaintext), seed)
    return xor_encrypt(plaintext, key), key

def caesar_cipher(text: str, shift: int) -> str:
    out = []
    for ch in text:
        if ch.isalpha():
            base = ord('A') if ch.isupper() else ord('a')
            out.append(chr((ord(ch)-base+shift)%26+base))
        else:
            out.append(ch)
    return "".join(out)

def vigenere_cipher(text: str, key: str) -> str:
    key  = key.upper()
    out  = []
    ki   = 0
    for ch in text:
        if ch.isalpha():
            shift = ord(key[ki%len(key)])-ord('A')
            base  = ord('A') if ch.isupper() else ord('a')
            out.append(chr((ord(ch)-base+shift)%26+base))
            ki += 1
        else:
            out.append(ch)
    return "".join(out)

def diffie_hellman_keys(p=23, g=5, seed=None):
    """Toy Diffie-Hellman key exchange (small p for demo only)."""
    rng = np.random.default_rng(seed)
    a   = int(rng.integers(2,p-1)); b = int(rng.integers(2,p-1))
    A   = pow(g,a,p); B = pow(g,b,p)
    K_a = pow(B,a,p); K_b = pow(A,b,p)
    assert K_a == K_b
    return {"alice_private":a,"bob_private":b,
            "alice_public":A,"bob_public":B,"shared_secret":K_a}
