# Updated UI Generator with Better Open-Source Model and DartPad/Expo Preview Support

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

# Use a smaller model for CPU support
model_id = "deepseek-ai/deepseek-coder-1.3b-instruct"

# Use CPU or GPU if available
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 generation pipeline
generator = TextGenerationPipeline(model=model, tokenizer=tokenizer, device=device)

def generate_ui(platform, framework, ui_prompt):
    prompt_template = f"""
You are a professional mobile developer.
Create a full {framework} UI code for {platform} based on this request:
"""
    full_prompt = prompt_template + f"\n{ui_prompt}\n" + "\nEnsure to include imports, widgets, and app structure."
    response = generator(full_prompt, max_new_tokens=512, do_sample=True, temperature=0.7)[0]['generated_text']

    # Extract only code block (removes echoed prompt if needed)
    code_match = re.search(r"(?s)(import .*?)$", response)
    code = code_match.group(1).strip() if code_match else response.strip()
    return code, generate_preview_link(code, framework)

def generate_preview_link(code, framework):
    if framework == "Flutter":
        return f"https://dartpad.dev/?id=custom&code={gr.utils.sanitize_html(code)[:1500]}"
    elif framework == "React Native":
        return f"https://snack.expo.dev/?code={gr.utils.sanitize_html(code)[:1500]}"
    else:
        return "Preview not available for this framework."

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"),
        gr.Textbox(label="Preview Link")
    ],
    title="Prompt-to-UI Code Generator",
    description="Describe your screen and get UI code + preview for Flutter/React Native apps."
)

interface.launch()