import gradio as gr
import requests
from smolagents import CodeAgent, Tool, HfApiModel

# Fetch the API data once at startup
API_URL = "https://dt074px2e9qbh.cloudfront.net/exploit-api.json"
try:
    response = requests.get(API_URL)
    flagged_addresses = response.json()
except Exception as e:
    flagged_addresses = []
    print(f"Failed to load API data: {e}")

# Function to infer chain from address (basic heuristics)
def infer_chain(address):
    address = address.lower()
    if address.startswith("0x") and len(address) == 42:
        return "Ethereum"
    elif address.startswith("bc1") or address.startswith("1") or address.startswith("3"):
        return "Bitcoin"
    elif address.startswith("T") and len(address) == 34:
        return "Tron"
    else:
        return "Unknown"

# Tool to check if an address is flagged
def check_flagged_address(address: str) -> str:
    """Checks if a cryptocurrency address is flagged in the Bybit hack database."""
    chain = infer_chain(address)
    for entry in flagged_addresses:
        if isinstance(entry, dict) and "address" in entry:
            if entry["address"].lower() == address.lower():
                chain = entry.get("chain", chain)  # Use API chain if available
                return f"Address {address} is flagged as part of the Bybit hack on {chain}."
    return f"Address {address} is not flagged as part of the Bybit hack (inferred chain: {chain})."

# Define the tool for the agent
flagged_address_tool = Tool(
    name="check_flagged_address",
    description="Checks if a cryptocurrency address is flagged in the Bybit hack database.",
    fn=check_flagged_address
)

# Gradio chat function
def chat_with_agent(user_input, chat_history, hf_token):
    if not chat_history:
        chat_history = []
    
    # Check if token is provided
    if not hf_token or hf_token.strip() == "":
        return chat_history + [[user_input, "Please enter a valid Hugging Face API token."]], ""
    
    # Initialize the CodeAgent with the user-provided HF token
    agent = CodeAgent(
        model=HfApiModel(model_id="mistralai/Mixtral-8x7B-Instruct-v0.1", token=hf_token),
        tools=[flagged_address_tool],
        system_prompt="You are an agent that checks if cryptocurrency addresses are flagged as part of the Bybit hack. Given an address, use the check_flagged_address tool to verify its status and report the result, including the blockchain it's on."
    )
    
    # Run the agent with the user's input
    try:
        response = agent.run(user_input)
    except Exception as e:
        response = f"Error: {str(e)}. Check your HF token or try again."
    
    # Append as a [user, bot] pair for Chatbot
    chat_history.append([user_input, response])
    
    # Return updated history and clear the input box
    return chat_history, ""

# Create the Gradio interface with tabs
with gr.Blocks(title="Bybit Hack Address Checker") as demo:
    gr.Markdown("# Bybit Hack Address Checker")
    
    with gr.Tabs():
        # Tab 1: Address Checker UI
        with gr.Tab(label="Check Address"):
            gr.Markdown("Enter a cryptocurrency address to check if it's flagged in the Bybit hack database and identify its blockchain.")
            gr.Markdown("Provide your Hugging Face API token below (get it from [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)).")
            
            # HF Token input
            hf_token_input = gr.Textbox(
                placeholder="Enter your Hugging Face API token here",
                label="Hugging Face API Token",
                type="password",
                lines=1
            )
            
            # Chatbot component
            chatbot = gr.Chatbot(label="Conversation")
            
            # Address input
            msg = gr.Textbox(
                placeholder="Enter address here (e.g., 0x123..., bc1q...)",
                label="Address",
                lines=1
            )
            
            # Clear button
            clear = gr.Button("Clear")
            
            # Submit event
            msg.submit(
                fn=chat_with_agent,
                inputs=[msg, chatbot, hf_token_input],
                outputs=[chatbot, msg]
            )
            
            # Clear event
            clear.click(
                fn=lambda: ([], ""),
                inputs=None,
                outputs=[chatbot, msg]
            )
        
        # Tab 2: Explanation
        with gr.Tab(label="About This App"):
            gr.Markdown("""
                ## What This App Does
                This application helps you verify if a cryptocurrency address is associated with the Bybit hack, as tracked by the Elliptic dataset. Here’s how it works:

                1. **Input an Address**: Enter a wallet address (e.g., Ethereum `0x...`, Bitcoin `bc1q...`, Tron `T...`) in the 'Check Address' tab.
                2. **Provide HF Token**: Supply a Hugging Face API token to access the AI model powering the agent.
                3. **Agent Processing**: The app uses a `CodeAgent` from the `smolagents` library, powered by the free `Mixtral-8x7B-Instruct-v0.1` model via the Hugging Face Inference API.
                4. **Tool Usage**: The agent calls a custom tool (`check_flagged_address`) to compare your input against a list of flagged addresses from the Elliptic Bybit hack dataset (`https://dt074px2e9qbh.cloudfront.net/exploit-api.json`).
                5. **Chain Identification**: If the address isn’t explicitly tagged with a blockchain in the dataset, the app infers it (Ethereum, Bitcoin, Tron, or Unknown) based on common address formats.
                6. **Result**: You’ll see whether the address is flagged and its inferred or confirmed blockchain in the chat interface.

                ### Why It’s Useful
                - **Security**: Quickly check if an address you’re dealing with is linked to a known exploit.
                - **Free & Open**: Uses free-tier AI and public data, making it accessible to anyone with an HF account.
                - **Educational**: Demonstrates AI-driven blockchain analysis with `smolagents`.

                Get started by entering your HF token and an address in the 'Check Address' tab!
            """)

# Launch the app
demo.launch()