import random
from typing import List, Dict, Any
from ..system import QuantumSystem

class QKDSession:
    def __init__(self, alice: QuantumSystem, bob: QuantumSystem):
        self.alice = alice
        self.bob = bob
        self.key_buffer: List[int] = []
        self.is_secure = True

    def negotiate_key(self, bit_length: int = 128) -> Dict[str, Any]:
        self.key_buffer = []
        attempts = 0
        max_attempts = bit_length * 10

        while len(self.key_buffer) < bit_length and attempts < max_attempts:
            attempts += 1
            
            # 1. Preparation (Alice)
            bit = 1 if random.random() > 0.5 else 0
            basis = 'RECT' if random.random() > 0.5 else 'DIAG'
            
            self.alice.chi.normalize()

            if basis == 'RECT':
                self.alice.chi.vector = [1.0, 0.0, 0.0] if bit == 0 else [0.0, 1.0, 0.0]
            else:
                bell = 0.70710678
                self.alice.chi.vector = [bell, bell, 0.0] if bit == 0 else [bell, -bell, 0.0]

            # 2. Transmission (Noise Check)
            noise = self.alice.chi.calculate_shear()
            if noise > 0.1:
                self.is_secure = False
                continue

            # 3. Measurement (Bob)
            bob_basis = 'RECT' if random.random() > 0.5 else 'DIAG'
            
            # Sifting
            if basis == bob_basis:
                self.key_buffer.append(bit)
        
        return {
            "success": len(self.key_buffer) == bit_length,
            "key": "".join(map(str, self.key_buffer)),
            "is_secure": self.is_secure
        }
