# examples/blockchain.py

import hashlib
import json
import time
import random

import os
import sys

# Try to import from installed package first, fall back to direct imports
try:
    from chaincraft.shared_object import SharedObject, SharedObjectException
    from chaincraft.shared_message import SharedMessage
except ImportError:
    # Add parent directory to path as fallback
    sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
    from chaincraft.shared_object import SharedObject, SharedObjectException
    from chaincraft.shared_message import SharedMessage
from typing import List, Dict, Any, Optional
import secrets
from dataclasses import dataclass
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.exceptions import InvalidSignature
import base64


class BlockchainUtils:
    """Utility functions for blockchain operations"""

    @staticmethod
    def calculate_hash(data: Any) -> str:
        """Calculate SHA-256 hash of JSON-serialized data"""
        if isinstance(data, dict) or isinstance(data, list):
            data_str = json.dumps(data, sort_keys=True)
        elif not isinstance(data, str):
            data_str = str(data)
        else:
            data_str = data

        return hashlib.sha256(data_str.encode()).hexdigest()

    @staticmethod
    def verify_proof_of_work(block_data: Dict, nonce: int, difficulty: int) -> bool:
        """Verify if a given nonce meets the proof-of-work requirement"""
        # Create copy of block data without the nonce to create the challenge
        challenge = {k: v for k, v in block_data.items() if k != "nonce"}
        challenge_hash = BlockchainUtils.calculate_hash(challenge)

        # Combine challenge hash with nonce to get final hash
        final_hash = BlockchainUtils.calculate_hash(challenge_hash + str(nonce))

        # Check if hash meets difficulty (has required number of leading zeros)
        return final_hash.startswith("0" * difficulty)

    @staticmethod
    def find_proof_of_work(block_data: Dict, difficulty: int) -> tuple:
        """Find a nonce that satisfies the proof-of-work requirement"""
        # Create copy of block data without the nonce to create the challenge
        challenge = {k: v for k, v in block_data.items() if k != "nonce"}
        challenge_hash = BlockchainUtils.calculate_hash(challenge)

        nonce = 0
        max_nonce = 2**32  # Prevent infinite loops

        while nonce < max_nonce:
            # Combine challenge hash with nonce to get final hash
            final_hash = BlockchainUtils.calculate_hash(challenge_hash + str(nonce))

            # Check if hash meets difficulty
            if final_hash.startswith("0" * difficulty):
                return nonce, final_hash

            nonce += 1

        raise Exception(f"Couldn't find valid proof after {max_nonce} attempts")

    @staticmethod
    def generate_keypair() -> tuple:
        """Generate ECDSA keypair for transaction signing"""
        private_key = ec.generate_private_key(ec.SECP256K1)
        public_key = private_key.public_key()

        # Convert to strings for storage/transmission
        private_key_str = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        ).decode()
        public_key_str = public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        ).decode()

        return private_key_str, public_key_str

    @staticmethod
    def get_address_from_public_key(public_key: str) -> str:
        """Generate Ethereum-like address from public key"""
        try:
            # Load the public key from PEM format
            key = serialization.load_pem_public_key(public_key.encode())
            
            # Get public key in raw bytes format
            public_key_bytes = key.public_bytes(
                encoding=serialization.Encoding.X962,
                format=serialization.PublicFormat.UncompressedPoint
            )
            
            # Remove the first byte (0x04 for uncompressed keys) and hash
            public_key_bytes = public_key_bytes[1:]
            
            # Hash the public key and take last 20 bytes (like Ethereum)
            address = hashlib.sha256(public_key_bytes).digest()[-20:].hex()
            return f"0x{address}"
        except Exception as e:
            print(f"Error generating address: {e}")
            return f"0x{'0' * 40}"  # Return invalid address in case of error

    @staticmethod
    def sign_transaction(tx_data: Dict, private_key_str: str) -> str:
        """Sign transaction data with private key"""
        # Remove signature field if present when signing
        tx_copy = {k: v for k, v in tx_data.items() if k != "signature"}
        message = json.dumps(tx_copy, sort_keys=True).encode()

        # Load private key from PEM string
        private_key = serialization.load_pem_private_key(
            private_key_str.encode(), 
            password=None
        )

        # Sign the message using deterministic ECDSA (RFC 6979)
        signature = private_key.sign(
            message,
            ec.ECDSA(hashes.SHA256())
        )
        
        # Convert to hex
        return signature.hex()

    @staticmethod
    def verify_signature(tx_data: Dict, signature: str, public_key_str: str) -> bool:
        """Verify transaction signature with public key"""
        # Remove signature field if present when verifying
        tx_copy = {k: v for k, v in tx_data.items() if k != "signature"}
        message = json.dumps(tx_copy, sort_keys=True).encode()

        try:
            # Convert hex string to bytes for signature
            signature_bytes = bytes.fromhex(signature)
            
            # Load public key from PEM string
            public_key = serialization.load_pem_public_key(public_key_str.encode())
            
            # Verify the signature
            public_key.verify(
                signature_bytes,
                message,
                ec.ECDSA(hashes.SHA256())
            )
            return True
        except Exception as e:
            print(f"Signature verification error: {e}")
            return False


