"""BB84 QKD protocol implementation."""

import numpy as np

from ..core import Measurement, QuantumChannel, Qubit
from .base import BaseProtocol


class BB84(BaseProtocol):
    """Implementation of the BB84 quantum key distribution protocol.

    BB84 is the first and most well-known quantum key distribution protocol,
    developed by Charles Bennett and Gilles Brassard in 1984.
    """

    def __init__(
        self,
        channel: QuantumChannel,
        key_length: int = 100,
        security_threshold: float = 0.11,
    ):
        """Initialize the BB84 protocol.

        Args:
            channel: Quantum channel for qubit transmission
            key_length: Desired length of the final key
            security_threshold: Maximum QBER value considered secure

        """
        super().__init__(channel, key_length)

        # BB84-specific parameters
        self.bases = ["computational", "hadamard"]
        self.security_threshold = security_threshold  # 11% QBER threshold for BB84

        # Number of qubits to send (we'll send more than needed to account for sifting)
        self.num_qubits = key_length * 3  # Send 3x more qubits than needed

        # Alice's random bits and bases
        self.alice_bits = []
        self.alice_bases = []

        # Bob's measurement results and bases
        self.bob_results = []
        self.bob_bases = []

    def prepare_states(self) -> list[Qubit]:
        """Prepare quantum states for transmission.

        In BB84, Alice randomly chooses bits and bases, and prepares qubits
        in the corresponding states.

        Returns:
            List of qubits to be sent through the quantum channel

        """
        qubits = []
        self.alice_bits = []
        self.alice_bases = []

        for _ in range(self.num_qubits):
            # Alice randomly chooses a bit (0 or 1)
            bit = np.random.randint(0, 2)
            self.alice_bits.append(bit)

            # Alice randomly chooses a basis
            basis = np.random.choice(self.bases)
            self.alice_bases.append(basis)

            # Prepare the qubit in the appropriate state
            if basis == "computational":
                # Computational basis: |0> or |1>
                qubit = Qubit.zero() if bit == 0 else Qubit.one()
            else:  # hadamard basis
                # Hadamard basis: |+> or |->
                qubit = Qubit.plus() if bit == 0 else Qubit.minus()

            qubits.append(qubit)

        return qubits

    def measure_states(self, qubits: list[Qubit]) -> list[int]:
        """Measure received quantum states.

        In BB84, Bob randomly chooses bases to measure in.

        Args:
            qubits: List of received qubits

        Returns:
            List of measurement results

        """
        self.bob_results = []
        self.bob_bases = []

        for qubit in qubits:
            if qubit is None:
                # Qubit was lost in the channel
                self.bob_results.append(None)
                self.bob_bases.append(None)
                continue

            # Bob randomly chooses a basis
            basis = np.random.choice(self.bases)
            self.bob_bases.append(basis)

            # Measure in the chosen basis
            result = Measurement.measure_in_basis(qubit, basis)
            self.bob_results.append(result)

        return self.bob_results

    def sift_keys(self) -> tuple[list[int], list[int]]:
        """Sift the raw keys to keep only measurements in matching bases.

        In BB84, Alice and Bob publicly compare their bases and keep only
        the bits where they used the same basis.

        Returns:
            Tuple of (alice_sifted_key, bob_sifted_key)

        """
        alice_sifted = []
        bob_sifted = []

        for i in range(self.num_qubits):
            # Skip if Bob didn't receive the qubit
            if self.bob_bases[i] is None:
                continue

            # Check if Alice and Bob used the same basis
            if self.alice_bases[i] == self.bob_bases[i]:
                sifted_alice_key.append(int(self.alice_bits[i]))
                sifted_bob_key.append(int(self.bob_results[i]))

        return alice_sifted, bob_sifted

    def estimate_qber(self) -> float:
        """Estimate the Quantum Bit Error Rate (QBER).

        In BB84, Alice and Bob publicly compare a subset of their sifted keys
        to estimate the error rate.

        Returns:
            Estimated QBER value

        """
        alice_sifted, bob_sifted = self.sift_keys()

        # If we don't have enough bits for estimation, return a high QBER
        if len(alice_sifted) < 10:
            return 1.0

        # Use the full sifted key for QBER estimation in tests
        sample_size = len(alice_sifted)
        indices = range(sample_size)

        # Count errors in the sample
        errors = 0
        for idx in indices:
            if alice_sifted[idx] != bob_sifted[idx]:
                errors += 1

        # Calculate QBER
        qber = errors / sample_size
        print(f"BB84 Estimated QBER: {qber}")
        return qber

    def _get_security_threshold(self) -> float:
        """Get the security threshold for the BB84 protocol.

        Returns:
            Maximum QBER value considered secure

        """
        return self.security_threshold

    def get_basis_reconciliation_rate(self) -> float:
        """Calculate the basis reconciliation rate.

        Returns:
            Fraction of qubits where Alice and Bob used the same basis

        """
        matches = 0
        total = 0

        for i in range(self.num_qubits):
            if self.bob_bases[i] is not None:
                total += 1
                if self.alice_bases[i] == self.bob_bases[i]:
                    matches += 1

        return matches / total if total > 0 else 0

    def get_key_rate(self) -> float:
        """Calculate the key generation rate.

        Returns:
            Fraction of transmitted qubits that result in secure key bits

        """
        if not self.is_complete:
            return 0.0

        return len(self.final_key) / self.num_qubits
