"""Piece management — tracks download state of torrent pieces.

Manages which pieces have been downloaded, which are in progress,
and selects the next piece to request from peers.

Ported from org.klomp.snark.PeerCoordinator / Storage piece tracking.
"""

from __future__ import annotations

import enum
import random


class PieceState(enum.Enum):
    MISSING = "missing"
    REQUESTED = "requested"
    COMPLETE = "complete"


class PieceManager:
    """Tracks piece download state for a torrent."""

    def __init__(
        self,
        total_pieces: int,
        piece_length: int,
        total_size: int,
    ) -> None:
        self.total_pieces = total_pieces
        self.piece_length = piece_length
        self.total_size = total_size
        self._states: list[PieceState] = [PieceState.MISSING] * total_pieces

    @property
    def completed_count(self) -> int:
        return sum(1 for s in self._states if s == PieceState.COMPLETE)

    @property
    def is_complete(self) -> bool:
        return self.completed_count == self.total_pieces

    @property
    def percent_complete(self) -> float:
        if self.total_pieces == 0:
            return 100.0
        return (self.completed_count / self.total_pieces) * 100.0

    def state(self, piece_index: int) -> PieceState:
        return self._states[piece_index]

    def mark_complete(self, piece_index: int) -> None:
        self._states[piece_index] = PieceState.COMPLETE

    def mark_requested(self, piece_index: int) -> None:
        if self._states[piece_index] == PieceState.MISSING:
            self._states[piece_index] = PieceState.REQUESTED

    def mark_missing(self, piece_index: int) -> None:
        if self._states[piece_index] != PieceState.COMPLETE:
            self._states[piece_index] = PieceState.MISSING

    def next_needed(self, available: set[int]) -> int | None:
        """Select the next piece to download from available pieces.

        Uses random selection from missing pieces that the peer has.
        """
        candidates = [
            i for i in available
            if i < self.total_pieces and self._states[i] == PieceState.MISSING
        ]
        if not candidates:
            return None
        return random.choice(candidates)

    def bitfield(self) -> bytes:
        """Generate a bitfield representing completed pieces.

        Each bit represents one piece (MSB first).
        """
        num_bytes = (self.total_pieces + 7) // 8
        result = bytearray(num_bytes)
        for i, state in enumerate(self._states):
            if state == PieceState.COMPLETE:
                byte_idx = i // 8
                bit_idx = 7 - (i % 8)
                result[byte_idx] |= (1 << bit_idx)
        return bytes(result)