"""
Proof of Authority consensus mechanism for HieraChain Ledger.

This module implements a Proof of Authority (PoA) consensus mechanism suitable
for the HieraChain Ledger where specific authorities (Main Chain,
Sub-Chains) have designated roles and permissions for block creation.
"""

import time
import logging
from typing import Any

from hierachain.consensus.base_consensus import (
    BaseConsensus, _verify_block_zk_proof
)
from hierachain.core.block import Block
from hierachain.security.security_utils import KeyPair

logger = logging.getLogger(__name__)


class ProofOfAuthority(BaseConsensus):
    """
    Proof of Authority consensus mechanism.
    This consensus mechanism is ideal for the HieraChain Ledger
    where:
    - Main Chain acts as the root authority
    - Sub-Chains are authorized domain-specific authorities
    - Block creation is controlled by authorized entities
    - No energy-intensive mining (suitable for business applications)
    """

    def __init__(self, name: str = "ProofOfAuthority"):
        """
        Initialize Proof of Authority consensus.
        Args:
            name: Name of the consensus mechanism
        """
        super().__init__(name)
        self.authorities: set[str] = set()
        self.authority_metadata: dict[str, dict[str, Any]] = {}
        self.block_interval: float = 10.0  # Default 10 seconds between blocks
        self.config = {
            "block_interval": self.block_interval,
            "require_authority_signature": True,
            "max_authorities": 100
        }

    # Forbidden event types for business logic
    FORBIDDEN_EVENT_TYPES = {"transaction", "mining", "coin_transfer", "wallet_update"}

    def add_authority(
        self,
        authority_id: str,
        metadata: dict[str, Any] | None = None
    ) -> bool:
        """
        Add a new authority to the consensus mechanism.

        Args:
            authority_id: Unique identifier for the authority
            metadata: Additional metadata about the authority
        Returns:
            True if authority was added successfully, False otherwise
        """
        if len(self.authorities) >= self.config["max_authorities"]:
            return False

        self.authorities.add(authority_id)
        self.authority_metadata[authority_id] = metadata or {}
        return True

    def remove_authority(self, authority_id: str) -> bool:
        """
        Remove an authority from the consensus mechanism.

        Args:
            authority_id: Authority identifier to remove

        Returns:
            True if authority was removed successfully, False otherwise
        """
        if authority_id in self.authorities:
            self.authorities.remove(authority_id)
            self.authority_metadata.pop(authority_id, None)
            return True
        return False

    def is_authority(self, authority_id: str) -> bool:
        """
        Check if an entity is an authorized authority.

        Args:
            authority_id: Authority identifier to check

        Returns:
            True if entity is an authority, False otherwise
        """
        return authority_id in self.authorities

    def can_create_block(self, authority_id: str | None = None) -> bool:
        """Check if a block can be created by the given authority."""
        if authority_id is None:
            return False
        return self.is_authority(authority_id)

    def _check_block_timing(self, block: Block, previous_block: Block) -> bool:
        """Check if block was created too fast."""
        time_diff = block.timestamp - previous_block.timestamp
        return time_diff >= self.config["block_interval"] / 2

    def _validate_block_events(self, block: Block) -> bool:
        """Validate all events in a block for consensus."""
        return all(
            self.validate_event_for_consensus(event)
            for event in block.to_event_list()
        )

    def validate_block(self, block: Block, previous_block: Block) -> bool:
        """
        Validate a block according to PoA consensus rules.

        Returns:
            True if block is valid, False otherwise
        """
        if not block.validate_structure():
            return False

        return (
            self._check_block_timing(block, previous_block) and
            self._validate_block_events(block) and
            (
                not self.config["require_authority_signature"]
                or self._has_valid_authority_signature(block)
            ) and
            _verify_block_zk_proof(block)
        )

    def finalize_block(
        self,
        block: Block,
        authority_id: str | None = None,
        private_key: str | None = None
    ) -> Block:
        """
        Finalize a block according to PoA consensus.

        Returns:
            Finalized block with PoA consensus data
        """
        if authority_id and self.is_authority(authority_id):
            authority_signature = _create_authority_signature(
                block, authority_id, private_key
            )
            consensus_event = {
                "event": "consensus_finalization",
                "entity_id": "system_consensus",
                "timestamp": time.time(),
                "details": {
                    "consensus_type": "proof_of_authority",
                    "authority_id": authority_id,
                    "authority_signature": authority_signature,
                    "finalized_at": time.time()
                }
            }
            current_events = block.to_event_list()
            current_events.append(consensus_event)
            block = Block(
                index=block.index,
                events=current_events,
                previous_hash=block.previous_hash,
                timestamp=block.timestamp,
                nonce=block.nonce
            )
        return block

    def _has_valid_authority_signature(self, block: Block) -> bool:
        """Check if block has a valid authority signature."""
        consensus_event = next(
            (
                e for e in block.to_event_list()
                if e.get("event") == "consensus_finalization"
            ),
            None
        )

        if not consensus_event:
            return False

        details = consensus_event.get("details", {})
        authority_id = details.get("authority_id")

        if not authority_id or not self.is_authority(authority_id):
            return False

        signature = details.get("authority_signature")
        public_key = self.authority_metadata.get(authority_id, {}).get("public_key")

        if not public_key or not signature:
            return True

        from hierachain.security.security_utils import verify_signature
        sig_str = f"{block.hash}{authority_id}{details.get('timestamp')}"
        return verify_signature(public_key, sig_str.encode(), signature)

    def get_next_authority(self, current_block_index: int) -> str | None:
        """Get the next authority that should create a block (round-robin)."""
        if not self.authorities:
            return None

        authorities_list = sorted(list(self.authorities))
        next_index = (current_block_index + 1) % len(authorities_list)
        return authorities_list[next_index]

    def get_authority_stats(self) -> dict[str, Any]:
        """Get statistics about authorities."""
        return {
            "total_authorities": len(self.authorities),
            "authorities": list(self.authorities),
            "authority_metadata": self.authority_metadata,
            "max_authorities": self.config["max_authorities"]
        }

    def validate_event_for_consensus(self, event: dict[str, Any]) -> bool:
        """
        Validate an event according to PoA consensus rules.

        Returns:
            True if event is valid, False otherwise
        """
        if not super().validate_event_for_consensus(event):
            return False

        return (
            _validate_entity_id(event.get("entity_id")) and
            event.get("event", "") not in self.FORBIDDEN_EVENT_TYPES
        )

    def reset_consensus_state(self) -> None:
        """Reset PoA consensus state."""
        # Keep authorities but reset any temporary state

    def get_block_creation_difficulty(self) -> float:
        """
        Get block creation difficulty for PoA (always 1.0 since no mining).

        Returns:
            Difficulty value (1.0 for PoA)
        """
        return 1.0

    def estimate_block_time(self) -> float:
        """Estimate block creation time for PoA."""
        return self.config["block_interval"]

    def __str__(self) -> str:
        """String representation of PoA consensus."""
        return f"ProofOfAuthority(authorities={len(self.authorities)})"

    def __repr__(self) -> str:
        """Detailed string representation of PoA consensus."""
        return (f"ProofOfAuthority(name={self.name}, "
                f"authorities={len(self.authorities)}, "
                f"block_interval={self.config['block_interval']})")


def _create_authority_signature(
    block: Block, authority_id: str, private_key: str | None = None
) -> str:
    """Create an authority signature for the block."""
    signature_data = {
        "block_hash": block.hash,
        "authority_id": authority_id,
        "timestamp": time.time(),
        "block_index": block.index
    }
    # Create plain text signature data
    sig_str = (
        f"{signature_data['block_hash']}"
        f"{authority_id}"
        f"{signature_data['timestamp']}"
    )
    
    if private_key:
        try:
            kp = KeyPair.from_private_key(private_key)
            return kp.sign(sig_str.encode())
        except Exception as e:
            logger.error("Failed to sign block with private key: %s", e)
            
    # Fallback to random signature for tests running without proper keys setup
    return KeyPair().sign(sig_str.encode())


def _validate_entity_id(entity_id: Any) -> bool:
    """Validate optional entity_id structure."""
    if entity_id is None:
        return True
    return isinstance(entity_id, str)
