import os
import gradio as gr
from huggingface_hub import InferenceClient

def get_hf_token():
    """Get the HF token from environment variables or Gradio secrets"""
    try:
        # Try Gradio secrets first (works on Spaces)
        return gr.secrets["HF_TOKEN"]
    except (AttributeError, KeyError):
        # Fall back to environment variable (works locally)
        return os.getenv("HF_TOKEN")

def generate_code(prompt):
    try:
        token = get_hf_token()
        if not token:
            return "Error: Missing Hugging Face token."
            
        client = InferenceClient(
            model="codellama/CodeLlama-34b-Instruct-hf",  # Smaller alternative
            token=token
        )
        
        system_prompt = """You are an expert Next.js 15 or higher with modern App Router, TypeScript and TailwindCSS developer. 
        Generate clean, efficient code following these rules:
        1. Use Next.js App Router
        2. Strict TypeScript typing
        3. Modern TailwindCSS classes
        4. Include proper error boundaries
        5. Add JSDoc comments"""
        
        response = client.chat_completion(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": prompt}
            ],
            max_tokens=1500,
            temperature=0.3
        )
        
        return response.choices[0].message.content
        
    except Exception as e:
        return f"Error generating code: {str(e)}"

# Gradio interface with improved layout
with gr.Blocks(
    title="Next.js Code Generator",
    theme=gr.themes.Soft(),
    css=".gradio-container {max-width: 900px !important}"
) as demo:
    
    gr.Markdown("## 🚀 Next.js + TypeScript + TailwindCSS Code Generator")
    
    with gr.Row():
        with gr.Column(scale=3):
            input_prompt = gr.Textbox(
                label="Describe your component",
                placeholder="Create a responsive navbar with dark mode toggle...",
                lines=5,
                max_lines=10
            )
            with gr.Row():
                submit_btn = gr.Button("Generate Code", variant="primary")
                clear_btn = gr.Button("Clear")
        
        with gr.Column(scale=7):
            output_code = gr.Code(
                label="Generated Component",
                language="typescript",
                interactive=False,
                lines=25
            )
    
    # Form submission handling
    submit_btn.click(
        fn=generate_code,
        inputs=input_prompt,
        outputs=output_code
    )
    
    clear_btn.click(
        fn=lambda: ("", ""),
        inputs=None,
        outputs=[input_prompt, output_code]
    )

# Launch with sharing enabled
if __name__ == "__main__":
    demo.launch(
        auth=None,  # Disable authentication for local development
        share=False  # Disable sharing when running locally
    )