# Updated UI Generator with Better Open-Source Model

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

# Use a better model for code generation, such as DeepSeek-Coder or Codestral
model_id = "deepseek-ai/deepseek-coder-6.7b-instruct"  # Change to another if desired

# Ensure torch uses CPU if no GPU
device = 0 if torch.cuda.is_available() else -1

# Load tokenizer and model
print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)

# Build the text generation pipeline
generator = TextGenerationPipeline(model=model, tokenizer=tokenizer, device=device)

def generate_ui(platform, framework, ui_prompt):
    prompt = f"""
You are an expert mobile app developer.
Generate a complete {framework} UI code snippet for a {platform} app based on the description:
"{ui_prompt}"
Include all required imports, a main method, and best practices for UI structure.
"""
    response = generator(prompt, max_new_tokens=512, do_sample=True, temperature=0.7)[0]['generated_text']
    # Trim the echoed prompt and just return the generated code
    return response.split(""""""")[-1].strip()

interface = gr.Interface(
    fn=generate_ui,
    inputs=[
        gr.Dropdown(["Android", "iOS"], label="Platform"),
        gr.Dropdown(["Flutter", "Kotlin XML", "SwiftUI", "React Native"], label="Framework"),
        gr.Textbox(lines=4, label="UI Prompt", placeholder="e.g. Login screen with email & password, dark theme")
    ],
    outputs=gr.Code(label="Generated UI Code"),
    title="Prompt-to-UI Code Generator",
    description="Generate Android/iOS UI code in Flutter, SwiftUI, XML, or React Native by just describing the layout."
)

interface.launch()