@dataclass
class Transaction:
    """Represents a signed transaction transferring value between addresses"""

    sender: str  # Sender's address
    recipient: str  # Recipient's address
    amount: float  # Amount to transfer
    fee: float  # Transaction fee
    timestamp: float  # Transaction creation time
    public_key: str  # Sender's public key (for verification)
    signature: str  # Transaction signature
    tx_id: str  # Transaction ID (hash)

    @classmethod
    def create(
        cls,
        sender: str,
        recipient: str,
        amount: float,
        fee: float,
        private_key: str,
        public_key: str,
    ) -> "Transaction":
        """Create and sign a new transaction"""
        # Basic transaction data
        tx_data = {
            "sender": sender,
            "recipient": recipient,
            "amount": amount,
            "fee": fee,
            "timestamp": time.time(),
        }

        # Sign the transaction
        signature = BlockchainUtils.sign_transaction(tx_data, private_key)

        # Add signature and public key
        tx_data["signature"] = signature
        tx_data["public_key"] = public_key

        # Generate transaction ID
        tx_id = BlockchainUtils.calculate_hash(tx_data)
        tx_data["tx_id"] = tx_id

        return cls(**tx_data)

    def to_dict(self) -> Dict:
        """Convert transaction to dictionary"""
        return {
            "tx_id": self.tx_id,
            "sender": self.sender,
            "recipient": self.recipient,
            "amount": self.amount,
            "fee": self.fee,
            "timestamp": self.timestamp,
            "public_key": self.public_key,
            "signature": self.signature,
        }

    @classmethod
    def from_dict(cls, tx_dict: Dict) -> "Transaction":
        """Create Transaction object from dictionary"""
        return cls(
            tx_id=tx_dict["tx_id"],
            sender=tx_dict["sender"],
            recipient=tx_dict["recipient"],
            amount=tx_dict["amount"],
            fee=tx_dict["fee"],
            timestamp=tx_dict["timestamp"],
            public_key=tx_dict["public_key"],
            signature=tx_dict["signature"],
        )

    def is_valid(self) -> bool:
        """Verify transaction integrity and signature"""
        # Check transaction has valid structure
        if not all(
            hasattr(self, attr)
            for attr in [
                "sender",
                "recipient",
                "amount",
                "fee",
                "timestamp",
                "public_key",
                "signature",
                "tx_id",
            ]
        ):
            return False

        # Check that amount and fee are positive
        if self.amount <= 0 or self.fee < 0:
            return False

        # Check that sender address matches public key
        derived_address = BlockchainUtils.get_address_from_public_key(self.public_key)
        if derived_address != self.sender:
            return False

        # Verify the signature
        tx_data = self.to_dict()
        return BlockchainUtils.verify_signature(
            tx_data, self.signature, self.public_key
        )


