from flask import Flask, request, render_template, send_file, Response
import ast
import io
import zipfile
from datetime import datetime
import base64

app = Flask(__name__)

# Expanded mapping of Gradio components to Flask-compatible HTML inputs
COMPONENT_MAP = {
    'gr.Textbox': {'html_type': 'input', 'attributes': {'type': 'text', 'class': 'form-control'}},
    'gr.MultimodalTextbox': {'html_type': 'input', 'attributes': {'type': 'file', 'multiple': 'true', 'class': 'form-control'}},
    'gr.Slider': {'html_type': 'input', 'attributes': {'type': 'range', 'class': 'form-range', 'min': '0', 'max': '100'}},
    'gr.Image': {'html_type': 'input', 'attributes': {'type': 'file', 'accept': 'image/*', 'class': 'form-control'}},
    'gr.Dropdown': {'html_type': 'select', 'attributes': {'class': 'form-select'}},
    'gr.Checkbox': {'html_type': 'input', 'attributes': {'type': 'checkbox', 'class': 'form-check-input'}},
    'gr.Radio': {'html_type': 'input', 'attributes': {'type': 'radio', 'class': 'form-check-input'}},
    'gr.Label': {'html_type': 'div', 'attributes': {'class': 'output-label'}},
    'gr.Audio': {'html_type': 'input', 'attributes': {'type': 'file', 'accept': 'audio/*', 'class': 'form-control'}},
    'gr.Video': {'html_type': 'input', 'attributes': {'type': 'file', 'accept': 'video/*', 'class': 'form-control'}},
    'gr.File': {'html_type': 'input', 'attributes': {'type': 'file', 'class': 'form-control'}},
    'gr.Number': {'html_type': 'input', 'attributes': {'type': 'number', 'class': 'form-control'}},
    'gr.ColorPicker': {'html_type': 'input', 'attributes': {'type': 'color', 'class': 'form-control'}},
    'gr.Markdown': {'html_type': 'div', 'attributes': {'class': 'markdown-output'}},
    'gr.Dataframe': {'html_type': 'textarea', 'attributes': {'class': 'form-control', 'readonly': 'true'}},
    'gr.HTML': {'html_type': 'div', 'attributes': {'class': 'html-output'}},
    'gr.Chatbot': {'html_type': 'div', 'attributes': {'class': 'chatbot-output'}},
    'gr.Button': {'html_type': 'button', 'attributes': {'type': 'submit', 'class': 'btn btn-primary'}},
    'gr.ClearButton': {'html_type': 'button', 'attributes': {'type': 'button', 'class': 'btn btn-secondary', 'onclick': 'clearChatbot()'}}
}

def parse_gradio_script(script):
    """Parse the Gradio script to extract components, layout, and events."""
    tree = ast.parse(script)
    components = {
        'inputs': [],
        'outputs': [],
        'functions': {},
        'layout': [],
        'is_blocks': False,
        'html_content': [],
        'events': {}
    }
    
    for node in ast.walk(tree):
        if isinstance(node, ast.With):
            for item in node.items:
                if isinstance(item.context_expr, ast.Call) and hasattr(item.context_expr.func, 'attr') and item.context_expr.func.attr == 'Blocks':
                    components['is_blocks'] = True
                    for body_node in node.body:
                        if isinstance(body_node, ast.Expr) and isinstance(body_node.value, ast.Call):
                            call = body_node.value
                            if hasattr(call.func, 'attr'):
                                comp_type = f"gr.{call.func.attr}"
                                if comp_type in COMPONENT_MAP:
                                    comp_name = None
                                    for kw in call.keywords:
                                        if kw.arg == '_js':
                                            continue
                                        if isinstance(kw.value, ast.Name):
                                            comp_name = kw.value.id
                                    if comp_name:
                                        components['layout'].append({'type': comp_type, 'name': comp_name})
                                    else:
                                        components['html_content'].append(ast.unparse(call))
                                elif call.func.attr in ['Row', 'Column', 'Group']:
                                    components['layout'].append({'type': f"gr.{call.func.attr}", 'children': []})
        
        if isinstance(node, ast.Assign):
            target = node.targets[0]
            if isinstance(node.value, ast.Call) and hasattr(node.value.func, 'attr') and node.value.func.attr == 'click':
                event_name = target.id if isinstance(target, ast.Name) else None
                if event_name:
                    fn = ast.unparse(node.value.args[0]) if node.value.args else None
                    inputs = [ast.unparse(arg) for arg in node.value.keywords[0].value.elts] if node.value.keywords else []
                    outputs = [ast.unparse(arg) for arg in node.value.keywords[1].value.elts] if len(node.value.keywords) > 1 else []
                    components['events'][event_name] = {'fn': fn, 'inputs': inputs, 'outputs': outputs}
                    components['functions'][fn] = fn
                    components['inputs'].extend(inputs)
                    components['outputs'].extend(outputs)
    
    components['inputs'] = list(dict.fromkeys(components['inputs']))
    components['outputs'] = list(dict.fromkeys(components['outputs']))
    return components

