"""
seed_server.py - A lightweight emulated BitTorrent client (seed side) for
load-testing the peer/engine against "as many peers as possible".

Rather than spinning up N heavyweight libtorrent sessions, this serves one
torrent's data from a single asyncio event loop across many listener sockets.
Each listener is a distinct endpoint, so from the engine's point of view each is
a separate peer: one `add_peer(tid, ip, port)` per listener.

It speaks the plaintext BitTorrent v1 wire protocol a real client would on the
seed side: validates the handshake + info-hash, advertises a full bitfield,
unchokes, and answers `request` messages with `piece` data read (mmap'd) from
disk. That is exactly the subset the engine drives, and it is plaintext because
the engine does not negotiate MSE encryption yet.

Usage
-----
Serve an existing torrent (data already on disk under --data):
    python harness/seed_server.py some.torrent --data /path/to/datadir --peers 64

Generate a random test torrent, serve it, and print the .torrent path:
    python harness/seed_server.py --generate 256M --peers 64 --out /tmp/seedtest

Generate + serve + drive the engine against every peer and report throughput:
    python harness/seed_server.py --generate 256M --peers 64 --out /tmp/seedtest --self-test

In a test, use the SeedSwarm class directly to start listeners in-process and
read `swarm.endpoints`.
"""
from __future__ import annotations

import argparse
import asyncio
import mmap
import os
import secrets
import struct
import threading
import time

PSTR = b"BitTorrent protocol"
HANDSHAKE_LEN = 68

# Wire message ids.
MSG_CHOKE, MSG_UNCHOKE, MSG_INTERESTED, MSG_NOT_INTERESTED = 0, 1, 2, 3
MSG_HAVE, MSG_BITFIELD, MSG_REQUEST, MSG_PIECE, MSG_CANCEL = 4, 5, 6, 7, 8

DRAIN_HIGH_WATER = 1 << 20  # let blocks queue up, apply backpressure past 1 MiB
MAX_BLOCK = 1 << 17         # reject absurd request lengths (128 KiB)


# --------------------------------------------------------------------------- #
# Data source: read (offset, length) from the torrent's concatenated files.
# --------------------------------------------------------------------------- #
class DataSource:
    """Maps the linear piece space onto one or more on-disk files (BT v1 lays
    files out back-to-back). Files are mmap'd for cheap repeated reads."""

    def __init__(self, files: list[tuple[str, int]]):
        self._maps = []  # (global_offset, size, mmap_or_none)
        self._handles = []
        off = 0
        for path, size in files:
            mm = None
            if size > 0:
                fh = open(path, "rb")
                self._handles.append(fh)
                mm = mmap.mmap(fh.fileno(), size, prot=mmap.PROT_READ)
            self._maps.append((off, size, mm))
            off += size
        self.total = off
        self._single = self._maps[0][2] if len(self._maps) == 1 else None

    def read(self, offset: int, length: int) -> bytes:
        if self._single is not None:                 # fast path: one file
            return self._single[offset:offset + length]
        out = bytearray()
        remaining = length
        for foff, size, mm in self._maps:
            if remaining <= 0:
                break
            if offset >= foff + size or offset < foff:
                continue
            local = offset - foff
            take = min(size - local, remaining)
            out += mm[local:local + take]
            offset += take
            remaining -= take
        return bytes(out)

    def close(self):
        for _, _, mm in self._maps:
            if mm is not None:
                mm.close()
        for fh in self._handles:
            fh.close()


def _full_bitfield(num_pieces: int) -> bytes:
    nbytes = (num_pieces + 7) // 8
    bf = bytearray(b"\xff" * nbytes)
    rem = num_pieces & 7
    if rem:                                          # clear pad bits past the end
        bf[-1] = (0xFF << (8 - rem)) & 0xFF
    return bytes(bf)


def _msg(mid: int, payload: bytes = b"") -> bytes:
    return struct.pack(">IB", 1 + len(payload), mid) + payload


