"""
Key bundle export/import for USB key transfer between devices.

Bundles contain all key material encrypted with a user-provided password.
Format: AES-256-GCM encrypted JSON, key derived via Argon2id.
"""

from __future__ import annotations

import json
import os
import struct
from pathlib import Path

from soosef.exceptions import KeystoreError

# Bundle file magic bytes
BUNDLE_MAGIC = b"SOOBNDL\x00"
BUNDLE_VERSION = 1


def export_bundle(
    identity_dir: Path,
    channel_key_file: Path,
    output_path: Path,
    password: bytes,
) -> None:
    """Export all key material to an encrypted bundle file."""
    from argon2.low_level import Type, hash_secret_raw
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM

    payload = {}

    # Collect identity keys
    priv_path = identity_dir / "private.pem"
    pub_path = identity_dir / "public.pem"
    if priv_path.exists():
        payload["identity_private_pem"] = priv_path.read_text()
    if pub_path.exists():
        payload["identity_public_pem"] = pub_path.read_text()

    # Collect channel key
    if channel_key_file.exists():
        payload["channel_key"] = channel_key_file.read_text().strip()

    if not payload:
        raise KeystoreError("No key material to export.")

    # Encrypt
    salt = os.urandom(32)
    key = hash_secret_raw(
        secret=password,
        salt=salt,
        time_cost=4,
        memory_cost=262144,  # 256 MB
        parallelism=4,
        hash_len=32,
        type=Type.ID,
    )

    nonce = os.urandom(12)
    plaintext = json.dumps(payload).encode()
    aesgcm = AESGCM(key)
    ciphertext = aesgcm.encrypt(nonce, plaintext, None)

    # Write bundle: magic + version + salt + nonce + ciphertext
    with open(output_path, "wb") as f:
        f.write(BUNDLE_MAGIC)
        f.write(struct.pack("<B", BUNDLE_VERSION))
        f.write(salt)
        f.write(nonce)
        f.write(ciphertext)


def import_bundle(
    bundle_path: Path,
    identity_dir: Path,
    channel_key_file: Path,
    password: bytes,
) -> dict[str, bool]:
    """Import key material from an encrypted bundle. Returns what was imported."""
    from argon2.low_level import Type, hash_secret_raw
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM

    data = bundle_path.read_bytes()

    if not data.startswith(BUNDLE_MAGIC):
        raise KeystoreError("Not a valid SooSeF key bundle.")

    offset = len(BUNDLE_MAGIC)
    version = struct.unpack_from("<B", data, offset)[0]
    offset += 1

    if version != BUNDLE_VERSION:
        raise KeystoreError(f"Unsupported bundle version: {version}")

    salt = data[offset : offset + 32]
    offset += 32
    nonce = data[offset : offset + 12]
    offset += 12
    ciphertext = data[offset:]

    key = hash_secret_raw(
        secret=password,
        salt=salt,
        time_cost=4,
        memory_cost=262144,
        parallelism=4,
        hash_len=32,
        type=Type.ID,
    )

    aesgcm = AESGCM(key)
    try:
        plaintext = aesgcm.decrypt(nonce, ciphertext, None)
    except Exception:
        raise KeystoreError("Decryption failed — wrong password or corrupted bundle.")

    payload = json.loads(plaintext)
    imported = {}

    if "identity_private_pem" in payload:
        identity_dir.mkdir(parents=True, exist_ok=True)
        priv_path = identity_dir / "private.pem"
        priv_path.write_text(payload["identity_private_pem"])
        priv_path.chmod(0o600)
        imported["identity_private"] = True

    if "identity_public_pem" in payload:
        identity_dir.mkdir(parents=True, exist_ok=True)
        (identity_dir / "public.pem").write_text(payload["identity_public_pem"])
        imported["identity_public"] = True

    if "channel_key" in payload:
        channel_key_file.parent.mkdir(parents=True, exist_ok=True)
        channel_key_file.write_text(payload["channel_key"])
        channel_key_file.chmod(0o600)
        imported["channel_key"] = True

    return imported


# Channel-key-only export/import — for sharing with collaborators
# without exposing identity keys.
CHANNEL_MAGIC = b"SOOCHNL\x00"


def export_channel_key(
    channel_key: str,
    output_path: Path,
    password: bytes,
) -> None:
    """Export only the channel key to an encrypted file."""
    from argon2.low_level import Type, hash_secret_raw
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM

    salt = os.urandom(32)
    key = hash_secret_raw(
        secret=password,
        salt=salt,
        time_cost=4,
        memory_cost=65536,  # 64 MB — lighter than full bundle for faster sharing
        parallelism=4,
        hash_len=32,
        type=Type.ID,
    )

    nonce = os.urandom(12)
    aesgcm = AESGCM(key)
    ciphertext = aesgcm.encrypt(nonce, channel_key.encode(), None)

    with open(output_path, "wb") as f:
        f.write(CHANNEL_MAGIC)
        f.write(struct.pack("<B", 1))
        f.write(salt)
        f.write(nonce)
        f.write(ciphertext)


def import_channel_key(
    bundle_path: Path,
    password: bytes,
) -> str:
    """Decrypt and return a channel key from an exported bundle."""
    from argon2.low_level import Type, hash_secret_raw
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM

    data = bundle_path.read_bytes()

    if not data.startswith(CHANNEL_MAGIC):
        raise KeystoreError("Not a valid SooSeF channel key bundle.")

    offset = len(CHANNEL_MAGIC) + 1  # magic + version byte
    salt = data[offset : offset + 32]
    offset += 32
    nonce = data[offset : offset + 12]
    offset += 12
    ciphertext = data[offset:]

    key = hash_secret_raw(
        secret=password,
        salt=salt,
        time_cost=4,
        memory_cost=65536,
        parallelism=4,
        hash_len=32,
        type=Type.ID,
    )

    aesgcm = AESGCM(key)
    try:
        plaintext = aesgcm.decrypt(nonce, ciphertext, None)
    except Exception:
        raise KeystoreError("Decryption failed — wrong password or corrupted bundle.")

    return plaintext.decode()


def channel_key_to_qr_data(channel_key: str) -> str:
    """Encode a channel key for QR code display.

    Returns a URI string that can be rendered as a QR code for
    in-person key exchange (e.g., scan from phone/laptop).
    """
    import base64
    import zlib

    compressed = zlib.compress(channel_key.encode(), 9)
    b64 = base64.urlsafe_b64encode(compressed).decode()
    return f"soosef-channel:{b64}"


def channel_key_from_qr_data(qr_data: str) -> str:
    """Decode a channel key from QR code data."""
    import base64
    import zlib

    if not qr_data.startswith("soosef-channel:"):
        raise KeystoreError("Not a valid SooSeF channel key QR code.")
    b64 = qr_data[len("soosef-channel:"):]
    compressed = base64.urlsafe_b64decode(b64)
    return zlib.decompress(compressed).decode()