"""
Regression tests for swarm_download's endgame helpers.

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

import hashlib
import os
import sys

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

from engine_ffi import BLOCK_SIZE  # noqa: E402
from swarm_download import EndgameController, PieceAssembler  # noqa: E402


class FakeMeta:
    def __init__(self, payloads: list[bytes]):
        self.payloads = payloads
        self.num_pieces = len(payloads)
        self.piece_length = max(len(p) for p in payloads)
        self.total_size = sum(len(p) for p in payloads)
        self.piece_hashes = [hashlib.sha1(p).digest() for p in payloads]

    def piece_len(self, piece: int) -> int:
        return len(self.payloads[piece])


class FakeEngine:
    def __init__(self):
        self.priorities = []
        self.rearms = []

    def set_priority(self, tid: int, piece: int, priority: int) -> None:
        self.priorities.append((tid, piece, priority))

    def request_piece(self, tid: int, piece: int) -> None:
        self.rearms.append((tid, piece))


def test_piece_assembler_ignores_duplicate_blocks():
    payload = (b"a" * BLOCK_SIZE) + b"tail"
    meta = FakeMeta([payload])
    done = bytearray(meta.num_pieces)
    asm = PieceAssembler(meta, done)

    first = payload[:BLOCK_SIZE]
    tail = payload[BLOCK_SIZE:]

    added, complete = asm.add_block(0, 0, first)
    assert added
    assert not complete
    assert asm.received[0] == len(first)

    added, complete = asm.add_block(0, 0, first)
    assert not added
    assert not complete
    assert asm.received[0] == len(first)

    added, complete = asm.add_block(0, BLOCK_SIZE, tail)
    assert added
    assert complete
    assert hashlib.sha1(asm.piece_bytes(0)).digest() == meta.piece_hashes[0]


def test_endgame_rearms_only_unfinished_pieces_on_interval():
    meta = FakeMeta([b"a", b"b", b"c", b"d"])
    done = bytearray([0, 1, 0, 0])
    eng = FakeEngine()
    ctl = EndgameController(meta, min_pieces=3, peer_factor=2.0, interval=3.0)

    ctl.maybe_rearm(eng, 7, done, done_count=1, connected=2, now=10.0)
    assert eng.priorities == [(7, 0, 255), (7, 2, 255), (7, 3, 255)]
    assert eng.rearms == [(7, 0), (7, 2), (7, 3)]

    ctl.maybe_rearm(eng, 7, done, done_count=1, connected=2, now=11.0)
    assert eng.rearms == [(7, 0), (7, 2), (7, 3)]

    done[2] = 1
    ctl.maybe_rearm(eng, 7, done, done_count=2, connected=2, now=13.0)
    assert eng.rearms[-2:] == [(7, 0), (7, 3)]


if __name__ == "__main__":
    test_piece_assembler_ignores_duplicate_blocks()
    test_endgame_rearms_only_unfinished_pieces_on_interval()
    print("swarm endgame OK")