import gradio as gr
from transformers import pipeline, AutoTokenizer

# Load the StarChat beta model
MODEL_NAME = "HuggingFaceH4/starchat-beta"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
generator = pipeline(
    "text-generation",
    model=MODEL_NAME,
    tokenizer=tokenizer,
    device="cpu",  # Free tier uses CPU
)

def generate_code(natural_language_input):
    try:
        # Create a proper prompt for the model
        chat_prompt = [
            {"role": "system", "content": "You are a helpful AI assistant that generates Python code based on natural language descriptions."},
            {"role": "user", "content": f"Write Python code that: {natural_language_input}"}
        ]
        
        # Generate the prompt for the model
        prompt = tokenizer.apply_chat_template(chat_prompt, tokenize=False, add_generation_prompt=True)
        
        # Generate code with appropriate parameters for code generation
        generated = generator(
            prompt,
            max_new_tokens=256,
            temperature=0.2,
            top_p=0.95,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
        
        # Extract and clean the generated code
        full_response = generated[0]['generated_text']
        code_response = full_response.replace(prompt, "").strip()
        
        # Sometimes the model adds non-code text - we'll try to extract just the code blocks
        if "```python" in code_response:
            code_response = code_response.split("```python")[1].split("```")[0]
        elif "```" in code_response:
            code_response = code_response.split("```")[1].split("```")[0]
        
        return code_response
    except Exception as e:
        return f"Error generating code: {str(e)}"

# Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🚀 Text to Python Code Generator
    ### Powered by StarChat Beta
    Describe what you want the code to do in natural language, and get Python code!
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            input_text = gr.Textbox(
                label="Describe your code task",
                placeholder="e.g., 'create a Flask web server with two routes'",
                lines=3,
                max_lines=6
            )
            with gr.Row():
                generate_btn = gr.Button("Generate Code", variant="primary")
                clear_btn = gr.Button("Clear")
            
            gr.Markdown("**Tips:** Be specific in your description for better results.")
            
        with gr.Column(scale=4):
            output_code = gr.Code(
                label="Generated Python Code",
                language="python",
                interactive=True,
                lines=10
            )
            with gr.Row():
                copy_btn = gr.Button("Copy to Clipboard")
                refresh_btn = gr.Button("Try Again")
    
    # Examples
    examples = gr.Examples(
        examples=[
            ["create a function to calculate factorial recursively"],
            ["write code to download a file from URL and save it locally"],
            ["create a pandas DataFrame from a dictionary and plot a bar chart"],
            ["implement a simple REST API using FastAPI with GET and POST endpoints"]
        ],
        inputs=input_text,
        label="Click any example to try it"
    )
    
    # Button actions
    generate_btn.click(
        fn=generate_code,
        inputs=input_text,
        outputs=output_code
    )
    
    clear_btn.click(
        fn=lambda: ("", ""),
        inputs=None,
        outputs=[input_text, output_code]
    )
    
    refresh_btn.click(
        fn=generate_code,
        inputs=input_text,
        outputs=output_code
    )
    
    copy_btn.click(
        fn=lambda code: gr.Clipboard().copy(code),
        inputs=output_code,
        outputs=None,
        api_name="copy_code"
    )

demo.launch(debug=True)