# frontend_agent/ui_generator.py
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

MODEL_NAME = "facebook/opt-1.3b"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

def generate_react_component_llm(task_name):
    """
    Uses LLM to generate a React functional component for a frontend task.
    """
    prompt = f"""
Generate a React functional component for this frontend task.
Task: {task_name}
Include props if needed and simple placeholder JSX.
"""
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=250)
    code = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Remove prompt part if included
    if "Task:" in code:
        code = code.split("Task:")[-1].strip()
    return code