"""
proof_of_bandwidth.py

This module implements a simple proof-of-bandwidth mechanism.
Whenever a node successfully serves data (e.g. a file or a chunk),
the node generates a “proof” record that includes:
  • A unique session identifier (e.g. derived from the file hash
    and the remote node’s identifier)
  • The number of bytes served
  • A timestamp of when the transfer finished
  • A digital signature over these values (using the node’s private key)

The ProofOfBandwidthManager then verifies the proof and awards the node
credits proportionally to the number of bytes served.
"""

import time
import logging
from cryptography.exceptions import InvalidSignature
from .crypto import sign_message, verify_signature

# If no logging handlers have been configured yet, configure basic logging.
if not logging.getLogger().hasHandlers():
    logging.basicConfig(
        level=logging.DEBUG,
        format='%(asctime)s %(name)s %(levelname)s: %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S',
    )

logger = logging.getLogger("proof_of_bandwidth")

class ProofOfBandwidthManager:
    def __init__(self, node, credit_manager, reward_rate=1e-6):
        """
        Initializes the proof-of-bandwidth manager.

        Args:
            node: The StorageNode instance serving data.
            credit_manager: An instance of CreditManager to award credits.
            reward_rate: Rate of credits per byte served (default: 1e-6 credits/byte).
        """
        self.node = node
        self.credit_manager = credit_manager
        self.reward_rate = reward_rate
        self.logger = logging.getLogger("proof_of_bandwidth")
        # Maintain a set of processed proofs to avoid duplicates.
        self.processed_proofs = set()
        self.logger.info(
            f"ProofOfBandwidthManager initialized for node {self.node.node_id} with reward rate {self.reward_rate} credits/byte."
        )

    async def generate_proof(self, session_id, bytes_served):
        """
        Generates a bandwidth proof.

        Args:
            session_id (str): A unique session identifier (for example: "{file_hash}_{remote_node_id}").
            bytes_served (int): The number of bytes successfully sent.

        Returns:
            dict: Proof data containing session_id, bytes_served, timestamp, signature, and node_id.
        """
        timestamp = time.time()
        proof_payload = f"{session_id}:{bytes_served}:{timestamp}".encode("utf-8")
        signature = sign_message(self.node.private_key, proof_payload)
        proof = {
            "session_id": session_id,
            "bytes_served": bytes_served,
            "timestamp": timestamp,
            "signature": signature.hex(),
            "node_id": self.node.node_id,
        }
        self.logger.info(
            f"Proof generated for session '{session_id}', bytes served: {bytes_served}."
        )
        self.logger.debug(f"Generated proof-of-bandwidth payload: {proof}")
        return proof

    async def verify_and_award(self, proof):
        """
        Verifies the proof and awards credits proportional to the bytes served.

        In a production setup, a validator network might verify the signature
        using the node’s public key. For this demo, we assume the proof is self-signed.

        Args:
            proof (dict): The proof data as generated by generate_proof().

        Returns:
            bool: True if the proof is valid and credits were awarded, False otherwise.
        """
        # Create a unique identifier for this proof (using session_id and timestamp)
        proof_id = proof.get("session_id") + str(proof.get("timestamp"))
        if proof_id in self.processed_proofs:
            self.logger.info("Proof already processed, skipping duplicate award.")
            return False

        # Ensure the proof has the required fields.
        required_fields = ["session_id", "bytes_served", "timestamp", "signature", "node_id"]
        if not all(field in proof for field in required_fields):
            self.logger.error("Invalid proof structure; missing fields.")
            return False

        # Here you could verify the signature. For this example, we assume
        # the proof is valid. (Optionally, uncomment the following lines if
        # your public key is available and verification is desired.)
        #
        # payload = f"{proof['session_id']}:{proof['bytes_served']}:{proof['timestamp']}".encode("utf-8")
        # valid = verify_signature(self.node.public_key_bytes, payload, bytes.fromhex(proof["signature"]))
        # if not valid:
        #     self.logger.error("Proof signature invalid.")
        #     return False

        self.logger.debug("Proof signature assumed valid (or externally verified).")

        # Compute the reward based on bytes served.
        try:
            bytes_served = float(proof["bytes_served"])
        except (TypeError, ValueError):
            self.logger.error("Invalid bytes_served value in proof.")
            return False

        reward = bytes_served * self.reward_rate
        self.logger.info(
            f"Awarding {reward} credits to node {proof['node_id']} for serving {bytes_served} bytes."
        )

        # Award credits to the node using the credit manager’s method.
        request_id = proof_id  # Using proof_id to ensure the same proof is not processed twice.
        await self.credit_manager.add_credit_once(proof["node_id"], request_id, reward)
        self.processed_proofs.add(proof_id)
        self.logger.info(f"Proof processed and credits awarded. Proof ID: {proof_id}.")
        return True