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]:
        return self._request('/pqc/kyber/keygen', 'POST', {'level': level})

    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]:
        return self._request('/pqc/kyber/decapsulate', 'POST', {'ciphertext': ciphertext, 'privateKey': private_key, 'level': level})


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 Client:
    """
    Official Python SDK for the AnuAstra Virtual Quantum Processor.
    Supports native circuit building, OpenQASM 2.0, and Post-Quantum Cryptography (Kyber/Dilithium).
    """
    def __init__(self, api_key: str, base_url: str = 'https://astra.cryptopix.in/api/v1'):
        if not api_key:
            raise ValueError("AnuAstra API Key is required")
        self.api_key = api_key
        self.base_url = base_url
        self.kyber = KyberPQC(self._request)
        self.dilithium = DilithiumPQC(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"AnuAstra 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.
        (Legacy method. Use `.execute()` for preferred Object-Oriented flow).
        
        Args:
            qubits: Number of qubits or a QuantumCircuit object.
            circuit: List of gate operations (if not using QuantumCircuit).
            shots: Number of times to measure the circuit. Returns 'state_integers' array.
        """
        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.
        Returns a result dictionary containing 'state_integers' and 'state_binaries' for multi-shot analysis.
        """
        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.
        Returns results for the requested number of shots.
        """
        return self._request('/quantum/execute-qasm', 'POST', {
            'qasm': qasm_string,
            'shots': shots
        })
