"""
Bootnode Service - Peer Discovery for PoA Blockchain Validators

This service helps validators discover each other and bootstrap into the network.
It maintains a registry of known peers and provides discovery endpoints.

Features:
- Peer registration endpoint (POST /register_peer)
- Peer discovery endpoint (GET /discover)
- Peer listing endpoint (GET /peers)
- Health check endpoint
- Periodic cleanup of stale peers
"""

import os
import time
import json
import logging
from typing import Dict, List, Optional
from datetime import datetime, timedelta
from fastapi import FastAPI, HTTPException, status
from pydantic import BaseModel
import asyncio

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# ============================================================================
# Data Models
# ============================================================================

class PeerInfo(BaseModel):
    """Information about a validator peer"""
    node_id: str
    ip: str
    p2p_port: int
    rpc_port: int
    public_key: Optional[str] = None


class PeerRegistration(BaseModel):
    """Request to register a peer"""
    node_id: str
    ip: str
    p2p_port: int
    rpc_port: int
    public_key: Optional[str] = None


class PeerDiscoveryResponse(BaseModel):
    """Response from peer discovery"""
    peers: List[PeerInfo]
    count: int


class HealthResponse(BaseModel):
    """Health check response"""
    status: str
    timestamp: str
    peers_count: int


# ============================================================================
# Bootnode Service
# ============================================================================

class PeerRegistry:
    """In-memory registry of known peers with expiration"""

    def __init__(self, peer_timeout_seconds: int = 300):
        self.peers: Dict[str, dict] = {}  # node_id -> peer info with timestamp
        self.peer_timeout = peer_timeout_seconds

    def register_peer(self, peer: PeerInfo) -> None:
        """Register or update a peer"""
        self.peers[peer.node_id] = {
            "info": peer,
            "registered_at": time.time(),
            "last_heartbeat": time.time()
        }
        logger.info(
            f"Peer registered: {peer.node_id} "
            f"({peer.ip}:{peer.p2p_port}, RPC:{peer.rpc_port})"
        )

    def update_heartbeat(self, node_id: str) -> None:
        """Update heartbeat timestamp for a peer"""
        if node_id in self.peers:
            self.peers[node_id]["last_heartbeat"] = time.time()

    def get_peer(self, node_id: str) -> Optional[PeerInfo]:
        """Get a peer by node_id"""
        if node_id in self.peers:
            return self.peers[node_id]["info"]
        return None

    def get_all_peers(self) -> List[PeerInfo]:
        """Get all active peers"""
        return [entry["info"] for entry in self.peers.values()]

    def get_peers_except(self, exclude_node_id: str) -> List[PeerInfo]:
        """Get all peers except the specified one"""
        return [
            entry["info"]
            for node_id, entry in self.peers.items()
            if node_id != exclude_node_id
        ]

    def cleanup_stale_peers(self) -> int:
        """Remove peers that haven't sent heartbeat recently"""
        current_time = time.time()
        stale_peers = [
            node_id for node_id, entry in self.peers.items()
            if (current_time - entry["last_heartbeat"]) > self.peer_timeout
        ]

        for node_id in stale_peers:
            logger.warning(f"Removing stale peer: {node_id}")
            del self.peers[node_id]

        return len(stale_peers)


# ============================================================================
# FastAPI Application
# ============================================================================

app = FastAPI(
    title="E-Voting Bootnode",
    description="Peer discovery service for PoA validators",
    version="1.0.0"
)

# Global peer registry
peer_registry = PeerRegistry(peer_timeout_seconds=300)


