# app.py
from transformers import AutoTokenizer, AutoModelForCausalLM
import requests
import gradio as gr
import torch

# Load the CodeParrot model and tokenizer (only once)
model_name = "codeparrot/codeparrot-small"  # CodeParrot model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Groq API configuration
GROQ_API_KEY = "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"
GROQ_API_URL = "https://api.groq.com/v1/completions"

# Function to query Groq API
def query_groq(prompt):
    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json"
    }
    data = {
        "prompt": prompt,
        "max_tokens": 150
    }
    response = requests.post(GROQ_API_URL, headers=headers, json=data)
    return response.json()["choices"][0]["text"]

# Function to generate smart contract code
def generate_smart_contract(language, requirements):
    # Create a prompt for the model
    prompt = f"Generate a {language} smart contract with the following requirements: {requirements}"
    
    # Use the CodeParrot model to generate code
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
    outputs = model.generate(**inputs, max_length=300)  # Increased max_length for better results
    generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Enhance the code using Groq API
    enhanced_code = query_groq(generated_code)
    
    return enhanced_code

# Custom CSS for a 3D CGI Figma-like feel
custom_css = """
body {
    font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif;
    background: linear-gradient(135deg, #1e3c72 0%, #2a5298 100%);
    color: #fff;
    perspective: 1000px;
    overflow: hidden;
}

.gradio-container {
    background: rgba(255, 255, 255, 0.1);
    border-radius: 15px;
    padding: 20px;
    box-shadow: 0 4px 30px rgba(0, 0, 0, 0.1);
    backdrop-filter: blur(10px);
    border: 1px solid rgba(255, 255, 255, 0.3);
    transform-style: preserve-3d;
    transform: rotateY(0deg) rotateX(0deg);
    transition: transform 0.5s ease;
}

.gradio-container:hover {
    transform: rotateY(10deg) rotateX(10deg);
}

.gradio-input, .gradio-output {
    background: rgba(255, 255, 255, 0.2);
    border: none;
    border-radius: 10px;
    padding: 10px;
    color: #fff;
    transform-style: preserve-3d;
    transition: transform 0.3s ease;
}

.gradio-input:focus, .gradio-output:focus {
    background: rgba(255, 255, 255, 0.3);
    outline: none;
    transform: translateZ(20px);
}

.gradio-button {
    background: linear-gradient(135deg, #6a11cb 0%, #2575fc 100%);
    border: none;
    border-radius: 10px;
    color: #fff;
    padding: 10px 20px;
    font-size: 16px;
    cursor: pointer;
    transition: background 0.3s ease, transform 0.3s ease;
    transform-style: preserve-3d;
}

.gradio-button:hover {
    background: linear-gradient(135deg, #2575fc 0%, #6a11cb 100%);
    transform: translateZ(10px);
}

h1 {
    text-align: center;
    font-size: 2.5em;
    margin-bottom: 20px;
    color: white; /* White title color */
    transform-style: preserve-3d;
    transform: translateZ(30px);
}

@keyframes float {
    0% {
        transform: translateY(0) translateZ(0);
    }
    50% {
        transform: translateY(-10px) translateZ(10px);
    }
    100% {
        transform: translateY(0) translateZ(0);
    }
}

.gradio-container {
    animation: float 4s ease-in-out infinite;
}
"""

# Gradio interface for the app
def generate_contract(language, requirements):
    return generate_smart_contract(language, requirements)

# Dropdown options for programming languages
languages = ["Solidity", "Vyper", "Rust", "JavaScript", "Python"]

interface = gr.Interface(
    fn=generate_contract,
    inputs=[
        gr.Dropdown(label="Programming Language", choices=languages, value="Solidity"),  # Dropdown menu
        gr.Textbox(label="Requirements", placeholder="e.g., ERC20 token with minting functionality")
    ],
    outputs=gr.Textbox(label="Generated Smart Contract"),
    title="Smart Contract Generator",
    description="Generate smart contracts using AI.",
    css=custom_css
)

# Launch the Gradio app
if __name__ == "__main__":
    interface.launch()