"""Encryption helpers for Mira archives."""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import os
import zipfile
from contextlib import contextmanager
from typing import Iterator

from mira.temp import make_temp_file
MAGIC = b"MIRA1"
HEADER_LEN_SIZE = 4
CHUNK_SIZE = 64 * 1024
DEFAULT_ITERATIONS = 200_000
SALT_SIZE = 16

CIPHER_STREAM = "hmac-sha256-stream"
CIPHER_AES_GCM = "aes-256-gcm"

STREAM_NONCE_SIZE = 16
STREAM_TAG_SIZE = 32
AES_NONCE_SIZE = 12
AES_TAG_SIZE = 16


def is_encrypted_mira(path: str) -> bool:
    """Return True if a file looks like an encrypted .mira archive."""
    try:
        with open(path, "rb") as handle:
            return handle.read(len(MAGIC)) == MAGIC
    except OSError:
        return False


def encrypt_zip(
    zip_path: str,
    output_path: str,
    key: str,
    cipher: str = CIPHER_STREAM,
) -> None:
    """Encrypt a zip file into a .mira container.

    Args:
        zip_path: Path to a zip archive with build contents.
        output_path: Destination .mira file path.
        key: Encryption passphrase.
        cipher: Cipher name to use.
    """
    if not key:
        raise RuntimeError("Encryption key is required")

    cipher = cipher or CIPHER_STREAM
    if cipher == CIPHER_AES_GCM:
        _encrypt_aes_gcm(zip_path, output_path, key)
        return
    if cipher == CIPHER_STREAM:
        _encrypt_stream(zip_path, output_path, key)
        return
    raise RuntimeError(f"Unsupported cipher: {cipher}")


def decrypt_to_temp_zip(mira_path: str, key: str) -> str:
    """Decrypt an encrypted .mira archive into a temporary zip file."""
    if not key:
        raise RuntimeError("Encryption key is required to decrypt the archive")

    with open(mira_path, "rb") as source:
        header, header_end = _read_header(source)
        cipher = header.get("cipher", CIPHER_STREAM)
        salt = _decode_b64(header.get("salt"), "salt")
        nonce = _decode_b64(header.get("nonce"), "nonce")
        iterations = int(header.get("iterations", DEFAULT_ITERATIONS))
        tag_size = int(header.get("tag_size", _default_tag_size(cipher)))

        source.seek(0, os.SEEK_END)
        file_size = source.tell()
        tag_pos = file_size - tag_size
        if tag_pos <= header_end:
            raise RuntimeError("Invalid .mira archive structure")

        source.seek(tag_pos)
        tag = source.read(tag_size)

        source.seek(header_end)
        remaining = tag_pos - header_end

        temp_path = make_temp_file("mira_decrypt_", ".zip")
        try:
            if cipher == CIPHER_STREAM:
                enc_key, mac_key = _derive_keys(key, salt, iterations)
                mac = hmac.new(mac_key, digestmod=hashlib.sha256)
                counter = 0
                with open(temp_path, "wb") as temp_file:
                    while remaining > 0:
                        read_size = min(CHUNK_SIZE, remaining)
                        chunk = source.read(read_size)
                        if not chunk:
                            break
                        remaining -= len(chunk)
                        mac.update(chunk)
                        plain_chunk, counter = _xor_stream(
                            chunk, enc_key, nonce, counter
                        )
                        temp_file.write(plain_chunk)

                if not hmac.compare_digest(mac.digest(), tag):
                    os.remove(temp_path)
                    raise RuntimeError("Invalid encryption key or corrupted archive")

                return temp_path

            if cipher == CIPHER_AES_GCM:
                _decrypt_aes_gcm_to_path(
                    source,
                    temp_path,
                    nonce,
                    tag,
                    remaining,
                    key,
                    salt,
                    iterations,
                )
                return temp_path

            raise RuntimeError(f"Unsupported cipher: {cipher}")
        except Exception:
            if os.path.exists(temp_path):
                os.remove(temp_path)
            raise


@contextmanager
def open_mira_zip(mira_path: str, key: str | None) -> Iterator[zipfile.ZipFile]:
    """Open a .mira archive as a ZipFile, decrypting if needed."""
    if is_encrypted_mira(mira_path):
        if not key:
            raise RuntimeError("Encryption key required to open this .mira archive")
        temp_path = decrypt_to_temp_zip(mira_path, key)
        zip_file = zipfile.ZipFile(temp_path, "r")
        try:
            yield zip_file
        finally:
            zip_file.close()
            if os.path.exists(temp_path):
                os.remove(temp_path)
    else:
        zip_file = zipfile.ZipFile(mira_path, "r")
        try:
            yield zip_file
        finally:
            zip_file.close()


def _encrypt_stream(zip_path: str, output_path: str, key: str) -> None:
    salt = os.urandom(SALT_SIZE)
    nonce = os.urandom(STREAM_NONCE_SIZE)
    iterations = DEFAULT_ITERATIONS
    enc_key, mac_key = _derive_keys(key, salt, iterations)

    header = {
        "version": 1,
        "kdf": "pbkdf2-sha256",
        "iterations": iterations,
        "salt": base64.b64encode(salt).decode("ascii"),
        "nonce": base64.b64encode(nonce).decode("ascii"),
        "cipher": CIPHER_STREAM,
        "compression": "zip",
        "tag_size": STREAM_TAG_SIZE,
    }
    header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")

    with open(zip_path, "rb") as source, open(output_path, "wb") as target:
        target.write(MAGIC)
        target.write(len(header_bytes).to_bytes(HEADER_LEN_SIZE, "big"))
        target.write(header_bytes)

        mac = hmac.new(mac_key, digestmod=hashlib.sha256)
        counter = 0
        while True:
            chunk = source.read(CHUNK_SIZE)
            if not chunk:
                break
            cipher_chunk, counter = _xor_stream(chunk, enc_key, nonce, counter)
            mac.update(cipher_chunk)
            target.write(cipher_chunk)

        target.write(mac.digest())


