import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "GinieAI/Solidity-LLM"

print("Loading Ginie Solidity LLM...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True,
    device_map="cpu",
)
model.eval()
print("Model ready!")


def generate_contract(instruction, max_tokens=600, temperature=0.7):
    if not instruction.strip():
        return "⚠️ Please enter a contract description."

    prompt = f"### Instruction:\n{instruction.strip()}\n\n### Response:\n"

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=512,
    )

    with torch.no_grad():
        outputs = model.generate(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            max_new_tokens=int(max_tokens),
            temperature=float(temperature),
            do_sample=temperature > 0,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    generated = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[1] :],
        skip_special_tokens=True,
    ).strip()

    return generated


examples = [
    ["Write a Solidity ERC20 token with minting, burning, and owner controls."],
    ["Write a Solidity multisig wallet requiring 2 of 3 signatures."],
    ["Write a Solidity staking contract with 10% APY rewards."],
    ["Write a Solidity DAO governance contract with proposal voting."],
    ["Write a Solidity NFT contract with whitelist minting and reveal."],
]

css = """
.output-code { font-family: monospace; font-size: 13px; }
footer { display: none !important; }
"""

with gr.Blocks(theme=gr.themes.Soft(), title="Ginie AI — Smart Contract Generator", css=css) as demo:

    gr.Markdown(
        """
# 🔷 Ginie AI — Solidity Smart Contract Generator
Describe the contract you need in plain English. Ginie generates production-ready Solidity.
> ⚠️ Always review AI-generated contracts before deploying to mainnet.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            instruction = gr.Textbox(
                label="📝 Describe your contract",
                placeholder="e.g. Write a Solidity ERC20 token with minting, burning, and owner controls.",
                lines=5,
            )
            with gr.Row():
                max_tokens = gr.Slider(
                    minimum=100, maximum=800, value=600, step=50,
                    label="Max tokens",
                )
                temperature = gr.Slider(
                    minimum=0.1, maximum=1.0, value=0.7, step=0.1,
                    label="Temperature",
                )
            btn = gr.Button("⚡ Generate Contract", variant="primary", size="lg")

        with gr.Column(scale=2):
            output = gr.Code(
                label="Generated Solidity",
                language="javascript",   # Gradio has no Solidity lexer; JS gives syntax highlighting
                lines=30,
                elem_classes=["output-code"],
            )

    gr.Examples(
        examples=examples,
        inputs=instruction,
        label="💡 Example prompts",
    )

    btn.click(
        fn=generate_contract,
        inputs=[instruction, max_tokens, temperature],
        outputs=output,
    )

    gr.Markdown(
        """
---
Model: [GinieAI/Solidity-LLM](https://huggingface.co/GinieAI/Solidity-LLM) · Base: Chain-GPT/Solidity-LLM · License: MIT
        """
    )

demo.launch()