app.py

# -*- coding: utf-8 -*-
import os
import warnings
import io
import zipfile
import tempfile
from pathlib import Path

import gradio as gr

# ----------------------------------------------------------------------
# Suppress noisy warnings in the Space environment
# ----------------------------------------------------------------------
warnings.filterwarnings("ignore", category=UserWarning)

# ----------------------------------------------------------------------
# Groq client initialization (reads API key from environment)
# ----------------------------------------------------------------------
groq_api_key = os.getenv("GROQ_API_KEY")
groq_client = None

if groq_api_key:
    try:
        from groq import Groq
        groq_client = Groq(api_key=groq_api_key)
    except Exception as e:
        print(f"Failed to import or initialise Groq client: {e}")
        groq_client = None
else:
    print("GROQ_API_KEY not set. The app will work in demo mode with placeholder messages.")

# ----------------------------------------------------------------------
# Supported languages (common set)
# ----------------------------------------------------------------------
SUPPORTED_LANGUAGES = [
    "Python",
    "JavaScript",
    "Java",
    "C++",
    "C#",
    "Go",
    "Ruby",
    "PHP",
    "TypeScript",
    "HTML",
    "CSS",
    "Swift",
    "Kotlin",
    "Rust",
    "Scala",
    "Perl",
    "R",
    "Shell",
]

# ----------------------------------------------------------------------
# Function: improve existing code
# ----------------------------------------------------------------------
def improve_code(original_code: str, language: str) -> str:
    """
    Sends the original code to the LLM and returns an improved version.
    """
    if not original_code.strip():
        return "Please provide some code to improve."

    if not groq_client:
        return "No Groq API key configured. Set the GROQ_API_KEY environment variable."

    prompt = (
        f"You are a senior software engineer. Improve the following {language} code for "
        "readability, efficiency and best practices. Return only the improved code "
        "without any explanation.\n\n"
        f"Original code:\n{original_code}"
    )

    try:
        response = groq_client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.2,
            max_tokens=2000,
        )
        improved = response.choices[0].message.content
        return improved.strip() if improved else "No content returned from Groq."
    except Exception as e:
        return f"Error communicating with Groq: {e}"

# ----------------------------------------------------------------------
# Helper: create a zip file with given files
# ----------------------------------------------------------------------
def _create_zip(file_dict: dict) -> str:
    """
    file_dict: mapping of relative file path -> file content (string)
    Returns the path to the created zip file.
    """
    tmp_dir = tempfile.mkdtemp()
    zip_path = Path(tmp_dir) / "generated_project.zip"

    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
        for rel_path, content in file_dict.items():
            zipf.writestr(rel_path, content)

    return str(zip_path)

# ----------------------------------------------------------------------
# Function: generate full web project from an idea
# ----------------------------------------------------------------------
def generate_web_code(idea: str, language: str):
    """
    Generates a complete web project (HTML/CSS/JS or a single-file script)
    based on the user's idea and the selected language.
    Returns a tuple: (preview_text, zip_file_path)
    """
    if not idea.strip():
        return ("Please describe the web idea you want to build.", None)

    if not groq_client:
        return ("No Groq API key configured. Set the GROQ_API_KEY environment variable.", None)

    prompt = (
        f"You are a senior full-stack developer. Based on the following description, "
        f"write a complete, runnable {language} web application. Include all necessary "
        "files (HTML, CSS, JavaScript, or a single script) and comment the code. "
        "Also provide a minimal requirements.txt, a README.md explaining how to run the project, "
        "and a Dockerfile that builds and runs the app. Return the files in the following format:\n"
        "=== filename ===\n"
        "<file content>\n"
        "=== filename ===\n"
        "<file content>\n"
        "Do not add any extra explanation.\n\n"
        f"Description:\n{idea}"
    )

    try:
        response = groq_client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3,
            max_tokens=4000,
        )
        generated = response.choices[0].message.content
        if not generated:
            return ("No content returned from Groq.", None)

        # Parse the generated sections
        sections = {}
        current_file = None
        for line in generated.splitlines():
            if line.startswith("=== ") and line.endswith(" ==="):
                current_file = line.strip("= ").strip()
                sections[current_file] = ""
            elif current_file:
                sections[current_file] += line + "\n"

        # Determine a preview file (HTML, Python or JS if present)
        preview_key = None
        for key in sections:
            if key.lower().endswith((".html", ".py", ".js")):
                preview_key = key
                break
        preview_text = sections.get(preview_key, "Generated files:\n" + "\n".join(sections.keys()))

        # Create zip with all files
        zip_path = _create_zip(sections)

        return (preview_text, zip_path)

    except Exception as e:
        return (f"Error communicating with Groq: {e}", None)

# ----------------------------------------------------------------------
# Gradio UI
# ----------------------------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("# Code Improver & Web Generator")

    with gr.Tabs():
        # --------------------------------------------------------------
        # Tab 1: Improve existing code
        # --------------------------------------------------------------
        with gr.TabItem("Improve Code"):
            with gr.Row():
                with gr.Column():
                    lang_select_improve = gr.Dropdown(
                        choices=SUPPORTED_LANGUAGES,
                        label="Language",
                        value="Python",
                    )
                    code_input = gr.Textbox(
                        label="Original Code",
                        placeholder="Paste your code here",
                        lines=15,
                    )
                with gr.Column():
                    improved_output = gr.Textbox(
                        label="Improved Code",
                        placeholder="Improved code will appear here",
                        lines=15,
                    )
            improve_btn = gr.Button("Improve")
            improve_btn.click(
                fn=improve_code,
                inputs=[code_input, lang_select_improve],
                outputs=improved_output,
            )

        # --------------------------------------------------------------
        # Tab 2: Generate web project from idea
        # --------------------------------------------------------------
        with gr.TabItem("Generate Web"):
            with gr.Row():
                with gr.Column():
                    lang_select_generate = gr.Dropdown(
                        choices=SUPPORTED_LANGUAGES,
                        label="Target Language",
                        value="HTML",
                    )
                    idea_input = gr.Textbox(
                        label="Web Idea",
                        placeholder="Describe the web page or app you want",
                        lines=12,
                    )
                with gr.Column():
                    generated_preview = gr.Textbox(
                        label="Generated Files Preview",
                        placeholder="The generated code will appear here",
                        lines=12,
                    )
                    download_file = gr.File(
                        label="Download Project Zip",
                        type="filepath",
                    )
            generate_btn = gr.Button("Generate")
            generate_btn.click(
                fn=generate_web_code,
                inputs=[idea_input, lang_select_generate],
                outputs=[generated_preview, download_file],
            )

demo.launch()