import re

import gradio as gr
from openai import OpenAI
from config import API_KEY, MODEL, SYSTEM_PROMPT, ENDPOINT, EXAMPLES, DEFAULT_LOCALE, DEFAULT_THEME

client = OpenAI(api_key=API_KEY, base_url=ENDPOINT)

react_imports = {
    "lucide-react": "https://esm.sh/lucide-react@0.525.0",
    "recharts": "https://esm.sh/recharts@3.1.0",
    "framer-motion": "https://esm.sh/framer-motion@12.23.6",
    "matter-js": "https://esm.sh/matter-js@0.20.0",
    "p5": "https://esm.sh/p5@2.0.3",
    "konva": "https://esm.sh/konva@9.3.22",
    "react-konva": "https://esm.sh/react-konva@19.0.7",
    "three": "https://esm.sh/three@0.178.0",
    "@react-three/fiber": "https://esm.sh/@react-three/fiber@9.2.0",
    "@react-three/drei": "https://esm.sh/@react-three/drei@10.5.2",
    "@tailwindcss/browser": "https://esm.sh/@tailwindcss/browser@4.1.11",
    "react": "https://esm.sh/react@^19.0.0",
    "react/": "https://esm.sh/react@^19.0.0/",
    "react-dom": "https://esm.sh/react-dom@^19.0.0",
    "react-dom/": "https://esm.sh/react-dom@^19.0.0/",
    "react-dom/client": "https://esm.sh/react-dom@^19.0.0/client"
}


class GradioEvents:

    @staticmethod
    def generate_code(input_value, system_prompt_input_value, state_value):

        def get_generated_files(text):
            patterns = {
                'html': r'```html\n(.+?)\n```',
                'jsx': r'```jsx\n(.+?)\n```',
                'tsx': r'```tsx\n(.+?)\n```',
            }
            result = {}

            for ext, pattern in patterns.items():
                matches = re.findall(pattern, text, re.DOTALL)
                if matches:
                    content = '\n'.join(matches).strip()
                    result[f'index.{ext}'] = content

            if len(result) == 0:
                result["index.html"] = text.strip()
            return result

        def process_react_imports(code):
            """Process React code to work with ES modules and import map"""
            # Remove export default and replace with const App
            processed_code = code.replace("export default", "const App =")
            
            # Process imports to use the import map
            import_lines = []
            code_lines = []
            
            for line in processed_code.split('\n'):
                line = line.strip()
                if line.startswith('import '):
                    # Extract import statement
                    import_lines.append(line)
                else:
                    code_lines.append(line)
            
            # Combine processed imports and code
            if import_lines:
                processed_code = '\n'.join(import_lines) + '\n\n' + '\n'.join(code_lines)
            else:
                processed_code = '\n'.join(code_lines)
            
            return processed_code

        yield {
            empty_output: gr.update(visible=False),
            loading_output: gr.update(visible=True),
            sandbox: gr.update(visible=False),
            output: gr.update(value=None)
        }

        if input_value is None:
            input_value = ''

        messages = [{
            'role': "system",
            "content": SYSTEM_PROMPT
            # 'content': system_prompt_input_value
        }] + state_value["history"]

        messages.append({'role': "user", 'content': input_value})

        generator = client.chat.completions.create(model=MODEL,
                                                   messages=messages,
                                                   stream=True)
        response = ""
        for chunk in generator:
            content = chunk.choices[0].delta.content
            if content is not None:
                response += content
            if chunk.choices[0].finish_reason == 'stop':
                state_value["history"] = messages + [{
                    'role': "assistant",
                    'content': response
                }]
                generated_files = get_generated_files(response)
                react_code = generated_files.get(
                    "index.tsx") or generated_files.get("index.jsx")
                html_code = generated_files.get("index.html")
                # Completed
                # Create HTML content for sandbox
                if react_code:
                    # For React code, create a complete HTML page that renders React with ES modules
                    processed_react_code = process_react_imports(react_code)
                    escaped_react_code = processed_react_code.replace("`", "`").replace("'", "'")
                    
                    # Create import map from react_imports
                    import_map = '{\n    "imports": {\n'
                    for package, url in react_imports.items():
                        import_map += f'      "{package}": "{url}",\n'
                    import_map = import_map.rstrip(',\n') + '\n    }\n  }'
                    
                    sandbox_content = f"""
                    <iframe style="width: 100%; height: 600px; border: 1px solid #ddd; border-radius: 8px;" 
                            srcdoc='<!DOCTYPE html>
                    <html>
                    <head>
                        <meta charset="UTF-8">
                        <meta name="viewport" content="width=device-width, initial-scale=1.0">
                        <script type="importmap">
                          {import_map}
                        </script>

                    </head>
                    <body>
                        <div id="root"></div>
                        <script type="module">
                            import React from "react";
                            import { createRoot } from "react-dom/client";
                            
                            {escaped_react_code}
                            
                            const root = createRoot(document.getElementById("root"));
                            root.render(React.createElement(App));
                        </script>
                    </body>
                    </html>'>
                    </iframe>
                    """
                else:
                    # For HTML code, just display it in iframe
                    escaped_html_code = html_code.replace("'", "&#39;")
                    sandbox_content = f"""
                    <iframe style="width: 100%; height: 600px; border: 1px solid #ddd; border-radius: 8px;" 
                            srcdoc='{escaped_html_code}'>
                    </iframe>
                    """
                
                yield {
                    output: gr.update(value=response),
                    download_content: gr.update(value=react_code or html_code),
                    empty_output: gr.update(visible=False),
                    loading_output: gr.update(visible=False),
                    sandbox: gr.update(value=sandbox_content, visible=True),
                    download_btn: gr.update(interactive=True),
                    state: gr.update(value=state_value)
                }

            else:
                # Generating
                yield {
                    output: gr.update(value=response),
                }

    @staticmethod
    def select_example(example: dict):
        return lambda: gr.update(value=example["description"])

    @staticmethod
    def toggle_visibility(visible_state):
        return gr.update(visible=not visible_state)

    @staticmethod
    def disable_btns(btns: list):
        return lambda: [gr.update(disabled=True) for _ in btns]

    @staticmethod
    def enable_btns(btns: list):
        return lambda: [gr.update(disabled=False) for _ in btns]

    @staticmethod
    def update_system_prompt(system_prompt_input_value, state_value):
        state_value["system_prompt"] = system_prompt_input_value
        return gr.update(value=state_value)

    @staticmethod
    def reset_system_prompt(state_value):
        return gr.update(value=state_value["system_prompt"])

    @staticmethod
    def render_history(statue_value):
        return gr.update(value=statue_value["history"])

    @staticmethod
    def clear_history(state_value):
        gr.Success("History Cleared.")
        state_value["history"] = []
        return gr.update(value=state_value)


