import os
import importlib.util
import sys
import tempfile
import re
import ast
import gradio as gr
from openai import OpenAI

# Extract code blocks from text
def extract_code_blocks(text):
    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]

# Validate that the code is safe to run
def validate_gradio_code(code):
    try:
        tree = ast.parse(code)
        
        # Check imports
        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}"
        
        # Check for harmful operations
        code_str = code.lower()
        harmful_ops = ['subprocess', 'system(', 'popen(', 'execve(', 'chmod(', 'rmdir(', 
                      'remove(', 'unlink(', 'rmtree(', 'socket', 'eval(', 'exec(']
        
        for op in harmful_ops:
            if op in code_str:
                return False, f"Potentially harmful operation detected: {op}"
        
        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)}"

# Global variables for app storage
generated_app = None
current_code = ""

# Generate app using OpenAI
def generate_gradio_app(api_key, prompt):
    if not api_key or len(api_key) < 20:
        return None, "Please provide a valid OpenAI API key"
    
    try:
        client = OpenAI(api_key=api_key)
        
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": """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.
                The code must:
                1. Import all necessary libraries
                2. Define a complete, functional Gradio interface
                3. Launch the interface with share=False
                4. Handle errors gracefully
                5. Be completely self-contained in a single script
                6. End with a simple if __name__ == "__main__": block that launches the app
                """
                },
                {"role": "user", "content": prompt}
            ],
            temperature=0.2,
            max_tokens=4000
        )
        
        generated_code = response.choices[0].message.content
        code_blocks = extract_code_blocks(generated_code)
        
        if code_blocks:
            return code_blocks[0], None
        else:
            return generated_code, None
            
    except Exception as e:
        return None, str(e)

# Load and run generated app
def load_and_run_gradio_app(code):
    global generated_app, current_code
    
    # Validate code
    is_valid, error_msg = validate_gradio_code(code)
    if not is_valid:
        return None, error_msg
    
    # Clean up previous app
    if generated_app:
        try:
            generated_app.close()
        except:
            pass
    
    # 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)
        
        # Find Gradio interface
        for attr_name in dir(module):
            attr = getattr(module, attr_name)
            if isinstance(attr, (gr.Blocks, gr.Interface)):
                generated_app = attr
                current_code = code
                return attr, None
        
        return None, "No Gradio interface found in the generated code"
        
    except Exception as e:
        return None, f"Error executing the generated code: {str(e)}"
    finally:
        try:
            os.unlink(temp_file)
        except:
            pass

# Create the main UI
with gr.Blocks(title="AI Gradio App Generator") as demo:
    gr.Markdown("# 🤖 AI Gradio App Generator")
    gr.Markdown("Describe the app you want, and I'll generate it for you using OpenAI's API.")
    
    with gr.Row():
        with gr.Column(scale=2):
            api_key = gr.Textbox(
                label="OpenAI API Key", 
                placeholder="sk-...",
                type="password",
                info="Your key is used for this session only and not stored"
            )
            
            prompt = gr.Textbox(
                label="App Description", 
                placeholder="Describe the Gradio app you want to create...",
                lines=5
            )
            
            with gr.Row():
                submit_btn = gr.Button("Generate App", variant="primary")
                clear_btn = gr.Button("Clear", variant="secondary")
            
            with gr.Accordion("Generated Code", open=False):
                code_output = gr.Code(language="python", label="Generated Code")
            
            status_output = gr.Markdown("")
        
        with gr.Column(scale=3):
            app_container = gr.HTML("<div style='text-align: center; padding: 50px;'><h3>Your generated app will appear here</h3></div>")
    
    def on_submit(api_key_input, prompt_text):
        # Generate code
        code, error = generate_gradio_app(api_key_input, prompt_text)
        if error:
            return None, f"⚠️ Error generating code: {error}", "<div style='text-align: center; padding: 50px;'><h3>Error generating app</h3></div>"
        
        # Load and run app
        app, run_error = load_and_run_gradio_app(code)
        if run_error:
            return code, f"⚠️ Error running the generated app: {run_error}", "<div style='text-align: center; padding: 50px;'><h3>Error running app</h3></div>"
        
        # Create iframe to display app
        iframe_html = f'''
        <div style="border: 1px solid #ddd; border-radius: 8px; padding: 0; overflow: hidden;">
            <iframe id="appFrame" src="/generated_app" width="100%" height="800px" frameborder="0"></iframe>
        </div>
        '''
        return code, "✅ App generated successfully!", iframe_html
    
    def on_clear():
        return "", "", "<div style='text-align: center; padding: 50px;'><h3>Your generated app will appear here</h3></div>"
    
    submit_btn.click(
        on_submit,
        inputs=[api_key, prompt],
        outputs=[code_output, status_output, app_container]
    )
    
    clear_btn.click(
        on_clear,
        inputs=[],
        outputs=[prompt, status_output, app_container]
    )

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