"""
MSE / PE encryption tests against a libtorrent seed forced into encrypted-only
mode (in/out_enc_policy = forced, allowed_enc_level = rc4). A plaintext engine
cannot complete the handshake with such a seed, so a successful byte-for-byte
download proves the MSE transport works end to end.

  * test_encrypted_download   - engine with encryption=1 (offer RC4+plaintext)
  * test_require_rc4          - engine with encryption=2 (RC4 only)

Run with:  python tests/test_encryption.py
"""
from __future__ import annotations

import hashlib
import os
import sys
import tempfile
import time

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(ROOT, "harness"))

import libtorrent as lt  # noqa: E402

from harness import load_metadata  # noqa: E402
from engine_ffi import Engine, EngineConfig, STATE_ERROR, ERROR_NAMES  # noqa: E402
from test_localseed import ensure_built, make_torrent, pick_listen_port  # noqa: E402

LIB = os.path.join(ROOT, "build", "libtorrentpeer.so")


def start_encrypted_seed(root: str, torrent_path: str):
    """A seed that REQUIRES MSE/RC4 encryption (refuses plaintext)."""
    port = pick_listen_port()
    ses = lt.session({
        "listen_interfaces": f"127.0.0.1:{port}",
        "enable_dht": False, "enable_lsd": False,
        "enable_upnp": False, "enable_natpmp": False,
        "in_enc_policy": int(lt.enc_policy.forced),
        "out_enc_policy": int(lt.enc_policy.forced),
        "allowed_enc_level": int(lt.enc_level.rc4),
        "prefer_rc4": True,
    })
    h = ses.add_torrent({"ti": lt.torrent_info(torrent_path), "save_path": root,
                         "flags": lt.torrent_flags.seed_mode})
    deadline = time.time() + 15
    while time.time() < deadline and not h.status().is_seeding:
        time.sleep(0.1)
    assert h.status().is_seeding, "seed did not become ready"
    return ses, h, ses.listen_port() or port


def _download(meta, port, encryption, timeout=60.0):
    with Engine(EngineConfig(loop_count=1, slots_per_loop=1024,
                             max_pipeline=512, encryption=encryption),
                lib_path=LIB) as eng:
        tid = eng.add_torrent(meta.info_hash, b"-PC0001-" + os.urandom(12),
                              meta.piece_length, meta.total_size, meta.num_pieces)
        eng.set_priorities(tid, [1] * meta.num_pieces)
        eng.add_peer(tid, "127.0.0.1", port)

        buffers = [bytearray(meta.piece_len(i)) for i in range(meta.num_pieces)]
        received = [0] * meta.num_pieces
        done = bytearray(meta.num_pieces)
        done_count = 0
        deadline = time.time() + timeout
        while done_count < meta.num_pieces:
            st = eng.status(tid)
            if st.state == STATE_ERROR:
                raise RuntimeError(f"engine error: {ERROR_NAMES[st.error]}")
            descs = eng.poll_ready()
            if not descs:
                eng.wait(200)
                if time.time() > deadline:
                    raise TimeoutError(f"stalled at {done_count}/{meta.num_pieces} "
                                       f"(connected={st.peers_connected} "
                                       f"failed={st.peers_failed})")
                continue
            for x in descs:
                buf = buffers[x.piece]
                buf[x.begin:x.begin + x.len] = eng.block_data(x.loop, x.slot, x.len)
                eng.release(x.loop, x.slot)
                received[x.piece] += x.len
                if not done[x.piece] and received[x.piece] >= meta.piece_len(x.piece):
                    if hashlib.sha1(bytes(buf)).digest() != meta.piece_hashes[x.piece]:
                        raise ValueError(f"piece {x.piece} hash mismatch")
                    done[x.piece] = 1
                    done_count += 1
                    eng.set_priority(tid, x.piece, 0)
        return b"".join(bytes(b) for b in buffers)


def test_encrypted_download():
    ensure_built()
    size = 8 * 1024 * 1024
    with tempfile.TemporaryDirectory() as root:
        torrent, original, _ = make_torrent(root, size)
        meta = load_metadata(torrent)
        ses, h, port = start_encrypted_seed(root, torrent)
        try:
            got = _download(meta, port, encryption=1)
        finally:
            ses.remove_torrent(h)
        assert got == original, "decrypted bytes differ from original"
    print(f"encrypted (offer RC4+plain) OK: {size/1e6:.1f} MB over MSE")


def test_require_rc4():
    ensure_built()
    size = 8 * 1024 * 1024
    with tempfile.TemporaryDirectory() as root:
        torrent, original, _ = make_torrent(root, size)
        meta = load_metadata(torrent)
        ses, h, port = start_encrypted_seed(root, torrent)
        try:
            got = _download(meta, port, encryption=2)
        finally:
            ses.remove_torrent(h)
        assert got == original, "decrypted bytes differ from original"
    print(f"encrypted (require RC4) OK: {size/1e6:.1f} MB over MSE")


if __name__ == "__main__":
    test_encrypted_download()
    test_require_rc4()