# --------------------------------------------------------------------------- #
# The swarm: many listeners sharing one event loop, on a background thread.
# --------------------------------------------------------------------------- #
class SeedSwarm:
    def __init__(self, info_hash: bytes, piece_length: int, num_pieces: int,
                 source: DataSource, host: str = "127.0.0.1"):
        self.info_hash = info_hash
        self.piece_length = piece_length
        self.num_pieces = num_pieces
        self.source = source
        self.host = host
        self.bitfield = _full_bitfield(num_pieces)
        self.endpoints: list[tuple[str, int]] = []
        self.served_bytes = 0
        self._loop = asyncio.new_event_loop()
        self._servers: list[asyncio.AbstractServer] = []
        self._thread = threading.Thread(target=self._run, daemon=True)

    # -- connection handler (one per accepted peer) ------------------------- #
    async def _handle(self, reader: asyncio.StreamReader,
                      writer: asyncio.StreamWriter):
        try:
            hs = await reader.readexactly(HANDSHAKE_LEN)
            if hs[1:20] != PSTR or hs[28:48] != self.info_hash:
                writer.close()
                return
            peer_id = b"-SD0001-" + secrets.token_bytes(12)
            writer.write(bytes([len(PSTR)]) + PSTR + b"\x00" * 8 +
                         self.info_hash + peer_id)
            writer.write(_msg(MSG_BITFIELD, self.bitfield))
            writer.write(_msg(MSG_UNCHOKE))
            await writer.drain()

            while True:
                (length,) = struct.unpack(">I", await reader.readexactly(4))
                if length == 0:
                    continue                          # keep-alive
                body = await reader.readexactly(length)
                mid = body[0]
                if mid == MSG_REQUEST and length >= 13:
                    index, begin, blen = struct.unpack(">III", body[1:13])
                    if blen > MAX_BLOCK:
                        continue
                    data = self.source.read(index * self.piece_length + begin, blen)
                    writer.write(struct.pack(">IB", 9 + len(data), MSG_PIECE) +
                                 struct.pack(">II", index, begin) + data)
                    self.served_bytes += len(data)
                    if writer.transport.get_write_buffer_size() > DRAIN_HIGH_WATER:
                        await writer.drain()
                # interested / not-interested / cancel / choke: nothing to do —
                # we are already unchoked and answer requests as they arrive.
        except (asyncio.IncompleteReadError, ConnectionResetError,
                BrokenPipeError, ConnectionError):
            pass
        finally:
            try:
                writer.close()
            except Exception:
                pass

    async def _make_servers(self, count: int):
        servers, endpoints = [], []
        for _ in range(count):
            srv = await asyncio.start_server(self._handle, self.host, 0)
            servers.append(srv)
            endpoints.append((self.host, srv.sockets[0].getsockname()[1]))
        return servers, endpoints

    def _run(self):
        asyncio.set_event_loop(self._loop)
        self._loop.run_forever()

    def start(self, count: int):
        self._thread.start()
        fut = asyncio.run_coroutine_threadsafe(self._make_servers(count), self._loop)
        self._servers, self.endpoints = fut.result(timeout=15)
        return self.endpoints

    def stop(self):
        def _shutdown():
            for srv in self._servers:
                srv.close()
            self._loop.stop()
        if self._loop.is_running():
            self._loop.call_soon_threadsafe(_shutdown)
        self._thread.join(timeout=5)
        try:
            self._loop.close()
        except Exception:
            pass
        self.source.close()


# --------------------------------------------------------------------------- #
# Torrent loading / generation.
# --------------------------------------------------------------------------- #
def load_torrent(torrent_path: str, data_dir: str):
    """Return (info_hash, piece_length, num_pieces, total_size, DataSource)."""
    from torrent_meta import load_torrent as parse_torrent
    tf = parse_torrent(torrent_path)
    files = []
    for path, size in tf.files:
        files.append((os.path.join(data_dir, path), size))
    missing = [p for p, _ in files if not os.path.exists(p)]
    if missing:
        raise FileNotFoundError(f"data files not found under {data_dir}: {missing}")
    meta = tf.metadata
    return (meta.info_hash, meta.piece_length, meta.num_pieces,
            meta.total_size, DataSource(files))


def generate_torrent(out_dir: str, size: int, piece_size: int = 256 * 1024):
    """Create a random data file + .torrent under out_dir. Returns
    (torrent_path, info_hash, piece_length, num_pieces, total_size, DataSource)."""
    import libtorrent as lt
    os.makedirs(out_dir, exist_ok=True)
    data_path = os.path.join(out_dir, "data.bin")
    with open(data_path, "wb") as f:
        rem = size
        while rem > 0:
            n = min(rem, 8 << 20)
            f.write(os.urandom(n))
            rem -= n
    fs = lt.file_storage()
    lt.add_files(fs, data_path)
    t = lt.create_torrent(fs, piece_size=piece_size)
    t.set_priv(False)
    lt.set_piece_hashes(t, out_dir)
    torrent_path = os.path.join(out_dir, "test.torrent")
    with open(torrent_path, "wb") as f:
        f.write(lt.bencode(t.generate()))
    ih, pl, npc, total, src = load_torrent(torrent_path, out_dir)
    return torrent_path, ih, pl, npc, total, src