def generate_flask_app(script, components):
    """Generate Flask app code and frontend templates."""
    flask_code = f"""from flask import Flask, request, render_template
import io
from PIL import Image
import base64
import pandas as pd
import json

app = Flask(__name__)

# Original functions from Gradio script
{script.split('with gr.Blocks')[0].strip()}

@app.route('/', methods=['GET', 'POST'])
def index():
    result = {{}}
    chatbot_messages = request.form.get('chatbot_messages', '[]')
    chatbot_messages = json.loads(chatbot_messages) if chatbot_messages else []
    
    if request.method == 'POST':
"""
    
    for i, input_comp in enumerate(components['inputs']):
        comp_type = input_comp.split('(')[0]
        if comp_type == 'gr.MultimodalTextbox':
            flask_code += f"""
        files = request.files.getlist('input_{i}')
        if files:
            inputs_{i} = [file.read() for file in files]
        else:
            inputs_{i} = request.form.get('input_{i}', '')
"""
        elif comp_type == 'gr.Image':
            flask_code += f"""
        file = request.files.get('input_{i}')
        if file:
            inputs_{i} = Image.open(file)
        else:
            inputs_{i} = None
"""
        elif comp_type in ['gr.Audio', 'gr.Video', 'gr.File']:
            flask_code += f"""
        file = request.files.get('input_{i}')
        if file:
            inputs_{i} = file.read()
        else:
            inputs_{i} = None
"""
        elif comp_type == 'gr.Checkbox':
            flask_code += f"""
        inputs_{i} = request.form.get('input_{i}') == 'on'
"""
        elif comp_type == 'gr.Number':
            flask_code += f"""
        inputs_{i} = float(request.form.get('input_{i}', 0))
"""
        else:
            flask_code += f"""
        inputs_{i} = request.form.get('input_{i}', '')
"""

    for event_name, event in components['events'].items():
        flask_code += f"""
        if '{event_name}' in request.form:
            inputs = [{', '.join(f'inputs_{components["inputs"].index(inp)}' for inp in event['inputs'])}]
            result['{event_name}'] = {event['fn']}(*inputs)
            if isinstance(result['{event_name}'], list):
                chatbot_messages.extend(result['{event_name}'])
            else:
                chatbot_messages.append({{ 'role': 'assistant', 'content': str(result['{event_name}']) }})
"""

    flask_code += """
        if 'clear_chatbot' in request.form:
            chatbot_messages = []
"""

    flask_code += """
    return render_template('index.html', result=result, chatbot_messages=json.dumps(chatbot_messages))
"""

    html_template = """<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Do-it-All Flask App</title>
    <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0/dist/css/bootstrap.min.css" rel="stylesheet">
    <style>
        body { background-color: #f8f9fa; font-family: Arial, sans-serif; }
        .container { max-width: 1200px; margin: 20px auto; padding: 20px; background: white; border-radius: 10px; box-shadow: 0 0 10px rgba(0,0,0,0.1); }
        .form-group { margin-bottom: 20px; }
        .btn-submit { width: 100%; }
        .output { margin-top: 20px; padding: 15px; background: #e9ecef; border-radius: 5px; }
        .markdown-output { white-space: pre-wrap; }
        .html-output { margin: 10px 0; }
        .chatbot-output { height: 800px; overflow-y: auto; border: 1px solid #ddd; padding: 10px; }
        .btn_row { display: flex; gap: 10px; justify-content: center; }
        .new_btn { padding: 10px 20px; background: #007bff; color: white; border-radius: 5px; cursor: pointer; }
    </style>
</head>
<body>
    <div class="container">
"""
    
    for html in components['html_content']:
        html_content = html.split('gr.HTML(')[1].rstrip(')')
        html_content = html_content.strip('"').strip("'")
        html_template += f"""
        <div class="html-output">{html_content}</div>
"""

    html_template += """
        <form method="POST" enctype="multipart/form-data">
            <input type="hidden" name="chatbot_messages" value="{{ chatbot_messages }}">
"""
    
    for comp in components['layout']:
        comp_type = comp['type']
        if comp_type in ['gr.Row', 'gr.Column', 'gr.Group']:
            html_template += f"""
            <div class="{'row' if comp_type == 'gr.Row' else 'col' if comp_type == 'gr.Column' else 'group'}">
"""
        elif comp_type in COMPONENT_MAP:
            comp_info = COMPONENT_MAP[comp_type]
            attrs = ' '.join([f'{k}="{v}"' for k, v in comp_info['attributes'].items()])
            html_template += f"""
            <div class="form-group">
                <label for="{comp['name']}">{comp['name'].replace('_', ' ').title()}</label>
                <{comp_info['html_type']} name="{comp['name']}" id="{comp['name']}" {attrs}>
"""
            if comp_type == 'gr.Dropdown':
                html_template += """
                    <option value="option1">Option 1</option>
                    <option value="option2">Option 2</option>
                </select>
"""
            elif comp_type == 'gr.Button':
                html_template += f"""
                </{comp_info['html_type']}>
"""
            else:
                html_template += f"""
                </{comp_info['html_type']}>
"""
            html_template += """
            </div>
"""

    for event_name in components['events']:
        html_template += f"""
            <button type="submit" name="{event_name}" class="btn btn-primary btn-submit">Submit {event_name.replace('_', ' ').title()}</button>
"""

    html_template += """
            <button type="submit" name="clear_chatbot" class="btn btn-secondary">Clear Chatbot</button>
        </form>
"""

    html_template += """
        <div class="chatbot-output">
            {% for msg in json.loads(chatbot_messages) %}
                <p><strong>{{ msg.role }}:</strong> {{ msg.content }}</p>
            {% endfor %}
        </div>
    </div>
    <script src="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0/dist/js/bootstrap.bundle.min.js"></script>
    <script>
        function clearChatbot() {
            document.querySelector('input[name="chatbot_messages"]').value = '[]';
        }
    </script>
</body>
</html>
"""

    requirements = "flask\npillow\npandas\n"

    return flask_code, html_template, requirements