css = """
.output-empty, .output-loading {
  display: flex;
  flex-direction: column;
  align-items: center;
  justify-content: center;
  width: 100%;
  min-height: 680px;
}

.output-html {
  display: flex;
  flex-direction: column;
  width: 100%;
  min-height: 680px;
  max-height: 1200px;
}

.output-html > iframe {
  flex: 1;
  width: 100% !important;
  margin: 0 !important;
  padding: 0 !important;
}

.output-code {
  flex: 1;
  min-height: 100%;
}

.gradio-container {
  max-width: none !important;
  width: 100% !important;
}
"""

with gr.Blocks(css=css, title="GPT5 Vibe Tester") as demo:
    # Global State
    state = gr.State({"system_prompt": "", "history": []})
    
    # Header
    with gr.Row():
        with gr.Column():
            gr.HTML("""
                <div style="text-align: center; margin: 20px 0;">
                    <div style="font-size: 64px; margin-bottom: 10px;">🎯</div>
                    <h1 style="color: #6A57FF; margin: 10px 0;">GPT5 Vibe Tester</h1>
                    <p style="color: #666; margin: 5px 0;">Test your creative vibes with GPT5</p>
                </div>
            """)
    
    with gr.Row():
        # Left Column - Input and Controls
        with gr.Column(scale=1):
            # Input
            input = gr.Textbox(
                lines=4,
                placeholder="Describe the idea you want to test",
                label="Your Idea",
                elem_id="input-container"
            )
            
            # Note
            gr.HTML("""
                <p style="color: #ff7f00; font-weight: bold; margin: 10px 0;">
                    <strong>Note:</strong> Test your vibes through interactive prototypes. The app supports multi-round dialogue to refine your ideas.
                </p>
            """)
            
            # Submit Button
            submit_btn = gr.Button(
                "Submit",
                variant="primary",
                size="lg",
                elem_id="submit-btn"
            )
            
            # Settings
            gr.HTML("<hr style='margin: 20px 0;'><h3>Settings</h3>")
            
            with gr.Row():
                history_btn = gr.Button("📜 View History", elem_id="history-btn")
                clear_history_btn = gr.Button("🧹 Clear History", variant="stop")
            
            # Examples
            gr.HTML("<hr style='margin: 20px 0;'><h3>Examples</h3>")
            
            # Create example buttons
            example_buttons = []
            for i, example in enumerate(EXAMPLES):
                btn = gr.Button(
                    f"{example['title']}", 
                    variant="secondary",
                    size="sm"
                )
                example_buttons.append(btn)
                # Set up click event for each example
                btn.click(
                    fn=lambda desc=example['description']: desc,
                    outputs=[input]
                )

        # Right Column - Output
        with gr.Column(scale=2):
            with gr.Group():
                gr.HTML("<h3>Output</h3>")
                
                with gr.Row():
                    download_btn = gr.Button("📥 Download Code", variant="secondary", interactive=False)
                    view_code_btn = gr.Button("🧑‍💻 View Code", variant="primary")
                
                # Output tabs using visibility
                output_state = gr.State("empty")  # "empty", "loading", "render"
                
                # Empty state
                empty_output = gr.HTML(
                    """
                    <div class="output-empty">
                        <div style="text-align: center; color: #666;">
                            <h3>🚀 Ready to Generate</h3>
                            <p>Enter your request to generate code</p>
                        </div>
                    </div>
                    """,
                    visible=True
                )
                
                # Loading state  
                loading_output = gr.HTML(
                    """
                    <div class="output-loading">
                        <div style="text-align: center; color: #666;">
                            <h3>⚡ Generating code...</h3>
                            <p>Please wait while I create your application</p>
                        </div>
                    </div>
                    """,
                    visible=False
                )
                
                # Render state
                sandbox = gr.HTML(
                    "",
                    elem_classes="output-html",
                    visible=False
                )
    
    # Hidden components for functionality
    download_content = gr.Text(visible=False)
    output_loading = gr.State(False)
    state_tab = gr.State("empty")
    
    # Modals using Gradio components
    with gr.Row(visible=False) as system_prompt_modal:
        with gr.Column():
            gr.HTML("<h3>System Prompt</h3>")
            system_prompt_input = gr.Textbox(
                lines=8,
                placeholder="Enter your system prompt here",
                label="System Prompt"
            )
    
    # Code viewer (will be shown/hidden)
    with gr.Row(visible=False) as output_code_drawer:
        with gr.Column():
            gr.HTML("<h3>Generated Code</h3>")
            output = gr.Markdown(label="Code Output")
    
    # History viewer
    with gr.Row(visible=False) as history_drawer:
        with gr.Column():
            gr.HTML("<h3>Chat History</h3>")
            history_output = gr.Chatbot(
                show_label=False,
                type="messages",
                height=400
            )
    # Event Handlers
    
    # Clear history
    clear_history_btn.click(
        fn=GradioEvents.clear_history,
        inputs=[state],
        outputs=[state]
    )
    
    # Show history
    history_btn.click(
        fn=lambda: gr.update(visible=True),
        outputs=[history_drawer]
    ).then(
        fn=GradioEvents.render_history,
        inputs=[state],
        outputs=[history_output]
    )
    
    # Download code
    download_btn.click(
        fn=None,
        inputs=[download_content],
        js="""(content) => {
            const blob = new Blob([content], { type: 'text/plain' })
            const url = URL.createObjectURL(blob)
            const a = document.createElement('a')
            a.href = url
            a.download = 'generated_code.txt'
            a.click()
        }"""
    )
    
    # View code
    view_code_btn.click(
        fn=lambda: gr.update(visible=True),
        outputs=[output_code_drawer]
    )
    
    # Submit button
    submit_btn.click(
        fn=lambda: [gr.update(interactive=False), gr.update(interactive=False)],
        outputs=[submit_btn, download_btn]
    ).then(
        fn=GradioEvents.generate_code,
        inputs=[input, system_prompt_input, state],
        outputs=[
            output, empty_output, loading_output, sandbox, 
            download_content, download_btn, state
        ]
    ).then(
        fn=lambda: [gr.update(interactive=True), gr.update(interactive=True)],
        outputs=[submit_btn, download_btn]
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=100,
               max_size=100).launch(ssr_mode=False, max_threads=100)