"""
Synor ZK SDK Client for Python

Zero-Knowledge proof systems for ZK-Rollups and privacy.
"""

import asyncio
import base64
from typing import Optional, List, Dict, Any

import httpx

from .types import (
    ProofSystem,
    CircuitType,
    Proof,
    VerificationKey,
    ProvingKey,
    CircuitConfig,
    AccountState,
    Transfer,
    Deposit,
    Withdrawal,
    Batch,
    RollupStats,
    ProofRequest,
    VerificationResult,
    MerkleProof,
    TrustedSetup,
    CeremonyContribution,
    ZkConfig,
    ZkError,
)


class SynorZk:
    """
    Synor ZK SDK Client

    Provides zero-knowledge proof generation and verification for ZK-Rollups.

    Example:
        ```python
        async with SynorZk(ZkConfig(api_key="your-api-key")) as zk:
            # Generate a proof
            proof = await zk.proofs.generate(ProofRequest(
                circuit_type=CircuitType.TRANSFER,
                public_inputs=[...],
                private_inputs=[...],
            ))

            # Verify the proof
            result = await zk.proofs.verify(proof)
            print(f"Valid: {result.valid}")
        ```
    """

    def __init__(self, config: ZkConfig):
        self.config = config
        self._client: Optional[httpx.AsyncClient] = None
        self._closed = False

        self.proofs = ProofsClient(self)
        self.circuits = CircuitsClient(self)
        self.rollup = RollupClient(self)
        self.state = StateClient(self)
        self.ceremony = CeremonyClient(self)

    async def __aenter__(self) -> "SynorZk":
        self._client = httpx.AsyncClient(timeout=self.config.timeout)
        return self

    async def __aexit__(self, *args) -> None:
        await self.close()

    @property
    def default_proof_system(self) -> ProofSystem:
        """Get the default proof system"""
        return self.config.default_proof_system

    async def health_check(self) -> bool:
        """Health check"""
        try:
            result = await self._get("/health")
            return result.get("status") == "healthy"
        except Exception:
            return False

    async def get_info(self) -> Dict[str, Any]:
        """Get service info"""
        return await self._get("/info")

    async def close(self) -> None:
        """Close the client"""
        self._closed = True
        if self._client:
            await self._client.aclose()
            self._client = None

    @property
    def is_closed(self) -> bool:
        """Check if client is closed"""
        return self._closed

    async def _get(self, path: str) -> Dict[str, Any]:
        """Internal GET request"""
        return await self._request("GET", path)

    async def _post(self, path: str, body: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Internal POST request"""
        return await self._request("POST", path, body)

    async def _delete(self, path: str) -> Dict[str, Any]:
        """Internal DELETE request"""
        return await self._request("DELETE", path)

    async def _request(
        self,
        method: str,
        path: str,
        body: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Internal HTTP request with retries"""
        if self._closed:
            raise ZkError("Client has been closed", code="CLIENT_CLOSED")

        if not self._client:
            self._client = httpx.AsyncClient(timeout=self.config.timeout)

        url = f"{self.config.endpoint}{path}"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.config.api_key}",
            "X-SDK-Version": "python/0.1.0",
        }

        last_error: Optional[Exception] = None

        for attempt in range(self.config.retries + 1):
            try:
                if method == "GET":
                    response = await self._client.get(url, headers=headers)
                elif method == "POST":
                    response = await self._client.post(url, headers=headers, json=body)
                elif method == "DELETE":
                    response = await self._client.delete(url, headers=headers)
                else:
                    raise ValueError(f"Unknown method: {method}")

                if response.status_code >= 400:
                    error = response.json() if response.content else {}
                    raise ZkError(
                        error.get("message", f"HTTP {response.status_code}"),
                        code=error.get("code"),
                        status=response.status_code,
                    )

                return response.json()

            except ZkError:
                raise
            except Exception as e:
                last_error = e
                if attempt < self.config.retries:
                    await asyncio.sleep(0.1 * (2 ** attempt))

        raise ZkError(
            str(last_error) if last_error else "Request failed",
            code="NETWORK_ERROR",
        )


