import requests
import os
import json
import base64
from typing import Optional, Dict, Any, Union
from .circuit import QuantumCircuit

class AnuAstra:
    """
    Official Python SDK for the Aṇu Astra Virtual Quantum Processor (v1.5.0)
    
    Enterprise-grade quantum computing API with:
    - QRNG, PQC (ML-KEM + ML-DSA), vQKD (BB84 + E91)
    - Quantum Circuit Execution (OpenQASM 2.0)
    - Immutable Audit Ledger with Ed25519 signatures
    - Credit-based usage tracking
    """
    def __init__(self, api_key: str, base_url: str = 'https://astra.cryptopix.in/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.rstrip('/')
        self.headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }

        # Initialize sub-modules
        self.kyber = self._Kyber(self)
        self.dilithium = self._Dilithium(self)
        self.qrng = self._QRNG(self)
        self.qkd = self._QKD(self)

    def _request(self, method: str, endpoint: str, data: Optional[Dict[str, Any]] = None, params: Optional[Dict[str, Any]] = None):
        url = f"{self.base_url}{endpoint}"
        try:
            response = requests.request(method, url, headers=self.headers, json=data, params=params)
            res_data = response.json()
        except Exception as e:
            raise Exception(f"Aṇu Astra Connection Error: {str(e)}")

        if not response.ok:
            request_id = res_data.get('requestId', '')
            request_id_str = f" (requestId: {request_id})" if request_id else ''
            error_msg = res_data.get('error', response.reason)
            raise Exception(f"Aṇu Astra API Error [{response.status_code}]: {error_msg}{request_id_str}")

        # Enterprise Envelope Unwrapping
        if all(k in res_data for k in ['data', 'compliance', 'metadata']):
            unwrapped = res_data['data']
            unwrapped['status'] = res_data.get('status', 'success')
            unwrapped['_metadata'] = res_data['metadata']
            unwrapped['_compliance'] = res_data['compliance']
            unwrapped['_cost_deducted'] = res_data.get('cost_deducted', 0)
            unwrapped['_verification_url'] = res_data['compliance'].get('verification_url', '')
            return unwrapped

        return res_data

    def get_balance(self) -> Dict[str, Any]:
        """View the remaining API execution credits"""
        return self._request('GET', '/billing/balance')

    def get_usage(self, limit: int = 50) -> Dict[str, Any]:
        """Get your API usage history (most recent first)"""
        return self._request('GET', '/usage/history', params={'limit': limit})

    def verify(self, transaction_id: str) -> Dict[str, Any]:
        """
        Verify any transaction against the immutable audit ledger.
        Returns the Cryptographic Certificate of Execution.
        """
        return self._request('GET', f'/verify/{transaction_id}')

    def execute_quantum(self, qubits_or_circuit: Union[int, QuantumCircuit], circuit_array: Optional[list] = None, shots: int = 1) -> Dict[str, Any]:
        """Submit a quantum payload circuit natively to the engine for physical resolution."""
        if isinstance(qubits_or_circuit, QuantumCircuit):
            qubits = qubits_or_circuit.qubits
            circuit = qubits_or_circuit.operations
        else:
            qubits = qubits_or_circuit
            circuit = circuit_array or []

        return self._request('POST', '/quantum/execute', data={
            'qubits': qubits,
            'circuit': circuit,
            'shots': shots
        })

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

    class _Kyber:
        def __init__(self, outer):
            self.outer = outer

        def keygen(self, level: int = 768):
            """Local Kyber key generation using FIPS 203 if available, otherwise API fallback."""
            try:
                from fips203 import ml_kem
                kem = {512: ml_kem.ML_KEM_512, 768: ml_kem.ML_KEM_768, 1024: ml_kem.ML_KEM_1024}.get(level, ml_kem.ML_KEM_768)
                dk, ek = kem.keygen()
                return {
                    'status': 'success',
                    'algorithm': f'ML-KEM-{level}',
                    'public_key': base64.b64encode(ek).decode('utf-8'),
                    'private_key': base64.b64encode(dk).decode('utf-8')
                }
            except ImportError:
                # Fallback to API if fips203 is not installed or import fails
                return self.outer._request('POST', '/quantum/pqc/kyber/keygen', data={'level': level})

        def encapsulate(self, public_key: str, level: int = 768):
            return self.outer._request('POST', '/quantum/pqc/kyber/encapsulate', data={
                'publicKey': public_key,
                'level': level
            })

        def decapsulate(self, ciphertext: str, private_key: str, level: int = 768):
            """Local Kyber decapsulation."""
            try:
                from fips203 import ml_kem
                kem = {512: ml_kem.ML_KEM_512, 768: ml_kem.ML_KEM_768, 1024: ml_kem.ML_KEM_1024}.get(level, ml_kem.ML_KEM_768)
                shared_secret = kem.decapsulate(base64.b64decode(ciphertext), base64.b64decode(private_key))
                return {
                    'status': 'success',
                    'algorithm': f'ML-KEM-{level}',
                    'shared_secret': base64.b64encode(shared_secret).decode('utf-8')
                }
            except ImportError:
                return self.outer._request('POST', '/quantum/pqc/kyber/decapsulate', data={
                    'ciphertext': ciphertext,
                    'privateKey': private_key,
                    'level': level
                })

    class _Dilithium:
        def __init__(self, outer):
            self.outer = outer

        def keygen(self, level: int = 3):
            return self.outer._request('POST', '/quantum/pqc/dilithium/keygen', data={'level': level})

        def sign(self, message: str, private_key: str, level: int = 3):
            return self.outer._request('POST', '/quantum/pqc/dilithium/sign', data={
                'message': message,
                'privateKey': private_key,
                'level': level
            })

        def verify(self, signature: str, message: str, public_key: str, level: int = 3):
            return self.outer._request('POST', '/quantum/pqc/dilithium/verify', data={
                'signature': signature,
                'message': message,
                'publicKey': public_key,
                'level': level
            })

    class _QRNG:
        def __init__(self, outer):
            self.outer = outer

        def get_random_bits(self, size: int = 32):
            return self.outer._request('POST', '/quantum/qrng', data={'type': 'bits', 'size': size})

        def get_random_bytes(self, size: int = 32):
            return self.outer._request('POST', '/quantum/qrng', data={'type': 'bytes', 'size': size})

        def get_random_int(self, min_val: int = 0, max_val: int = 100):
            return self.outer._request('POST', '/quantum/qrng', data={'type': 'int', 'min': min_val, 'max': max_val})

        def get_random_float(self):
            return self.outer._request('POST', '/quantum/qrng', data={'type': 'float'})

    class _QKD:
        def __init__(self, outer):
            self.outer = outer

        def bb84(self, key_length: int = 256):
            return self.outer._request('POST', '/quantum/qkd/bb84', data={'keyLength': key_length})

        def e91(self, key_length: int = 256):
            return self.outer._request('POST', '/quantum/qkd/e91', data={'keyLength': key_length})
