"""Test-only `.sientia` writer, mirrored field for field from the producer.

There is no production writer in this repository (change `import-model-workflow`, design D1): the
producer is the Streamlit app in `sientia-projects-templates`. This factory exists for one reason —
building the **hostile** variants the gate tests need, which the producer cannot emit. Its output is
pinned against the golden fixture by `test_format.py::test_factory_header_matches_fixture_layout`,
so a drift in either direction shows up as a failing test.

Mirrored from:
- `app/src/operations/model_export/crypto.py` — header layout, KDF, `secretstream` framing
- `app/src/operations/model_export/bundle.py` — zip layout and the two-top-level-key JSON

Deliberate deviation: the default KDF cost is a fraction of the producer's (8 MiB / 1 pass instead
of 256 MiB / 3), because the reader honours whatever the header declares and a test suite that
derives thirty 256 MiB Argon2id keys is a test suite nobody runs. The golden fixture is what covers
the production cost parameters.
"""

from __future__ import annotations

import base64
import hashlib
import io
import json
import os
import stat
import struct
import unicodedata
import zipfile
from dataclasses import dataclass, field
from typing import Any

import nacl.bindings as sodium
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

# --- mirrored from the producer's crypto.py ---
MAGIC = b'SIENTIA1'
FORMAT_VERSION = 1
AEAD_ID = 1
KDF_ID = 1
SALT_SIZE = 16
RESERVED_SIGNATURE_SIZE = 64
CHUNK_SIZE = 64 * 1024
HEADER_STRUCT = struct.Struct('>8sBBB16sQQB32s64s')
HEADER_SIZE = HEADER_STRUCT.size
PRODUCER_KDF_MEMLIMIT_BYTES = 256 * 1024 * 1024
PRODUCER_KDF_OPSLIMIT = 3

# Cheap-but-valid cost parameters for the suite; see the module docstring.
TEST_KDF_MEMLIMIT_BYTES = 8 * 1024 * 1024
TEST_KDF_OPSLIMIT = 1

DEFAULT_PASSWORD = 'Factory1Model!Test'
DEFAULT_ENVELOPE_KEY = bytes(range(32))

MODEL_DIRS = ('prediction_model', 'data_model')
MODEL_FILES = ('MLmodel', 'model.pkl', 'conda.yaml', 'python_env.yaml', 'requirements.txt')
CSV_FILES = ('train_data.csv', 'test_data.csv', 'raw_data.csv')

DEFAULT_METADATA = {
    'model_name': 'factory_model',
    'experiment_name': 'factory_experiment',
    'run_id': 'aa11bb22cc33dd44ee55ff6677889900',
    'model_version': '2',
    'model_project': 'factory_project',
    'export_timestamp': '2026-08-25T12:00:00+00:00',
}

DEFAULT_PARAMETERS = {
    'model_name': 'factory_model',
    'target_variable': 'SE',
    'input_variables': "['feature_a', 'feature_b']",
    'retrain': 'False',
}


def default_tree() -> dict[str, bytes]:
    """The artifact tree a well-formed bundle carries, as archive-name -> bytes."""
    files: dict[str, bytes] = {}
    for directory in MODEL_DIRS:
        for name in MODEL_FILES:
            files[f'artifacts/{directory}/{name}'] = f'{directory}/{name}\n'.encode()
    files['artifacts/model_card/model_card.json'] = json.dumps({'target': 'SE'}).encode()
    for name in CSV_FILES:
        files[f'artifacts/{name}'] = b'datetime,feature_a,SE\n2026-08-25T00:00:00,1.0,2.0\n'
    return files


def build_metadata_document(
    *,
    metadata: dict[str, Any] | None = None,
    parameters: dict[str, Any] | None = None,
    drop_metadata_field: str | None = None,
    drop_top_level_key: str | None = None,
    drop_parameter: str | None = None,
) -> dict[str, Any]:
    """Build `metadata.json`, with hooks for the malformed variants gate 6 must reject."""
    meta = {**DEFAULT_METADATA, **(metadata or {})}
    params = {**DEFAULT_PARAMETERS, **(parameters or {})}
    if drop_metadata_field:
        meta.pop(drop_metadata_field, None)
    if drop_parameter:
        params.pop(drop_parameter, None)
    document: dict[str, Any] = {'parameters': params, 'metadata': meta}
    if drop_top_level_key:
        document.pop(drop_top_level_key, None)
    return document