# ============================================================================
# Health Check
# ============================================================================

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint"""
    return HealthResponse(
        status="healthy",
        timestamp=datetime.utcnow().isoformat(),
        peers_count=len(peer_registry.get_all_peers())
    )


# ============================================================================
# Peer Registration
# ============================================================================

@app.post("/register_peer", response_model=PeerDiscoveryResponse)
async def register_peer(peer: PeerRegistration):
    """
    Register a peer node with the bootnode.

    The peer must provide:
    - node_id: Unique identifier (e.g., "validator-1")
    - ip: IP address or Docker service name
    - p2p_port: Port for P2P communication
    - rpc_port: Port for JSON-RPC communication
    - public_key: (Optional) Validator's public key for signing

    Returns: List of other known peers
    """
    try:
        # Register the peer
        peer_info = PeerInfo(**peer.dict())
        peer_registry.register_peer(peer_info)

        # Return other known peers
        other_peers = peer_registry.get_peers_except(peer.node_id)

        logger.info(f"Registration successful. Peer {peer.node_id} now knows {len(other_peers)} peers")

        return PeerDiscoveryResponse(
            peers=other_peers,
            count=len(other_peers)
        )

    except Exception as e:
        logger.error(f"Error registering peer: {e}")
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Failed to register peer: {str(e)}"
        )


# ============================================================================
# Peer Discovery
# ============================================================================

@app.get("/discover", response_model=PeerDiscoveryResponse)
async def discover_peers(node_id: str):
    """
    Discover peers currently in the network.

    Query Parameters:
    - node_id: The requesting peer's node_id (to exclude from results)

    Returns: List of all other known peers
    """
    try:
        # Update heartbeat for the requesting peer
        peer_registry.update_heartbeat(node_id)

        # Return all peers except the requester
        peers = peer_registry.get_peers_except(node_id)

        logger.info(f"Discovery request from {node_id}: returning {len(peers)} peers")

        return PeerDiscoveryResponse(
            peers=peers,
            count=len(peers)
        )

    except Exception as e:
        logger.error(f"Error discovering peers: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to discover peers: {str(e)}"
        )


# ============================================================================
# Peer Listing (Admin)
# ============================================================================

@app.get("/peers", response_model=PeerDiscoveryResponse)
async def list_all_peers():
    """
    List all known peers (admin endpoint).

    Returns: All registered peers
    """
    peers = peer_registry.get_all_peers()

    return PeerDiscoveryResponse(
        peers=peers,
        count=len(peers)
    )


# ============================================================================
# Peer Heartbeat
# ============================================================================

@app.post("/heartbeat")
async def peer_heartbeat(node_id: str):
    """
    Send a heartbeat to indicate the peer is still alive.

    Query Parameters:
    - node_id: The peer's node_id

    This keeps the peer in the registry and prevents timeout.
    """
    try:
        peer_registry.update_heartbeat(node_id)
        logger.debug(f"Heartbeat received from {node_id}")

        return {
            "status": "ok",
            "timestamp": datetime.utcnow().isoformat()
        }

    except Exception as e:
        logger.error(f"Error processing heartbeat: {e}")
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Failed to process heartbeat: {str(e)}"
        )


# ============================================================================
# Stats (Admin)
# ============================================================================

@app.get("/stats")
async def get_stats():
    """Get bootnode statistics"""
    peers = peer_registry.get_all_peers()

    return {
        "total_peers": len(peers),
        "peers": [
            {
                "node_id": p.node_id,
                "ip": p.ip,
                "p2p_port": p.p2p_port,
                "rpc_port": p.rpc_port
            }
            for p in peers
        ],
        "timestamp": datetime.utcnow().isoformat()
    }


# ============================================================================
# Background Tasks
# ============================================================================

async def cleanup_stale_peers_task():
    """Periodic task to cleanup stale peers"""
    while True:
        await asyncio.sleep(60)  # Cleanup every 60 seconds
        removed_count = peer_registry.cleanup_stale_peers()
        if removed_count > 0:
            logger.info(f"Cleaned up {removed_count} stale peers")


@app.on_event("startup")
async def startup_event():
    """Start background tasks on application startup"""
    logger.info("Bootnode starting up...")
    asyncio.create_task(cleanup_stale_peers_task())
    logger.info("Cleanup task started")


@app.on_event("shutdown")
async def shutdown_event():
    """Log shutdown"""
    logger.info("Bootnode shutting down...")


# ============================================================================
# Main
# ============================================================================

if __name__ == "__main__":
    import uvicorn

    port = int(os.getenv("BOOTNODE_PORT", "8546"))
    host = os.getenv("BOOTNODE_HOST", "0.0.0.0")

    logger.info(f"Starting bootnode on {host}:{port}")

    uvicorn.run(
        app,
        host=host,
        port=port,
        log_level="info"
    )