def parse_size(s: str) -> int:
    s = s.strip().upper()
    mult = 1
    if s and s[-1] in "KMG":
        mult = {"K": 1024, "M": 1024**2, "G": 1024**3}[s[-1]]
        s = s[:-1]
    return int(float(s) * mult)


# --------------------------------------------------------------------------- #
# Optional: drive the engine against every peer and verify + report rate.
# --------------------------------------------------------------------------- #
def _self_test(torrent_path: str, data_dir: str, endpoints, timeout: float):
    import hashlib
    from engine_ffi import Engine, EngineConfig, STATE_ERROR, ERROR_NAMES
    from harness import load_metadata

    meta = load_metadata(torrent_path)
    lib = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                       "build", "libtorrentpeer.so")
    loops = min(8, max(1, (os.cpu_count() or 1)))
    with Engine(EngineConfig(loop_count=loops, slots_per_loop=4096,
                             max_pipeline=1024), lib_path=lib) as eng:
        tid = eng.add_torrent(meta.info_hash, b"-PC0001-" + secrets.token_bytes(12),
                              meta.piece_length, meta.total_size, meta.num_pieces)
        eng.set_priorities(tid, [1] * meta.num_pieces)
        for ip, port in endpoints:
            eng.add_peer(tid, ip, 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
        t0 = time.time()
        deadline = t0 + 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}")
                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)
        dt = time.time() - t0
        mb = meta.total_size / 1e6
        print(f"self-test OK: {mb:.1f} MB from {len(endpoints)} peers "
              f"in {dt:.2f}s = {mb/dt:.1f} MB/s (peers={eng.status(tid).peers})")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("torrent", nargs="?", help="path to an existing .torrent")
    ap.add_argument("--data", help="directory (or file) holding the torrent's data")
    ap.add_argument("--generate", metavar="SIZE",
                    help="generate a random torrent of this size (e.g. 256M, 1G)")
    ap.add_argument("--out", default="/tmp/seedtest",
                    help="output dir for --generate (default: /tmp/seedtest)")
    ap.add_argument("--peers", type=int, default=32, help="number of listeners")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--piece-size", type=int, default=256 * 1024)
    ap.add_argument("--ports-file", help="write the endpoints (one ip:port/line) here")
    ap.add_argument("--self-test", action="store_true",
                    help="drive the engine against all peers, verify, report rate")
    ap.add_argument("--timeout", type=float, default=300.0)
    args = ap.parse_args()

    torrent_path = args.torrent
    if args.generate:
        size = parse_size(args.generate)
        torrent_path, ih, pl, npc, total, src = generate_torrent(
            args.out, size, args.piece_size)
        data_dir = args.out
        print(f"generated {torrent_path}  ({total} bytes, {npc} pieces x {pl})")
    else:
        if not torrent_path or not args.data:
            ap.error("provide a .torrent and --data, or use --generate SIZE")
        data_dir = args.data if os.path.isdir(args.data) else os.path.dirname(args.data)
        ih, pl, npc, total, src = load_torrent(torrent_path, data_dir)

    swarm = SeedSwarm(ih, pl, npc, src, host=args.host)
    swarm.start(args.peers)
    print(f"seeding {total} bytes on {len(swarm.endpoints)} peers "
          f"({args.host}:{swarm.endpoints[0][1]}..{swarm.endpoints[-1][1]})")
    if args.ports_file:
        with open(args.ports_file, "w") as f:
            f.writelines(f"{ip}:{port}\n" for ip, port in swarm.endpoints)
        print(f"endpoints written to {args.ports_file}")

    try:
        if args.self_test:
            _self_test(torrent_path, data_dir, swarm.endpoints, args.timeout)
        else:
            print("serving; press Ctrl-C to stop")
            while True:
                time.sleep(1.0)
    except KeyboardInterrupt:
        pass
    finally:
        swarm.stop()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())