@app.route('/', methods=['GET', 'POST'])
def convert_gradio():
    result = None
    
    if request.method == 'POST':
        gradio_script = request.form['gradio_script']
        try:
            components = parse_gradio_script(gradio_script)
            flask_code, html_template, requirements = generate_flask_app(gradio_script, components)
            
            buffer = io.BytesIO()
            with zipfile.ZipFile(buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
                zip_file.writestr('app.py', flask_code)
                zip_file.writestr('templates/index.html', html_template)
                zip_file.writestr('requirements.txt', requirements)
            buffer.seek(0)
            zip_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
            
            result = {
                'flask_code': flask_code,
                'html_template': html_template,
                'requirements': requirements,
                'zip_base64': zip_base64
            }
        except Exception as e:
            return render_template('index.html', error=str(e), gradio_script=gradio_script)
    
    return render_template('index.html', result=result, gradio_script=request.form.get('gradio_script', ''))

@app.route('/download/<filename>')
def download_file(filename):
    if filename == 'app.py':
        content = request.args.get('content')
        return Response(content, mimetype='text/plain', headers={"Content-Disposition": f"attachment;filename={filename}"})
    elif filename == 'index.html':
        content = request.args.get('content')
        return Response(content, mimetype='text/html', headers={"Content-Disposition": f"attachment;filename={filename}"})
    elif filename == 'requirements.txt':
        content = request.args.get('content')
        return Response(content, mimetype='text/plain', headers={"Content-Disposition": f"attachment;filename={filename}"})
    elif filename == 'zip':
        zip_base64 = request.args.get('content')
        buffer = io.BytesIO(base64.b64decode(zip_base64))
        buffer.seek(0)
        return send_file(buffer, mimetype='application/zip', as_attachment=True, download_name=f'flask_app_{datetime.now().strftime("%Y%m%d_%H%M%S")}.zip')
    return "File not found", 404

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=7860, debug=True)