class ProofsClient:
    """Proofs sub-client"""

    def __init__(self, zk: SynorZk):
        self._zk = zk

    async def generate(self, request: ProofRequest) -> Proof:
        """Generate a zero-knowledge proof"""
        body = request.to_dict()
        if not body.get("system"):
            body["system"] = self._zk.default_proof_system.value
        result = await self._zk._post("/proofs/generate", body)
        return Proof.from_dict(result)

    async def verify(self, proof: Proof) -> VerificationResult:
        """Verify a proof"""
        result = await self._zk._post("/proofs/verify", {
            "proof_id": proof.id,
            "data": proof.data,
            "public_inputs": proof.public_inputs,
            "system": proof.system.value,
            "circuit_type": proof.circuit_type.value,
        })
        return VerificationResult.from_dict(result)

    async def get(self, proof_id: str) -> Proof:
        """Get a proof by ID"""
        result = await self._zk._get(f"/proofs/{proof_id}")
        return Proof.from_dict(result)

    async def list(
        self,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
    ) -> List[Proof]:
        """List recent proofs"""
        params = []
        if limit:
            params.append(f"limit={limit}")
        if offset:
            params.append(f"offset={offset}")
        query = f"?{'&'.join(params)}" if params else ""
        result = await self._zk._get(f"/proofs{query}")
        return [Proof.from_dict(p) for p in result.get("proofs", [])]

    async def serialize(self, proof: Proof) -> bytes:
        """Serialize a proof to bytes"""
        result = await self._zk._post("/proofs/serialize", {"proof_id": proof.id})
        return base64.b64decode(result["data"])

    async def deserialize(self, data: bytes, system: ProofSystem) -> Proof:
        """Deserialize bytes to a proof"""
        result = await self._zk._post("/proofs/deserialize", {
            "data": base64.b64encode(data).decode(),
            "system": system.value,
        })
        return Proof.from_dict(result)


class CircuitsClient:
    """Circuits sub-client"""

    def __init__(self, zk: SynorZk):
        self._zk = zk

    async def get_verification_key(
        self,
        circuit_type: CircuitType,
        system: Optional[ProofSystem] = None,
    ) -> VerificationKey:
        """Get verification key for a circuit"""
        sys = system or self._zk.default_proof_system
        result = await self._zk._get(f"/circuits/{circuit_type.value}/vk?system={sys.value}")
        return VerificationKey.from_dict(result)

    async def get_proving_key(
        self,
        circuit_type: CircuitType,
        system: Optional[ProofSystem] = None,
    ) -> ProvingKey:
        """Get proving key for a circuit"""
        sys = system or self._zk.default_proof_system
        result = await self._zk._get(f"/circuits/{circuit_type.value}/pk?system={sys.value}")
        return ProvingKey.from_dict(result)

    async def list(self) -> List[CircuitConfig]:
        """List available circuits"""
        result = await self._zk._get("/circuits")
        return [CircuitConfig.from_dict(c) for c in result.get("circuits", [])]

    async def get_config(self, circuit_type: CircuitType) -> CircuitConfig:
        """Get circuit configuration"""
        result = await self._zk._get(f"/circuits/{circuit_type.value}/config")
        return CircuitConfig.from_dict(result)

    async def compile(self, config: CircuitConfig) -> str:
        """Compile a custom circuit, returns circuit ID"""
        result = await self._zk._post("/circuits/compile", {
            "type": config.type.value,
            "max_batch_size": config.max_batch_size,
            "tree_depth": config.tree_depth,
            "verify_signatures": config.verify_signatures,
            "custom_constraints": config.custom_constraints,
        })
        return result["circuit_id"]