def _encrypt_aes_gcm(zip_path: str, output_path: str, key: str) -> None:
    Cipher, algorithms, modes, _ = _require_crypto()

    salt = os.urandom(SALT_SIZE)
    nonce = os.urandom(AES_NONCE_SIZE)
    iterations = DEFAULT_ITERATIONS
    enc_key = _derive_key(key, salt, iterations, 32)

    header = {
        "version": 1,
        "kdf": "pbkdf2-sha256",
        "iterations": iterations,
        "salt": base64.b64encode(salt).decode("ascii"),
        "nonce": base64.b64encode(nonce).decode("ascii"),
        "cipher": CIPHER_AES_GCM,
        "compression": "zip",
        "tag_size": AES_TAG_SIZE,
    }
    header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")

    cipher = Cipher(algorithms.AES(enc_key), modes.GCM(nonce))
    encryptor = cipher.encryptor()

    with open(zip_path, "rb") as source, open(output_path, "wb") as target:
        target.write(MAGIC)
        target.write(len(header_bytes).to_bytes(HEADER_LEN_SIZE, "big"))
        target.write(header_bytes)

        while True:
            chunk = source.read(CHUNK_SIZE)
            if not chunk:
                break
            cipher_chunk = encryptor.update(chunk)
            if cipher_chunk:
                target.write(cipher_chunk)

        encryptor.finalize()
        target.write(encryptor.tag)


def _decrypt_aes_gcm_to_path(
    handle,
    output_path: str,
    nonce: bytes,
    tag: bytes,
    remaining: int,
    key: str,
    salt: bytes,
    iterations: int,
) -> None:
    Cipher, algorithms, modes, InvalidTag = _require_crypto()
    enc_key = _derive_key(key, salt, iterations, 32)

    cipher = Cipher(algorithms.AES(enc_key), modes.GCM(nonce, tag))
    decryptor = cipher.decryptor()

    with open(output_path, "wb") as out:
        while remaining > 0:
            read_size = min(CHUNK_SIZE, remaining)
            chunk = handle.read(read_size)
            if not chunk:
                break
            remaining -= len(chunk)
            plain = decryptor.update(chunk)
            if plain:
                out.write(plain)

        try:
            decryptor.finalize()
        except InvalidTag as exc:
            raise RuntimeError("Invalid encryption key or corrupted archive") from exc


def _derive_keys(password: str, salt: bytes, iterations: int) -> tuple[bytes, bytes]:
    key = hashlib.pbkdf2_hmac(
        "sha256", password.encode("utf-8"), salt, iterations, dklen=64
    )
    return key[:32], key[32:]


def _derive_key(password: str, salt: bytes, iterations: int, length: int) -> bytes:
    return hashlib.pbkdf2_hmac(
        "sha256", password.encode("utf-8"), salt, iterations, dklen=length
    )


def _xor_stream(
    data: bytes, key: bytes, nonce: bytes, counter: int
) -> tuple[bytes, int]:
    out = bytearray(len(data))
    offset = 0
    while offset < len(data):
        block = hmac.new(
            key, nonce + counter.to_bytes(8, "big"), hashlib.sha256
        ).digest()
        block_len = min(len(block), len(data) - offset)
        for i in range(block_len):
            out[offset + i] = data[offset + i] ^ block[i]
        offset += block_len
        counter += 1
    return bytes(out), counter


def _read_header(handle) -> tuple[dict, int]:
    magic = handle.read(len(MAGIC))
    if magic != MAGIC:
        raise RuntimeError("Not an encrypted .mira archive")

    header_len_bytes = handle.read(HEADER_LEN_SIZE)
    if len(header_len_bytes) != HEADER_LEN_SIZE:
        raise RuntimeError("Invalid .mira header length")
    header_len = int.from_bytes(header_len_bytes, "big")
    if header_len <= 0:
        raise RuntimeError("Invalid .mira header size")

    header_bytes = handle.read(header_len)
    if len(header_bytes) != header_len:
        raise RuntimeError("Invalid .mira header data")

    try:
        header = json.loads(header_bytes.decode("utf-8"))
    except json.JSONDecodeError as exc:
        raise RuntimeError("Invalid .mira header JSON") from exc

    header_end = len(MAGIC) + HEADER_LEN_SIZE + header_len
    return header, header_end


def _decode_b64(value: str | None, label: str) -> bytes:
    if not value:
        raise RuntimeError(f"Missing {label} in .mira header")
    try:
        return base64.b64decode(value.encode("ascii"))
    except (ValueError, UnicodeEncodeError) as exc:
        raise RuntimeError(f"Invalid base64 value for {label}") from exc


def _default_tag_size(cipher: str) -> int:
    if cipher == CIPHER_AES_GCM:
        return AES_TAG_SIZE
    return STREAM_TAG_SIZE


def _require_crypto():
    try:
        from cryptography.exceptions import InvalidTag
        from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
    except ModuleNotFoundError as exc:
        raise RuntimeError(
            "AES-GCM requires cryptography. Install with: pip install mira[crypto]"
        ) from exc
    return Cipher, algorithms, modes, InvalidTag