@dataclass
class Block:
    """Represents a block in the blockchain"""

    index: int  # Block height in the chain
    timestamp: float  # Block creation time
    transactions: List[Dict]  # List of transactions included in this block
    previous_hash: str  # Hash of the previous block
    miner: str  # Address of the miner (for reward)
    nonce: int  # Proof-of-work nonce
    hash: str  # Block hash

    @classmethod
    def create(
        cls,
        index: int,
        transactions: List[Dict],
        previous_hash: str,
        miner: str,
        difficulty: int,
    ) -> "Block":
        """Create and mine a new block"""
        block_data = {
            "index": index,
            "timestamp": time.time(),
            "transactions": transactions,
            "previous_hash": previous_hash,
            "miner": miner,
        }

        # Find proof of work
        nonce, block_hash = BlockchainUtils.find_proof_of_work(block_data, difficulty)

        # Add nonce and hash
        block_data["nonce"] = nonce
        block_data["hash"] = block_hash

        return cls(**block_data)

    def to_dict(self) -> Dict:
        """Convert block to dictionary"""
        return {
            "index": self.index,
            "timestamp": self.timestamp,
            "transactions": self.transactions,
            "previous_hash": self.previous_hash,
            "miner": self.miner,
            "nonce": self.nonce,
            "hash": self.hash,
        }

    @classmethod
    def from_dict(cls, block_dict: Dict) -> "Block":
        """Create Block object from dictionary"""
        return cls(
            index=block_dict["index"],
            timestamp=block_dict["timestamp"],
            transactions=block_dict["transactions"],
            previous_hash=block_dict["previous_hash"],
            miner=block_dict["miner"],
            nonce=block_dict["nonce"],
            hash=block_dict["hash"],
        )

    def is_valid(self, difficulty: int) -> bool:
        """Verify block integrity and proof-of-work"""
        # Check block has valid structure
        if not all(
            hasattr(self, attr)
            for attr in [
                "index",
                "timestamp",
                "transactions",
                "previous_hash",
                "miner",
                "nonce",
                "hash",
            ]
        ):
            return False

        # Check proof of work
        block_data = self.to_dict()
        return BlockchainUtils.verify_proof_of_work(block_data, self.nonce, difficulty)


class Mempool(SharedObject):
    """
    Mempool for holding pending transactions before they're included in blocks.
    Not merklelized since it's a temporary storage.
    """

    def __init__(self, difficulty: int = 4):
        """Initialize mempool with empty transactions dict"""
        self.transactions: Dict[str, Transaction] = {}  # tx_id -> Transaction
        self.difficulty = difficulty

    def is_valid(self, message: SharedMessage) -> bool:
        """
        Check if message contains a valid transaction or block
        """
        try:
            data = message.data

            # Handle transaction message
            if (
                isinstance(data, dict)
                and "type" in data
                and data["type"] == "transaction"
            ):
                tx_data = data["payload"]
                tx = Transaction.from_dict(tx_data)
                return tx.is_valid()

            # Handle block message (which will clear transactions from mempool)
            elif isinstance(data, dict) and "type" in data and data["type"] == "block":
                block_data = data["payload"]
                block = Block.from_dict(block_data)
                return block.is_valid(self.difficulty)

            return False
        except Exception as e:
            print(f"Error validating message: {e}")
            return False

    def add_message(self, message: SharedMessage) -> None:
        """
        Process a new message - either a transaction to add to mempool
        or a block that will clear transactions from the mempool
        """
        data = message.data

        # Handle transaction message
        if isinstance(data, dict) and "type" in data and data["type"] == "transaction":
            tx_data = data["payload"]
            tx = Transaction.from_dict(tx_data)

            # Add to mempool if not already there
            if tx.tx_id not in self.transactions:
                self.transactions[tx.tx_id] = tx
                print(f"Added transaction {tx.tx_id[:8]} to mempool")

        # Handle block message
        elif isinstance(data, dict) and "type" in data and data["type"] == "block":
            block_data = data["payload"]
            block = Block.from_dict(block_data)

            # Remove transactions included in the block from mempool
            for tx_dict in block.transactions:
                tx_id = tx_dict["tx_id"]
                if tx_id in self.transactions:
                    del self.transactions[tx_id]

            print(
                f"Cleared {len(block.transactions)} transactions from mempool after block {block.index}"
            )

    # These methods aren't needed for Mempool since it's non-merklelized,
    # but SharedObject requires them to be implemented
    def is_merkelized(self) -> bool:
        return False

    def get_latest_digest(self) -> str:
        return ""

    def has_digest(self, hash_digest: str) -> bool:
        return False

    def is_valid_digest(self, hash_digest: str) -> bool:
        return False

    def add_digest(self, hash_digest: str) -> bool:
        return False

    def gossip_object(self, digest) -> List[SharedMessage]:
        return []

    def get_messages_since_digest(self, digest: str) -> List[SharedMessage]:
        return []

    def get_transactions_by_fee(self, max_count: int = 10) -> List[Transaction]:
        """Get transactions sorted by fee (highest first), up to max_count"""
        sorted_txs = sorted(
            self.transactions.values(), key=lambda tx: tx.fee, reverse=True
        )
        return sorted_txs[:max_count]