class RollupClient:
    """Rollup sub-client"""

    def __init__(self, zk: SynorZk):
        self._zk = zk

    async def get_stats(self) -> RollupStats:
        """Get rollup statistics"""
        result = await self._zk._get("/rollup/stats")
        return RollupStats.from_dict(result)

    async def get_current_batch(self) -> Batch:
        """Get current batch"""
        result = await self._zk._get("/rollup/batch/current")
        return Batch.from_dict(result)

    async def get_batch(self, batch_number: int) -> Batch:
        """Get batch by number"""
        result = await self._zk._get(f"/rollup/batch/{batch_number}")
        return Batch.from_dict(result)

    async def submit_transfer(self, transfer: Transfer) -> str:
        """Submit a transfer, returns transaction ID"""
        result = await self._zk._post("/rollup/transfer", transfer.to_dict())
        return result["tx_id"]

    async def submit_deposit(self, deposit: Deposit) -> str:
        """Submit a deposit, returns transaction ID"""
        result = await self._zk._post("/rollup/deposit", deposit.to_dict())
        return result["tx_id"]

    async def submit_withdrawal(self, withdrawal: Withdrawal) -> str:
        """Submit a withdrawal, returns transaction ID"""
        result = await self._zk._post("/rollup/withdraw", withdrawal.to_dict())
        return result["tx_id"]

    async def finalize_batch(self) -> Batch:
        """Finalize current batch (generates proof and submits to L1)"""
        result = await self._zk._post("/rollup/batch/finalize", {})
        return Batch.from_dict(result)

    async def get_pending_transactions(self) -> List[Transfer]:
        """Get pending transactions"""
        result = await self._zk._get("/rollup/pending")
        return [Transfer.from_dict(t) for t in result.get("transactions", [])]


class StateClient:
    """State sub-client"""

    def __init__(self, zk: SynorZk):
        self._zk = zk

    async def get_root(self) -> str:
        """Get current state root"""
        result = await self._zk._get("/state/root")
        return result["root"]

    async def get_account(self, address: str) -> AccountState:
        """Get account state"""
        result = await self._zk._get(f"/state/account/{address}")
        return AccountState.from_dict(result)

    async def get_merkle_proof(self, address: str) -> MerkleProof:
        """Get merkle proof for account inclusion"""
        result = await self._zk._get(f"/state/proof/{address}")
        return MerkleProof.from_dict(result)

    async def verify_merkle_proof(self, proof: MerkleProof) -> bool:
        """Verify merkle proof"""
        result = await self._zk._post("/state/proof/verify", proof.to_dict())
        return result["valid"]

    async def get_state_at_batch(self, batch_number: int) -> Dict[str, Any]:
        """Get state at specific batch"""
        return await self._zk._get(f"/state/batch/{batch_number}")


class CeremonyClient:
    """Ceremony sub-client (trusted setup)"""

    def __init__(self, zk: SynorZk):
        self._zk = zk

    async def get_status(
        self,
        circuit_type: CircuitType,
        system: Optional[ProofSystem] = None,
    ) -> TrustedSetup:
        """Get ceremony status"""
        sys = system or self._zk.default_proof_system
        result = await self._zk._get(f"/ceremony/{circuit_type.value}?system={sys.value}")
        return TrustedSetup.from_dict(result)

    async def contribute(
        self,
        circuit_type: CircuitType,
        entropy: bytes,
        system: Optional[ProofSystem] = None,
    ) -> CeremonyContribution:
        """Contribute to ceremony"""
        sys = system or self._zk.default_proof_system
        result = await self._zk._post("/ceremony/contribute", {
            "circuit_type": circuit_type.value,
            "entropy": base64.b64encode(entropy).decode(),
            "system": sys.value,
        })
        return CeremonyContribution.from_dict(result)

    async def verify_contribution(
        self,
        circuit_type: CircuitType,
        contribution_hash: str,
    ) -> bool:
        """Verify a contribution"""
        result = await self._zk._post("/ceremony/verify", {
            "circuit_type": circuit_type.value,
            "contribution_hash": contribution_hash,
        })
        return result["valid"]

    async def list_contributions(
        self,
        circuit_type: CircuitType,
    ) -> List[CeremonyContribution]:
        """List contributions"""
        result = await self._zk._get(f"/ceremony/{circuit_type.value}/contributions")
        return [CeremonyContribution.from_dict(c) for c in result.get("contributions", [])]