@dataclass
class ZipSpec:
    """How to lay out the zip inside the bundle."""

    files: dict[str, bytes] = field(default_factory=default_tree)
    document: dict[str, Any] | None = None
    include_metadata: bool = True
    symlinks: dict[str, str] = field(default_factory=dict)
    absolute_entry: str | None = None
    traversal_entry: str | None = None
    bomb_bytes: int = 0
    extra_files: dict[str, bytes] = field(default_factory=dict)
    # name -> a raw `stat.S_IF*` mode, for an entry that is neither a symlink, a regular file nor a
    # directory (e.g. `stat.S_IFIFO`) — gate 4 must reject these on their file-type bits alone.
    special_mode_entries: dict[str, int] = field(default_factory=dict)
    # Directory entries the producer never writes (`build_zip` stores none), but a gate 5 defense
    # this repository still carries has to be exercised by something.
    directory_entries: list[str] = field(default_factory=list)
    # A file name already in `files`/`extra_files`, written a second time under the same name — two
    # archive entries claiming the same path, which only gate 5's extraction can catch.
    repeat_file: str | None = None


def build_zip(spec: ZipSpec | None = None) -> bytes:
    """Zip `artifacts/` plus `metadata.json`, mirroring the producer's `build_zip`.

    `metadata.json` is written first with `json.dumps` and no indent, entries use `/` separators,
    and no directory entries are stored — all three matter to the reader.
    """
    spec = spec or ZipSpec()
    buffer = io.BytesIO()
    with zipfile.ZipFile(buffer, 'w', zipfile.ZIP_DEFLATED) as archive:
        if spec.include_metadata:
            document = spec.document if spec.document is not None else build_metadata_document()
            archive.writestr('metadata.json', json.dumps(document))
        for name, content in {**spec.files, **spec.extra_files}.items():
            archive.writestr(_regular_file_info(name), content)
        for name, target in spec.symlinks.items():
            info = zipfile.ZipInfo(name)
            info.create_system = 3
            info.external_attr = (0o120777 << 16) | 0o200000
            archive.writestr(info, target)
        if spec.absolute_entry:
            archive.writestr(_regular_file_info(spec.absolute_entry), b'absolute\n')
        if spec.traversal_entry:
            archive.writestr(_regular_file_info(spec.traversal_entry), b'traversal\n')
        if spec.bomb_bytes:
            archive.writestr(_regular_file_info('artifacts/bomb.csv'), b'\0' * spec.bomb_bytes)
        for name, mode in spec.special_mode_entries.items():
            info = zipfile.ZipInfo(name)
            info.create_system = 3
            info.external_attr = mode << 16
            archive.writestr(info, b'')
        for name in spec.directory_entries:
            info = zipfile.ZipInfo(name if name.endswith('/') else f'{name}/')
            info.create_system = 3
            info.external_attr = (stat.S_IFDIR | 0o755) << 16
            archive.writestr(info, b'')
        if spec.repeat_file:
            content = {**spec.files, **spec.extra_files}[spec.repeat_file]
            archive.writestr(_regular_file_info(spec.repeat_file), content)
    return buffer.getvalue()


def _regular_file_info(name: str) -> zipfile.ZipInfo:
    """A zip entry carrying a Unix regular-file mode, as `ZipFile.write` would store it."""
    info = zipfile.ZipInfo(name)
    info.create_system = 3
    info.compress_type = zipfile.ZIP_DEFLATED
    info.external_attr = 0o100644 << 16
    return info


def derive_key(password: str, salt: bytes, opslimit: int, memlimit: int) -> bytes:
    """Argon2id over the NFC-normalised password — the producer's `_derive_key`."""
    normalized = unicodedata.normalize('NFC', password)
    return sodium.crypto_pwhash_alg(
        sodium.crypto_secretstream_xchacha20poly1305_KEYBYTES,
        normalized.encode('utf-8'),
        salt,
        opslimit,
        memlimit,
        sodium.crypto_pwhash_ALG_ARGON2ID13,
    )