class Ledger(SharedObject):
    """
    Blockchain ledger implementation that maintains the chain of blocks
    and tracks account balances. Merklelized for efficient state sync.
    """

    def __init__(self, difficulty: int = 4, reward: float = 10.0):
        """Initialize blockchain with genesis block"""
        self.chain: List[Block] = []
        self.balances: Dict[str, float] = {}  # address -> balance
        self.difficulty = difficulty
        self.mining_reward = reward

        # Create genesis block
        self._create_genesis_block()

        # Save chain hashes for merklelization
        self.chain_hashes: List[str] = [self.chain[0].hash]

    def _create_genesis_block(self) -> None:
        """Create the genesis (first) block in the chain"""
        genesis_address = "0x0000000000000000000000000000000000000000"
        genesis_block = Block(
            index=0,
            timestamp=time.time(),
            transactions=[],
            previous_hash="0" * 64,
            miner=genesis_address,
            nonce=0,
            hash=BlockchainUtils.calculate_hash(
                {
                    "index": 0,
                    "timestamp": time.time(),
                    "transactions": [],
                    "previous_hash": "0" * 64,
                    "miner": genesis_address,
                }
            ),
        )
        self.chain.append(genesis_block)

        # Give genesis address some coins
        self.balances[genesis_address] = 1000.0

    def is_valid(self, message: SharedMessage) -> bool:
        """
        Check if message contains a valid block that can be added to the chain
        """
        try:
            data = message.data

            # Only accept block messages
            if isinstance(data, dict) and "type" in data and data["type"] == "block":
                block_data = data["payload"]
                block = Block.from_dict(block_data)

                # Check block integrity and proof-of-work
                if not block.is_valid(self.difficulty):
                    print(f"Invalid block: Failed proof-of-work check")
                    return False

                # Check block index
                if block.index != len(self.chain):
                    print(
                        f"Invalid block: Expected index {len(self.chain)}, got {block.index}"
                    )
                    return False

                # Check previous hash
                if block.previous_hash != self.chain[-1].hash:
                    print(f"Invalid block: Previous hash doesn't match")
                    return False

                # Validate each transaction in the block
                for tx_dict in block.transactions:
                    tx = Transaction.from_dict(tx_dict)

                    # Check transaction signature
                    if not tx.is_valid():
                        print(f"Invalid transaction {tx.tx_id[:8]} in block")
                        return False

                    # Check sender has enough balance
                    sender_balance = self.balances.get(tx.sender, 0)
                    if sender_balance < tx.amount + tx.fee:
                        print(f"Insufficient balance for tx {tx.tx_id[:8]}")
                        return False

                return True

            return False
        except Exception as e:
            print(f"Error validating block message: {e}")
            return False

    def add_message(self, message: SharedMessage) -> None:
        """
        Process a new block and update the chain and balances
        """
        data = message.data

        # Only process block messages
        if isinstance(data, dict) and "type" in data and data["type"] == "block":
            block_data = data["payload"]
            block = Block.from_dict(block_data)

            # Process the block
            self._add_block(block)

            # Update merklelized chain hashes
            self.chain_hashes.append(block.hash)

            print(f"Added block {block.index} to chain, hash: {block.hash[:8]}")

    def _add_block(self, block: Block) -> None:
        """Add a validated block to the chain and update balances"""
        # Add the block to the chain
        self.chain.append(block)

        # Process mining reward
        self.balances[block.miner] = (
            self.balances.get(block.miner, 0) + self.mining_reward
        )

        # Process transactions
        for tx_dict in block.transactions:
            tx = Transaction.from_dict(tx_dict)

            # Deduct amount from sender
            self.balances[tx.sender] = (
                self.balances.get(tx.sender, 0) - tx.amount - tx.fee
            )

            # Add amount to recipient
            self.balances[tx.recipient] = self.balances.get(tx.recipient, 0) + tx.amount

            # Add fee to miner
            self.balances[block.miner] = self.balances.get(block.miner, 0) + tx.fee

    # Merklelized methods for state synchronization
    def is_merkelized(self) -> bool:
        return True

    def get_latest_digest(self) -> str:
        """Return the hash of the latest block"""
        return self.chain[-1].hash

    def has_digest(self, hash_digest: str) -> bool:
        """Check if a block with the given hash exists in the chain"""
        return hash_digest in self.chain_hashes

    def is_valid_digest(self, hash_digest: str) -> bool:
        """Check if a block hash is valid for sync"""
        return hash_digest in self.chain_hashes

    def add_digest(self, hash_digest: str) -> bool:
        """Not used in this implementation"""
        return False

    def gossip_object(self, digest) -> List[SharedMessage]:
        """
        Return blocks since the given digest hash for synchronization
        """
        try:
            # Find index of the digest in the chain
            if digest not in self.chain_hashes:
                return []

            index = self.chain_hashes.index(digest)

            # Return all blocks after this index
            messages = []
            for i in range(index + 1, len(self.chain)):
                block_dict = self.chain[i].to_dict()
                message_data = {"type": "block", "payload": block_dict}
                messages.append(SharedMessage(data=message_data))

            return messages
        except Exception as e:
            print(f"Error in gossip_object: {e}")
            return []

    def get_messages_since_digest(self, digest: str) -> List[SharedMessage]:
        """Same as gossip_object for this implementation"""
        return self.gossip_object(digest)

    def create_block(
        self, transactions: List[Transaction], miner_address: str
    ) -> Block:
        """
        Create a new block with the given transactions and miner
        """
        # Convert transactions to dictionaries for block
        tx_dicts = [tx.to_dict() for tx in transactions]

        # Create and return a new block
        return Block.create(
            index=len(self.chain),
            transactions=tx_dicts,
            previous_hash=self.chain[-1].hash,
            miner=miner_address,
            difficulty=self.difficulty,
        )


