import gradio as gr
import requests
import re
import tempfile
import importlib.util
import sys
import os
import ast

def call_openai_api(api_key, prompt):
    """Direct API call to OpenAI"""
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    
    # Determine Gradio version to ensure compatibility
    gradio_version = gr.__version__
    major_version = int(gradio_version.split('.')[0])
    
    system_prompt = f"""You are an expert Gradio developer. 
Create a standalone Gradio application based on the user's prompt.
Your response should ONLY include Python code without any explanation.

IMPORTANT: You must use Gradio version {gradio_version} (Gradio {major_version}.x) syntax.
For Gradio 3.x: Use the `.click()` method on buttons and elements to connect them to functions.
For Gradio 4.x: Use the event subscription pattern instead of .click().

The code must:
1. Import all necessary libraries
2. Define a complete, functional Gradio interface
3. DO NOT include a launch command or if __name__ == "__main__" block
4. The interface should be assigned to a variable named 'demo'
5. Handle errors gracefully
6. Be completely self-contained in a single script

Example for Gradio 3.x:
```python
import gradio as gr

def greet(name):
    return f"Hello, {name}!"

with gr.Blocks() as demo:
    name_input = gr.Textbox(label="Your Name")
    greet_btn = gr.Button("Greet")
    output = gr.Textbox(label="Output")
    
    greet_btn.click(fn=greet, inputs=name_input, outputs=output)
```
"""
    
    data = {
        "model": "gpt-4o",
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.2,
        "max_tokens": 4000
    }
    
    response = requests.post(
        "https://api.openai.com/v1/chat/completions",
        headers=headers,
        json=data
    )
    
    if response.status_code != 200:
        return None, f"API Error: {response.status_code} - {response.text}"
    
    result = response.json()
    return result["choices"][0]["message"]["content"], None

def extract_code_blocks(text):
    """Extract code blocks from markdown"""
    pattern = r'```(?:python)?\s*([\s\S]*?)```'
    matches = re.findall(pattern, text)
    
    if not matches and text.strip():
        if re.search(r'import\s+\w+|def\s+\w+\(|class\s+\w+:|if\s+__name__\s*==\s*[\'"]__main__[\'"]:', text):
            return [text.strip()]
    
    return [match.strip() for match in matches]

def validate_gradio_code(code):
    """Basic validation of the generated code"""
    try:
        tree = ast.parse(code)
        
        allowed_modules = ['gradio', 'numpy', 'pandas', 'matplotlib', 'PIL', 'os', 'io', 'base64', 
                          'time', 'datetime', 'json', 'random', 'math', 'sys', 're', 'pathlib',
                          'collections', 'typing', 'warnings']
        
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for name in node.names:
                    if name.name not in allowed_modules:
                        return False, f"Unauthorized import: {name.name}"
            
            elif isinstance(node, ast.ImportFrom):
                if node.module not in allowed_modules and node.module is not None:
                    return False, f"Unauthorized import from: {node.module}"
        
        return True, None
        
    except SyntaxError as e:
        return False, f"Syntax error in the code: {str(e)}"
    except Exception as e:
        return False, f"Error validating code: {str(e)}"

def load_generated_app(code):
    """Load and run the generated Gradio app"""
    # Save code to temp file
    with tempfile.NamedTemporaryFile(suffix='.py', delete=False) as f:
        f.write(code.encode('utf-8'))
        temp_file = f.name
    
    try:
        # Import module
        module_name = os.path.basename(temp_file).replace('.py', '')
        spec = importlib.util.spec_from_file_location(module_name, temp_file)
        module = importlib.util.module_from_spec(spec)
        sys.modules[module_name] = module
        spec.loader.exec_module(module)
        
        # Get the Gradio interface
        if hasattr(module, 'demo'):
            return module.demo, None
        else:
            return None, "No 'demo' variable found in the generated code"
            
    except Exception as e:
        return None, f"Error: {str(e)}"
    finally:
        try:
            os.unlink(temp_file)
        except:
            pass

def generate_app(api_key, prompt):
    """Main function to generate and load a Gradio app"""
    # Validate API key format
    if not api_key or len(api_key) < 20 or not api_key.startswith("sk-"):
        return None, "Please provide a valid OpenAI API key"
    
    # Call OpenAI API
    response, error = call_openai_api(api_key, prompt)
    if error:
        return None, error
    
    # Extract code blocks
    code_blocks = extract_code_blocks(response)
    if not code_blocks:
        return None, "No valid code found in the response"
    
    code = code_blocks[0]
    
    # Validate code
    is_valid, error_msg = validate_gradio_code(code)
    if not is_valid:
        return code, error_msg
    
    # Load the app
    app, app_error = load_generated_app(code)
    if app_error:
        return code, app_error
    
    return code, app

