import gradio as gr
import networkx as nx
import matplotlib.pyplot as plt
import hashlib
import time
import io
from PIL import Image

# ------------------- BLOCKCHAIN LOGIC -------------------- #

class Block:
    def __init__(self, index, data, previous_hash="0"):
        self.index = index
        self.timestamp = time.time()
        self.data = data
        self.previous_hash = previous_hash
        self.nonce = 0
        self.hash = self.calculate_hash()
    
    def calculate_hash(self):
        block_string = (
            str(self.index) +
            str(self.timestamp) +
            str(self.data) +
            str(self.previous_hash) +
            str(self.nonce)
        )
        return hashlib.sha256(block_string.encode()).hexdigest()

def mine_block(block, difficulty=1):
    """
    A quick "proof-of-work":
    The block's hash must start with '0' * difficulty, e.g. '0' for difficulty=1.
    If this is still slow for you, set difficulty=0 to skip "mining" entirely.
    """
    target_prefix = "0" * difficulty
    while not block.hash.startswith(target_prefix):
        block.nonce += 1
        block.hash = block.calculate_hash()

def create_genesis_block():
    genesis = Block(index=0, data="Genesis Block", previous_hash="0")
    mine_block(genesis, difficulty=1)
    return genesis

# The global chain
chain = []

# Initialize the chain with a genesis block if empty
if not chain:
    chain.append(create_genesis_block())

def add_block_to_chain(data):
    previous_block = chain[-1]
    new_block = Block(
        index=len(chain),
        data=data,
        previous_hash=previous_block.hash
    )
    mine_block(new_block, difficulty=1)
    chain.append(new_block)

# ------------------- VISUALIZATION -------------------- #

def visualize_chain(chain):
    """
    Creates a LINEAR graph of the blockchain using NetworkX
    and returns it as a PIL image.
    
    We'll place blocks in a line (x=index, y=0) to avoid 
    expensive layout computations each time.
    """
    G = nx.DiGraph()
    
    # Add nodes
    for blk in chain:
        # node label shows index + partial hash
        node_label = f"Block {blk.index}\nHash: {blk.hash[:6]}..."
        G.add_node(blk.index, label=node_label)

    # Add edges
    for i in range(len(chain) - 1):
        G.add_edge(chain[i].index, chain[i+1].index)

    # Fixed, linear layout: x = index, y = 0
    pos = {blk.index: (blk.index, 0) for blk in chain}

    # Draw the graph
    plt.figure(figsize=(10, 3))
    nx.draw_networkx_nodes(G, pos, node_color="lightblue", node_size=1500)
    nx.draw_networkx_edges(G, pos, arrowstyle='->', arrowsize=20)
    node_labels = nx.get_node_attributes(G, 'label')
    nx.draw_networkx_labels(G, pos, labels=node_labels, font_size=8)

    plt.title("Toy Blockchain Visualization (Linear Layout)")
    plt.axis("off")

    # Convert the Matplotlib figure to a PIL image
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight')
    buf.seek(0)
    plt.close()
    return Image.open(buf)

def add_and_visualize_block(data):
    # Add a new block with user-provided data
    add_block_to_chain(data)
    # Generate the updated visualization
    return visualize_chain(chain)

# ------------------- GRADIO INTERFACE -------------------- #

demo = gr.Interface(
    fn=add_and_visualize_block,
    inputs=gr.Textbox(lines=2, label="Data for New Block", placeholder="e.g. 'My transaction'"),
    outputs="image",
    title="Toy Blockchain Demo (Faster Visualization)",
    description=(
        "Enter any data to create a new block. "
        "Mining difficulty is set to 1, so it should be fast. "
        "Blocks are displayed in a linear chain."
    ),
)

demo.launch(debug=True)