CIA/e-voting-system/backend/blockchain_client.py

"""
BlockchainClient for communicating with PoA validator nodes.

This client submits votes to the distributed PoA blockchain network
and queries the state of votes on the blockchain.
"""

import logging
import httpx
import json
from typing import Optional, Dict, Any, List
from dataclasses import dataclass
from enum import Enum
import asyncio

logger = logging.getLogger(__name__)


class ValidatorStatus(str, Enum):
    """Status of a validator node"""
    HEALTHY = "healthy"
    DEGRADED = "degraded"
    UNREACHABLE = "unreachable"


@dataclass
class ValidatorNode:
    """Represents a PoA validator node"""
    node_id: str
    rpc_url: str  # JSON-RPC endpoint
    p2p_url: str  # P2P networking endpoint
    status: ValidatorStatus = ValidatorStatus.UNREACHABLE

    @property
    def health_check_url(self) -> str:
        """Health check endpoint"""
        return f"{self.rpc_url}/health"


class BlockchainClient:
    """
    Client for PoA blockchain network.

    Features:
    - Load balancing across multiple validators
    - Health monitoring
    - Automatic failover
    - Vote submission and confirmation tracking
    """

    # Default validator configuration
    # Use Docker service names for internal container communication
    # For external access (outside Docker), use localhost:PORT
    DEFAULT_VALIDATORS = [
        ValidatorNode(
            node_id="validator-1",
            rpc_url="http://validator-1:8001",
            p2p_url="http://validator-1:30303"
        ),
        ValidatorNode(
            node_id="validator-2",
            rpc_url="http://validator-2:8002",
            p2p_url="http://validator-2:30304"
        ),
        ValidatorNode(
            node_id="validator-3",
            rpc_url="http://validator-3:8003",
            p2p_url="http://validator-3:30305"
        ),
    ]

    def __init__(self, validators: Optional[List[ValidatorNode]] = None, timeout: float = 5.0):
        """
        Initialize blockchain client.

        Args:
            validators: List of validator nodes (uses defaults if None)
            timeout: HTTP request timeout in seconds
        """
        self.validators = validators or self.DEFAULT_VALIDATORS
        self.timeout = timeout
        self.healthy_validators: List[ValidatorNode] = []
        self._client: Optional[httpx.AsyncClient] = None

    async def __aenter__(self):
        """Async context manager entry"""
        logger.info("[BlockchainClient.__aenter__] Creating AsyncClient")
        self._client = httpx.AsyncClient(timeout=self.timeout)
        logger.info("[BlockchainClient.__aenter__] Refreshing validator status")
        await self.refresh_validator_status()
        logger.info(f"[BlockchainClient.__aenter__] Ready with {len(self.healthy_validators)} healthy validators")
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit"""
        if self._client:
            await self._client.aclose()

    async def refresh_validator_status(self) -> None:
        """
        Check health of all validators.

        Updates the list of healthy validators for load balancing.
        """
        if not self._client:
            self._client = httpx.AsyncClient(timeout=self.timeout)

        tasks = [self._check_validator_health(v) for v in self.validators]
        await asyncio.gather(*tasks, return_exceptions=True)

        self.healthy_validators = [
            v for v in self.validators
            if v.status == ValidatorStatus.HEALTHY
        ]

        logger.info(
            f"Validator health check: {len(self.healthy_validators)}/{len(self.validators)} healthy"
        )

    async def _check_validator_health(self, validator: ValidatorNode) -> None:
        """Check if a validator is healthy"""
        try:
            if not self._client:
                return

            response = await self._client.get(
                validator.health_check_url,
                timeout=self.timeout
            )

            if response.status_code == 200:
                validator.status = ValidatorStatus.HEALTHY
                logger.debug(f"✓ {validator.node_id} is healthy")
            else:
                validator.status = ValidatorStatus.DEGRADED
                logger.warning(f"⚠ {validator.node_id} returned status {response.status_code}")
        except Exception as e:
            validator.status = ValidatorStatus.UNREACHABLE
            logger.warning(f"✗ {validator.node_id} is unreachable: {e}")

    def _get_healthy_validator(self) -> Optional[ValidatorNode]:
        """
        Get a healthy validator for the next request.
        Uses round-robin for load balancing.
        """
        if not self.healthy_validators:
            logger.error("No healthy validators available!")
            return None

        # Simple round-robin: return first healthy validator
        # In production, implement proper round-robin state management
        return self.healthy_validators[0]

    async def submit_vote(
        self,
        voter_id: str,
        election_id: int,
        encrypted_vote: str,
        transaction_id: Optional[str] = None,
        ballot_hash: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Submit a vote to ALL PoA validators simultaneously.

        This ensures every validator receives the transaction directly,
        guaranteeing it will be included in the next block.

        Args:
            voter_id: Voter identifier
            election_id: Election ID
            encrypted_vote: Encrypted vote data
            transaction_id: Optional transaction ID (generated if not provided)
            ballot_hash: Optional ballot hash for verification

        Returns:
            Transaction receipt with block hash and index

        Raises:
            Exception: If all validators are unreachable
        """
        logger.info(f"[BlockchainClient.submit_vote] CALLED with voter_id={voter_id}, election_id={election_id}")
        logger.info(f"[BlockchainClient.submit_vote] healthy_validators count: {len(self.healthy_validators)}")

        if not self.healthy_validators:
            logger.error("[BlockchainClient.submit_vote] No healthy validators available!")
            raise Exception("No healthy validators available")

        # Generate transaction ID if not provided
        if not transaction_id:
            import uuid
            transaction_id = f"tx-{uuid.uuid4().hex[:12]}"

        # Generate ballot hash if not provided
        if not ballot_hash:
            import hashlib
            ballot_hash = hashlib.sha256(f"{voter_id}{election_id}{encrypted_vote}".encode()).hexdigest()

        import time

        # Create transaction data as JSON
        tx_data = {
            "voter_id": str(voter_id),
            "election_id": int(election_id),
            "encrypted_vote": str(encrypted_vote),
            "ballot_hash": str(ballot_hash),
            "timestamp": int(time.time())
        }

        # Encode transaction data as hex string with 0x prefix
        import json
        tx_json = json.dumps(tx_data)
        data_hex = "0x" + tx_json.encode().hex()

        # Prepare JSON-RPC request
        rpc_request = {
            "jsonrpc": "2.0",
            "method": "eth_sendTransaction",
            "params": [{
                "from": voter_id,
                "to": f"election-{election_id}",
                "data": data_hex,
                "gas": "0x5208"
            }],
            "id": transaction_id
        }

        # Submit to ALL healthy validators simultaneously
        logger.info(f"[BlockchainClient.submit_vote] Submitting to {len(self.healthy_validators)} validators")

        results = {}
        if not self._client:
            logger.error("[BlockchainClient.submit_vote] AsyncClient not initialized!")
            raise Exception("AsyncClient not initialized")

        for validator in self.healthy_validators:
            try:
                logger.info(f"[BlockchainClient.submit_vote] Submitting to {validator.node_id} ({validator.rpc_url}/rpc)")
                response = await self._client.post(
                    f"{validator.rpc_url}/rpc",
                    json=rpc_request,
                    timeout=self.timeout
                )
                logger.info(f"[BlockchainClient.submit_vote] Response from {validator.node_id}: status={response.status_code}")

                response.raise_for_status()
                result = response.json()

                # Check for JSON-RPC errors
                if "error" in result:
                    logger.error(f"RPC error from {validator.node_id}: {result['error']}")
                    results[validator.node_id] = f"RPC error: {result['error']}"
                else:
                    logger.info(f"✓ Vote accepted by {validator.node_id}: {result.get('result')}")
                    results[validator.node_id] = result.get("result")

            except Exception as e:
                logger.warning(f"Failed to submit to {validator.node_id}: {e}")
                results[validator.node_id] = str(e)

        # Check if at least one validator accepted the vote
        successful = [v for v in results.values() if not str(v).startswith(("RPC error", "Failed"))]
        if successful:
            logger.info(f"✓ Vote submitted successfully to {len(successful)} validators: {transaction_id}")
            return {
                "transaction_id": transaction_id,
                "block_hash": successful[0] if successful else None,
                "validator": self.healthy_validators[0].node_id,
                "status": "pending"
            }
        else:
            logger.error(f"Failed to submit vote to any validator")
            raise Exception(f"All validator submissions failed: {results}")

    async def get_transaction_receipt(
        self,
        transaction_id: str,
        election_id: int
    ) -> Optional[Dict[str, Any]]:
        """
        Get the receipt for a submitted vote.

        Args:
            transaction_id: Transaction ID returned from submit_vote
            election_id: Election ID

        Returns:
            Transaction receipt with confirmation status and block info
        """
        validator = self._get_healthy_validator()
        if not validator:
            return None

        rpc_request = {
            "jsonrpc": "2.0",
            "method": "eth_getTransactionReceipt",
            "params": [transaction_id],
            "id": transaction_id
        }

        try:
            if not self._client:
                raise Exception("AsyncClient not initialized")

            response = await self._client.post(
                f"{validator.rpc_url}/rpc",
                json=rpc_request,
                timeout=self.timeout
            )

            result = response.json()

            if "error" in result:
                logger.warning(f"RPC error: {result['error']}")
                return None

            receipt = result.get("result")
            if receipt:
                logger.debug(f"✓ Got receipt for {transaction_id}: block {receipt.get('blockNumber')}")

            return receipt

        except Exception as e:
            logger.warning(f"Failed to get receipt for {transaction_id}: {e}")
            return None

    async def get_vote_confirmation_status(
        self,
        transaction_id: str,
        election_id: int
    ) -> Dict[str, Any]:
        """
        Check if a vote has been confirmed on the blockchain.

        Args:
            transaction_id: Transaction ID
            election_id: Election ID

        Returns:
            Status information including block number and finality
        """
        receipt = await self.get_transaction_receipt(transaction_id, election_id)

        if receipt is None:
            return {
                "status": "pending",
                "confirmed": False,
                "transaction_id": transaction_id
            }

        return {
            "status": "confirmed",
            "confirmed": True,
            "transaction_id": transaction_id,
            "block_number": receipt.get("blockNumber"),
            "block_hash": receipt.get("blockHash"),
            "gas_used": receipt.get("gasUsed")
        }

    async def get_blockchain_state(self, election_id: int) -> Optional[Dict[str, Any]]:
        """
        Get the current state of the blockchain for an election.

        Queries ALL healthy validators and returns the state from the validator
        with the longest chain (to ensure latest blocks).

        Args:
            election_id: Election ID

        Returns:
            Blockchain state with block count and verification status
        """
        if not self.healthy_validators:
            return None

        if not self._client:
            raise Exception("AsyncClient not initialized")

        # Query all validators and get the one with longest chain
        best_state = None
        best_block_count = 0

        for validator in self.healthy_validators:
            try:
                logger.debug(f"Querying blockchain state from {validator.node_id}")
                response = await self._client.get(
                    f"{validator.rpc_url}/blockchain",
                    params={"election_id": election_id},
                    timeout=self.timeout
                )

                response.raise_for_status()
                state = response.json()

                # Get block count from this validator
                block_count = len(state.get("blocks", []))
                logger.debug(f"{validator.node_id} has {block_count} blocks")

                # Keep the state with the most blocks (longest chain)
                if block_count > best_block_count:
                    best_state = state
                    best_block_count = block_count
                    logger.info(f"Using state from {validator.node_id} ({block_count} blocks)")

            except Exception as e:
                logger.warning(f"Failed to get blockchain state from {validator.node_id}: {e}")
                continue

        return best_state if best_state else None

    async def verify_blockchain_integrity(self, election_id: int) -> bool:
        """
        Verify that the blockchain for an election is valid and unmodified.

        Args:
            election_id: Election ID

        Returns:
            True if blockchain is valid, False otherwise
        """
        state = await self.get_blockchain_state(election_id)

        if state is None:
            return False

        verification = state.get("verification", {})
        is_valid = verification.get("chain_valid", False)

        if is_valid:
            logger.info(f"✓ Blockchain for election {election_id} is valid")
        else:
            logger.error(f"✗ Blockchain for election {election_id} is INVALID")

        return is_valid

    async def get_election_results(self, election_id: int) -> Optional[Dict[str, Any]]:
        """
        Get the current vote counts for an election from the blockchain.

        Args:
            election_id: Election ID

        Returns:
            Vote counts by candidate and verification status
        """
        validator = self._get_healthy_validator()
        if not validator:
            return None

        try:
            if not self._client:
                raise Exception("AsyncClient not initialized")

            # Query results endpoint on validator
            response = await self._client.get(
                f"{validator.rpc_url}/results",
                params={"election_id": election_id},
                timeout=self.timeout
            )

            response.raise_for_status()
            return response.json()

        except Exception as e:
            logger.warning(f"Failed to get election results: {e}")
            return None

    async def wait_for_confirmation(
        self,
        transaction_id: str,
        election_id: int,
        max_wait_seconds: int = 30,
        poll_interval_seconds: float = 1.0
    ) -> bool:
        """
        Wait for a vote to be confirmed on the blockchain.

        Args:
            transaction_id: Transaction ID
            election_id: Election ID
            max_wait_seconds: Maximum time to wait in seconds
            poll_interval_seconds: Time between status checks

        Returns:
            True if vote was confirmed, False if timeout
        """
        import time

        start_time = time.time()
        while time.time() - start_time < max_wait_seconds:
            status = await self.get_vote_confirmation_status(transaction_id, election_id)

            if status.get("confirmed"):
                logger.info(f"✓ Vote confirmed: {transaction_id}")
                return True

            logger.debug(f"Waiting for confirmation... ({status['status']})")
            await asyncio.sleep(poll_interval_seconds)

        logger.warning(f"Confirmation timeout for {transaction_id}")
        return False


# Singleton instance for use throughout the backend
_blockchain_client: Optional[BlockchainClient] = None


async def get_blockchain_client() -> BlockchainClient:
    """
    Get or create the global blockchain client instance.

    Returns:
        BlockchainClient instance
    """
    global _blockchain_client

    if _blockchain_client is None:
        _blockchain_client = BlockchainClient()
        await _blockchain_client.refresh_validator_status()

    return _blockchain_client


def get_blockchain_client_sync() -> BlockchainClient:
    """
    Get the blockchain client (for sync contexts).

    Note: This returns the client without initializing it.
    Use with caution in async contexts.

    Returns:
        BlockchainClient instance
    """
    global _blockchain_client

    if _blockchain_client is None:
        _blockchain_client = BlockchainClient()

    return _blockchain_client