def encrypt_archive(  # noqa: PLR0913 - every knob here exists to build one hostile variant
    plaintext: bytes,
    password: str = DEFAULT_PASSWORD,
    *,
    format_version: int = FORMAT_VERSION,
    aead_id: int = AEAD_ID,
    kdf_id: int = KDF_ID,
    memlimit: int = TEST_KDF_MEMLIMIT_BYTES,
    opslimit: int = TEST_KDF_OPSLIMIT,
    declared_memlimit: int | None = None,
    declared_opslimit: int | None = None,
    parallelism: int = 1,
    reserved_signature: bytes | None = None,
    magic: bytes = MAGIC,
    digest_override: bytes | None = None,
    salt: bytes | None = None,
    drop_final_tag: bool = False,
) -> bytes:
    """Produce a full `.sientia` byte stream, mirroring the producer's `encrypt_archive`.

    The `declared_*` arguments write cost parameters into the header that differ from the ones
    actually used, which is how a "512 MiB header" is built without allocating 512 MiB.
    `drop_final_tag` pushes every chunk as `MESSAGE`, so the stream is framed as if more were
    coming — the truncation the digest cannot see.
    """
    salt = salt if salt is not None else os.urandom(SALT_SIZE)
    key = derive_key(password, salt, opslimit, memlimit)

    state = sodium.crypto_secretstream_xchacha20poly1305_state()
    stream_header = sodium.crypto_secretstream_xchacha20poly1305_init_push(state, key)

    chunks = []
    offset = 0
    total = len(plaintext)
    while offset < total:
        chunk = plaintext[offset : offset + CHUNK_SIZE]
        offset += len(chunk)
        is_last = offset >= total and not drop_final_tag
        tag = (
            sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL
            if is_last
            else sodium.crypto_secretstream_xchacha20poly1305_TAG_MESSAGE
        )
        chunks.append(sodium.crypto_secretstream_xchacha20poly1305_push(state, chunk, None, tag))
    ciphertext = b''.join(chunks)

    digest = digest_override if digest_override is not None else hashlib.sha256(ciphertext).digest()
    header = HEADER_STRUCT.pack(
        magic,
        format_version,
        aead_id,
        kdf_id,
        salt,
        declared_memlimit if declared_memlimit is not None else memlimit,
        declared_opslimit if declared_opslimit is not None else opslimit,
        parallelism,
        digest,
        reserved_signature if reserved_signature is not None else bytes(RESERVED_SIGNATURE_SIZE),
    )
    return header + stream_header + ciphertext


def make_bundle(
    zip_spec: ZipSpec | None = None,
    password: str = DEFAULT_PASSWORD,
    **header_kwargs: Any,
) -> bytes:
    """The common case: a well-formed tree, encrypted with the test cost parameters."""
    return encrypt_archive(build_zip(zip_spec), password, **header_kwargs)


def tamper_ciphertext(bundle: bytes, offset_from_stream: int = 8) -> bytes:
    """Flip one byte inside the chunk stream **and** repair the header digest.

    Repairing the digest is what makes this a gate-3 test rather than a gate-2 one: the file is
    self-consistent and only the AEAD can tell it was altered.
    """
    position = HEADER_SIZE + 24 + offset_from_stream
    altered = bytearray(bundle)
    altered[position] ^= 0xFF
    repaired = hashlib.sha256(bytes(altered[HEADER_SIZE + 24 :])).digest()
    altered[44:76] = repaired
    return bytes(altered)


def truncate_bundle(bundle: bytes, drop_bytes: int = 32) -> bytes:
    """Cut the tail off a bundle, leaving the header's digest describing bytes that are gone."""
    return bundle[:-drop_bytes]


def build_envelope(password: str = DEFAULT_PASSWORD, key: bytes = DEFAULT_ENVELOPE_KEY) -> str:
    """`base64(nonce || ciphertext || tag)` under AES-256-GCM, as the frontend will send it."""
    nonce = os.urandom(12)
    sealed = AESGCM(key).encrypt(nonce, password.encode('utf-8'), None)
    return base64.b64encode(nonce + sealed).decode('ascii')


def envelope_key_b64(key: bytes = DEFAULT_ENVELOPE_KEY) -> str:
    """The key as `IMPORT_PASSWORD_KEY` carries it."""
    return base64.b64encode(key).decode('ascii')