class BlockchainNode:
    """
    Helper class to manage blockchain operations for a node
    """

    def __init__(self, chaincraft_node, difficulty: int = 4, reward: float = 10.0):
        """Initialize with ChainCraft node and blockchain configuration"""
        self.node = chaincraft_node
        self.mempool = Mempool(difficulty)
        self.ledger = Ledger(difficulty, reward)

        # Add shared objects to the node
        self.node.add_shared_object(self.mempool)
        self.node.add_shared_object(self.ledger)

        # Generate key pair for this node
        self.private_key, self.public_key = BlockchainUtils.generate_keypair()
        self.address = BlockchainUtils.get_address_from_public_key(self.public_key)

        print(f"Node initialized with address: {self.address}")

    def create_transaction(
        self, recipient: str, amount: float, fee: float = 0.01
    ) -> str:
        """Create and broadcast a transaction"""
        # Create transaction
        tx = Transaction.create(
            sender=self.address,
            recipient=recipient,
            amount=amount,
            fee=fee,
            private_key=self.private_key,
            public_key=self.public_key,
        )

        # Prepare message
        message_data = {"type": "transaction", "payload": tx.to_dict()}

        # Broadcast transaction
        tx_hash, _ = self.node.create_shared_message(message_data)
        print(f"Transaction created and broadcast: {tx.tx_id[:8]}")

        return tx.tx_id

    def mine_block(self) -> Optional[str]:
        """Mine a new block with transactions from mempool"""
        # Get transactions from mempool
        transactions = self.mempool.get_transactions_by_fee(max_count=10)

        if not transactions:
            print("No transactions in mempool to mine")
            return None

        print(f"Mining block with {len(transactions)} transactions")

        # Create new block
        block = self.ledger.create_block(transactions, self.address)

        # Prepare message
        message_data = {"type": "block", "payload": block.to_dict()}

        # Broadcast block
        block_hash, _ = self.node.create_shared_message(message_data)
        print(f"Block mined and broadcast: {block.hash[:8]}")

        return block.hash

    def get_balance(self, address: Optional[str] = None) -> float:
        """Get balance for an address or self if None"""
        if address is None:
            address = self.address
        return self.ledger.balances.get(address, 0)

    def get_blockchain_info(self) -> Dict:
        """Get general information about the blockchain"""
        return {
            "chain_length": len(self.ledger.chain),
            "latest_block_hash": self.ledger.chain[-1].hash,
            "difficulty": self.ledger.difficulty,
            "mempool_size": len(self.mempool.transactions),
            "node_address": self.address,
            "node_balance": self.get_balance(),
        }


# Helper functions for tutorial
def generate_wallet():
    """Generate and return a new wallet with keys and address"""
    private_key, public_key = BlockchainUtils.generate_keypair()
    address = BlockchainUtils.get_address_from_public_key(public_key)

    return {"private_key": private_key, "public_key": public_key, "address": address}


def format_balance(balance: float) -> str:
    """Format balance with 2 decimal places"""
    return f"{balance:.2f}"


if __name__ == "__main__":
    print("Blockchain module loaded. Import this module in your application.")
