import requests
from typing import List, Dict, Any, Optional, Union
from .circuit import QuantumCircuit

class KyberPQC:
    def __init__(self, request_fn):
        self._request = request_fn

    def keygen(self, level: int = 768) -> Dict[str, Any]:
        """
        DEPRECATED: Server-side generation transmits private keys over HTTPS. 
        Use keygen_local() for true Zero-Knowledge Key Encapsulation.
        """
        return self._request('/pqc/kyber/keygen', 'POST', {'level': level})

    def keygen_local(self, level: int = 768) -> Dict[str, Any]:
        """
        Generates a Kyber keypair LOCALLY in the SDK to ensure zero-knowledge transmission.
        The private key never leaves the client device.
        """
        import pqcrypto.kem.kyber512 as kyber512
        import pqcrypto.kem.kyber768 as kyber768
        import pqcrypto.kem.kyber1024 as kyber1024
        import base64

        kem = kyber512 if level == 512 else kyber1024 if level == 1024 else kyber768
        public_key_bytes, secret_key_bytes = kem.generate_keypair()
        
        return {
            'status': 'success',
            'algorithm': f'ML-KEM-{level}',
            'public_key': base64.b64encode(public_key_bytes).decode('utf-8'),
            'private_key': base64.b64encode(secret_key_bytes).decode('utf-8')
        }

    def encapsulate(self, public_key: str, level: int = 768) -> Dict[str, Any]:
        return self._request('/pqc/kyber/encapsulate', 'POST', {'publicKey': public_key, 'level': level})

    def decapsulate(self, ciphertext: str, private_key: str, level: int = 768) -> Dict[str, Any]:
        """
        DEPRECATED: Server-side decapsulation requires transmitting private keys over HTTPS.
        Use decapsulate_local() for true Zero-Knowledge Key Encapsulation.
        """
        return self._request('/pqc/kyber/decapsulate', 'POST', {'ciphertext': ciphertext, 'privateKey': private_key, 'level': level})
        
    def decapsulate_local(self, ciphertext: str, private_key: str, level: int = 768) -> Dict[str, Any]:
        """
        Decapsulates a ciphertext LOCALLY in the SDK using the local private key.
        """
        import pqcrypto.kem.kyber512 as kyber512
        import pqcrypto.kem.kyber768 as kyber768
        import pqcrypto.kem.kyber1024 as kyber1024
        import base64

        kem = kyber512 if level == 512 else kyber1024 if level == 1024 else kyber768
        c_bytes = base64.b64decode(ciphertext)
        p_bytes = base64.b64decode(private_key)
        
        shared_secret = kem.decrypt(c_bytes, p_bytes)
        
        return {
            'status': 'success',
            'algorithm': f'ML-KEM-{level}',
            'shared_secret': base64.b64encode(shared_secret).decode('utf-8')
        }

class DilithiumPQC:
    def __init__(self, request_fn):
        self._request = request_fn

    def keygen(self, level: int = 3) -> Dict[str, Any]:
        return self._request('/pqc/dilithium/keygen', 'POST', {'level': level})

    def sign(self, message: str, private_key: str, level: int = 3) -> Dict[str, Any]:
        return self._request('/pqc/dilithium/sign', 'POST', {'message': message, 'privateKey': private_key, 'level': level})

    def verify(self, signature: str, message: str, public_key: str, level: int = 3) -> Dict[str, Any]:
        return self._request('/pqc/dilithium/verify', 'POST', {'signature': signature, 'message': message, 'publicKey': public_key, 'level': level})


class Qrng:
    """
    Quantum Random Number Generator (QRNG) native to the Aṇu Astra VQP.
    Uses pure quantum entropy via Hadamard-superposition.
    """
    def __init__(self, request_fn):
        self._request = request_fn

    def get_random_bits(self, size: int = 32) -> Dict[str, Any]:
        """Generate a random bit string of the specified size."""
        return self._request('/quantum/qrng', 'POST', {'type': 'bits', 'size': size})

    def get_random_bytes(self, length: int = 32) -> Dict[str, Any]:
        """Generate random bytes (returned as a hex string)."""
        return self._request('/quantum/qrng', 'POST', {'type': 'bytes', 'size': length})

    def get_random_int(self, min: int = 0, max: int = 100) -> Dict[str, Any]:
        """Generate a random integer within the range [min, max]."""
        return self._request('/quantum/qrng', 'POST', {'type': 'int', 'min': min, 'max': max})

    def get_random_float(self) -> Dict[str, Any]:
        """Generate a random float in the range [0.0, 1.0)."""
        return self._request('/quantum/qrng', 'POST', {'type': 'float'})


class Client:
    """
    Official Python SDK for the Aṇu Astra Virtual Quantum Processor
    """
    def __init__(self, api_key: str, base_url: str = 'https://api.anuastra.com/api/v1'):
        if not api_key:
            raise ValueError("Aṇu Astra API Key is required")
        self.api_key = api_key
        self.base_url = base_url
        self.kyber = KyberPQC(self._request)
        self.dilithium = DilithiumPQC(self._request)
        self.qrng = Qrng(self._request)

    def _request(self, endpoint: str, method: str, data: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        url = f"{self.base_url}{endpoint}"
        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }
        
        response = requests.request(method, url, headers=headers, json=data)
        
        try:
            result = response.json()
        except ValueError:
            response.raise_for_status()
            return {}

        if not response.ok:
            details = f" - {result.get('details', '')}" if 'details' in result else ''
            raise Exception(f"Aṇu Astra API Error: {result.get('error', response.reason)}{details}")
            
        return result

    def get_balance(self) -> Dict[str, Any]:
        return self._request('/billing/balance', 'GET')

    def execute_quantum(self, qubits: Union[int, QuantumCircuit], circuit: Optional[List[Dict[str, Any]]] = None, shots: int = 1) -> Dict[str, Any]:
        """
        Submit a quantum payload circuit natively to the engine for physical resolution.
        """
        if isinstance(qubits, QuantumCircuit):
            _q = qubits.qubits
            _circ = qubits.operations
        else:
            _q = qubits
            _circ = circuit or []
            
        return self._request('/quantum/execute', 'POST', {
            'qubits': _q,
            'circuit': _circ,
            'shots': shots
        })

    def execute(self, circuit: QuantumCircuit, shots: int = 1) -> Dict[str, Any]:
        """
        Execute a QuantumCircuit object gracefully.
        """
        return self.execute_quantum(circuit, shots=shots)

    def execute_qasm(self, qasm_string: str, shots: int = 1) -> Dict[str, Any]:
        """
        Submit a raw OpenQASM 2.0 string for server-side compilation and execution.
        """
        return self._request('/quantum/execute-qasm', 'POST', {
            'qasm': qasm_string,
            'shots': shots
        })
