"""
End-to-end tests against a local libtorrent seed:

  * test_localseed_roundtrip      - full seed, download everything, verify bytes
  * test_partial_availability      - seed has only some pieces; the peer must
                                     download exactly those and skip the rest
                                     (the bug a fixed in-order schedule hit)
  * test_priority_order            - pieces are selected highest-priority-first

Run with:  python tests/test_localseed.py   (or)  python -m pytest tests/
"""
from __future__ import annotations

import hashlib
import os
import socket
import subprocess
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 Downloader, load_metadata  # noqa: E402

LIB = os.path.join(ROOT, "build", "libtorrentpeer.so")
PIECE_SIZE = 256 * 1024


def pick_listen_port() -> int:
    """Reserve a free loopback TCP port briefly, then hand it to libtorrent."""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


def ensure_built():
    if os.path.exists(LIB):
        return
    build = os.path.join(ROOT, "build")
    subprocess.run(["cmake", "-S", ROOT, "-B", build,
                    "-DCMAKE_BUILD_TYPE=Release"], check=True)
    subprocess.run(["cmake", "--build", build], check=True)
    assert os.path.exists(LIB), "build did not produce libtorrentpeer.so"


def make_torrent(root: str, size: int) -> tuple[str, bytes, str]:
    data = os.urandom(size)
    path = os.path.join(root, "data.bin")
    with open(path, "wb") as f:
        f.write(data)
    fs = lt.file_storage()
    lt.add_files(fs, path)
    t = lt.create_torrent(fs, piece_size=PIECE_SIZE)
    t.set_priv(False)
    lt.set_piece_hashes(t, root)
    torrent_path = os.path.join(root, "test.torrent")
    with open(torrent_path, "wb") as f:
        f.write(lt.bencode(t.generate()))
    return torrent_path, data, path


def _session() -> tuple[lt.session, int]:
    port = pick_listen_port()
    return lt.session({
        "listen_interfaces": f"127.0.0.1:{port}",
        "enable_dht": False, "enable_lsd": False,
        "enable_upnp": False, "enable_natpmp": False,
        # Force plaintext so our minimal peer's handshake is accepted.
        "in_enc_policy": int(lt.enc_policy.disabled),
        "out_enc_policy": int(lt.enc_policy.disabled),
    }), port


def start_full_seed(root: str, torrent_path: str):
    ses, port = _session()
    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 start_partial_seed(root: str, torrent_path: str, data_path: str, size: int):
    """Corrupt the second half on disk so the seeder only HAS the first half."""
    with open(data_path, "r+b") as f:
        f.seek(size // 2)
        f.write(b"\x00" * (size - size // 2))
    ses, port = _session()
    h = ses.add_torrent({"ti": lt.torrent_info(torrent_path), "save_path": root})
    deadline = time.time() + 15
    while time.time() < deadline:
        state = str(h.status().state)
        if "checking" not in state and "queued" not in state:
            break
        time.sleep(0.1)
    time.sleep(0.3)
    avail = {i for i, b in enumerate(h.status().pieces) if b}
    return ses, h, ses.listen_port() or port, avail


def drive(dl: Downloader, meta, port: int, prio: bytearray, want: set,
          timeout: float = 30.0):
    """Drive the peer until `want` pieces are verified. Returns completion order."""
    for i in range(meta.num_pieces):
        dl.buffers[i] = bytearray(meta.piece_len(i))
    dl.peer.start("127.0.0.1", port)
    dl.peer.set_priorities(prio)
    order, got = [], set()
    deadline = time.time() + timeout
    while got != want and time.time() < deadline:
        st = dl.peer.status()
        if st.state == 6:  # STATE_ERROR
            raise RuntimeError(f"peer error {st.error}")
        descs = dl.peer.poll_ready()
        if not descs:
            dl.peer.wait(100)
            continue
        for x in descs:
            buf = dl.buffers[x.piece]
            buf[x.begin:x.begin + x.len] = dl.peer.block_data(x.slot, x.len)
            dl.peer.release(x.slot)
            dl.received[x.piece] += x.len
            if dl.received[x.piece] >= meta.piece_len(x.piece) and x.piece not in got:
                assert hashlib.sha1(bytes(buf)).digest() == meta.piece_hashes[x.piece]
                got.add(x.piece)
                order.append(x.piece)
                dl.peer.set_priority(x.piece, 0)
    return order, got


def test_localseed_roundtrip():
    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_full_seed(root, torrent)
        try:
            dl = Downloader(meta, num_slots=1024, max_pipeline=512, lib_path=LIB)
            try:
                t0 = time.time()
                got = dl.download("127.0.0.1", port, timeout=60.0)
                dt = time.time() - t0
            finally:
                dl.close()
        finally:
            ses.remove_torrent(h)
        assert got == original, "downloaded bytes differ from original"
        print(f"roundtrip OK: {size/1e6:.1f} MB in {dt:.2f}s "
              f"({size/1e6/dt:.0f} MB/s)")


def test_partial_availability():
    ensure_built()
    size = 8 * 1024 * 1024
    with tempfile.TemporaryDirectory() as root:
        torrent, _, data_path = make_torrent(root, size)
        meta = load_metadata(torrent)
        ses, h, port, avail = start_partial_seed(root, torrent, data_path, size)
        try:
            dl = Downloader(meta, num_slots=512, max_pipeline=256, lib_path=LIB)
            try:
                prio = bytearray([1] * meta.num_pieces)  # want everything
                _, got = drive(dl, meta, port, prio, want=avail, timeout=30.0)
                # peer must idle, not stall requesting missing pieces
                time.sleep(0.3)
                assert dl.peer.status().outstanding == 0
            finally:
                dl.close()
        finally:
            ses.remove_torrent(h)
        assert 0 < len(avail) < meta.num_pieces, "test needs a partial seed"
        assert got == avail, (sorted(got), sorted(avail))
        print(f"partial OK: downloaded {len(got)}/{meta.num_pieces} available "
              f"pieces, skipped {meta.num_pieces - len(avail)} missing")


def test_priority_order():
    ensure_built()
    size = 2 * 1024 * 1024  # 8 pieces
    with tempfile.TemporaryDirectory() as root:
        torrent, _, _ = make_torrent(root, size)
        meta = load_metadata(torrent)
        ses, h, port = start_full_seed(root, torrent)
        try:
            # pipeline=1 makes selection order observable; piece i gets priority
            # i+1, so we expect strictly descending completion order.
            dl = Downloader(meta, num_slots=4, max_pipeline=1, lib_path=LIB)
            try:
                prio = bytearray([i + 1 for i in range(meta.num_pieces)])
                want = set(range(meta.num_pieces))
                order, _ = drive(dl, meta, port, prio, want=want, timeout=30.0)
            finally:
                dl.close()
        finally:
            ses.remove_torrent(h)
        expected = list(range(meta.num_pieces - 1, -1, -1))
        assert order == expected, order
        print(f"priority order OK: {order}")


if __name__ == "__main__":
    test_localseed_roundtrip()
    test_partial_availability()
    test_priority_order()