# Example Gradio apps for common requests
EXAMPLE_APPS = {
    "calculator": """
import gradio as gr
import numpy as np

def calculate(num1, num2, operation):
    if operation == "Add":
        return num1 + num2
    elif operation == "Subtract":
        return num1 - num2
    elif operation == "Multiply":
        return num1 * num2
    elif operation == "Divide":
        if num2 == 0:
            return "Error: Division by zero"
        return num1 / num2
    else:
        return "Please select an operation"

with gr.Blocks() as demo:
    gr.Markdown("# Simple Calculator")
    
    with gr.Row():
        num1_input = gr.Number(label="First Number")
        num2_input = gr.Number(label="Second Number")
        
    op_dropdown = gr.Dropdown(
        ["Add", "Subtract", "Multiply", "Divide"], 
        label="Operation"
    )
    
    calculate_btn = gr.Button("Calculate")
    result = gr.Number(label="Result")
    
    calculate_btn.click(
        fn=calculate,
        inputs=[num1_input, num2_input, op_dropdown],
        outputs=result
    )
""",
    "image_filter": """
import gradio as gr
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter

def apply_filter(image, filter_type, intensity):
    if image is None:
        return None
    
    img = Image.fromarray(image)
    
    if filter_type == "Blur":
        filtered = img.filter(ImageFilter.GaussianBlur(radius=intensity))
    elif filter_type == "Sharpen":
        enhancer = ImageEnhance.Sharpness(img)
        filtered = enhancer.enhance(1 + intensity)
    elif filter_type == "Brightness":
        enhancer = ImageEnhance.Brightness(img)
        filtered = enhancer.enhance(1 + intensity)
    elif filter_type == "Contrast":
        enhancer = ImageEnhance.Contrast(img)
        filtered = enhancer.enhance(1 + intensity)
    else:
        filtered = img
    
    return np.array(filtered)

with gr.Blocks() as demo:
    gr.Markdown("# Image Filter App")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Original Image")
            filter_type = gr.Dropdown(
                ["Blur", "Sharpen", "Brightness", "Contrast"],
                label="Filter Type",
                value="Blur"
            )
            intensity = gr.Slider(0, 2, 0.5, label="Intensity")
            apply_btn = gr.Button("Apply Filter")
        
        with gr.Column():
            output_image = gr.Image(label="Filtered Image")
    
    apply_btn.click(
        fn=apply_filter,
        inputs=[input_image, filter_type, intensity],
        outputs=output_image
    )
"""
}

# Create the main UI
with gr.Blocks(title="AI Gradio App Generator") as demo:
    gr.Markdown("# 🤖 AI Gradio App Generator")
    gr.Markdown("Describe the Gradio app you want, and I'll generate it using OpenAI's API.")
    
    with gr.Row():
        with gr.Column():
            api_key = gr.Textbox(
                label="OpenAI API Key", 
                placeholder="sk-...",
                type="password",
                info="Your key is used only for this session"
            )
            
            prompt = gr.Textbox(
                label="App Description", 
                placeholder="Describe the Gradio app you want to create...",
                lines=5
            )
            
            submit_btn = gr.Button("Generate App", variant="primary")
            
            with gr.Accordion("Generated Code", open=False):
                code_output = gr.Code(language="python", label="Generated Code")
            
            status_output = gr.Markdown("")
    
    generated_app_container = gr.Group(visible=False)
    
    def on_submit(api_key_input, prompt_text):
        try:
            # If it's a simple request for a common app, use our built-in example
            lower_prompt = prompt_text.lower()
            for key, example_code in EXAMPLE_APPS.items():
                if key in lower_prompt:
                    try:
                        app, app_error = load_generated_app(example_code)
                        if app_error:
                            return example_code, f"⚠️ {app_error}", gr.update(visible=False)
                        return example_code, "✅ App generated successfully (from template)!", gr.update(visible=True, value=app)
                    except Exception as e:
                        pass  # Fall back to API generation if template fails
            
            # Generate custom app via API
            code, result = generate_app(api_key_input, prompt_text)
            
            if code is None:
                return None, f"⚠️ {result}", gr.update(visible=False)
            
            if isinstance(result, str):  # Error message
                return code, f"⚠️ {result}", gr.update(visible=False)
            
            # Success - result is the app
            return code, "✅ App generated successfully!", gr.update(visible=True, value=result)
        except Exception as e:
            return None, f"⚠️ Error: {str(e)}", gr.update(visible=False)
    
    submit_btn.click(
        on_submit,
        inputs=[api_key, prompt],
        outputs=[code_output, status_output, generated_app_container]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)