app.py

import gradio as gr
import json
import pandas as pd
import random
import spaces
import io
from PIL import Image, ImageDraw

# ==========================================
# 1. FALLBACK / MOCK ENGINE (When no API Key is provided)
# ==========================================
MOCK_SCHEMAS = {
    "retail": {
        "problem_domain": "Tabular Demand Forecasting",
        "technical_summary": "Uses regression techniques to estimate future weekly store demand metrics based on weather and marketing variables.",
        "recommended_architecture": "XGBoost Regressor",
        "confidence_score": 0.95,
        "fallback_assumptions_made": "",
        "inputs": [
            {"name": "weekly_marketing_spend", "type": "numeric", "range": [1000.0, 50000.0], "description": "Total ad spend in USD."},
            {"name": "is_holiday_week", "type": "categorical", "categories": ["Yes", "No"], "description": "Whether the week contains a national holiday."},
            {"name": "average_temperature_f", "type": "numeric", "range": [-10.0, 110.0], "description": "Average regional temperature in Fahrenheit."}
        ],
        "outputs": [
            {"name": "predicted_store_sales_usd", "type": "numeric", "range": [5000.0, 150000.0], "description": "Forecasted revenue."}
        ],
        "clarifications_needed": []
    },
    "default": {
        "problem_domain": "NLP Sentiment Analysis",
        "technical_summary": "Classifies incoming text reviews to determine operational urgency levels.",
        "recommended_architecture": "DistilBERT Sequence Classifier",
        "confidence_score": 0.8,
        "fallback_assumptions_made": "Assumed user wants a classification engine.",
        "inputs": [
            {"name": "customer_review", "type": "text", "description": "Raw text of the review."}
        ],
        "outputs": [
            {"name": "sentiment_label", "type": "categorical", "categories": ["Positive", "Neutral", "Negative"], "description": "Underlying emotional charge."},
            {"name": "urgency_score", "type": "numeric", "range": [0.0, 1.0], "description": "Required response speed."}
        ],
        "clarifications_needed": ["Would you benefit from named entity extraction as well?"]
    }
}

# ==========================================
# 2. PROMPT TEMPLATES & GEMINI API HANDLER
# ==========================================
DEFAULT_SYSTEM_PROMPT = """You are an expert AI Solutions Architect. Your job is to parse unstructured, chaotic, or vague AI problem statements and translate them into a rigorous, production-ready JSON data contract. This contract will directly drive automated mock data generation and UI scaffolding.

### OUTPUT FORMAT CONSTRAINT
You must output exactly one JSON object. Do not include introductory text, conversational pleasantries, or concluding notes.

### CRITICAL DATA DICTIONARY CONSTRAINTS
To prevent breaking downstream scripts, values for the "type" fields must strictly be one of these exact string literals:
- "text" (for unstructured text, reviews, descriptions)
- "numeric" (for continuous integers or floats, like prices, age, coordinates)
- "categorical" (for discrete classes, labels, choices, or binary classifications)
- "image" (for visual files, bounding box arrays, pixels)

### TARGET JSON SCHEMA
{
  "problem_domain": "string",
  "technical_summary": "string",
  "recommended_architecture": "string",
  "confidence_score": float,
  "fallback_assumptions_made": "string",
  "inputs": [
    {
      "name": "string (snake_case column name)",
      "type": "string (exactly 'text', 'numeric', 'categorical', or 'image')",
      "description": "string",
      "categories": ["string"], // REQUIRED ONLY IF type is 'categorical'
      "range": [float, float] // REQUIRED ONLY IF type is 'numeric'
    }
  ],
  "outputs": [
    {
      "name": "string (snake_case column name)",
      "type": "string (exactly 'text', 'numeric', 'categorical', or 'image')",
      "description": "string",
      "categories": ["string"], // REQUIRED ONLY IF type is 'categorical'
      "range": [float, float] // REQUIRED ONLY IF type is 'numeric'
    }
  ],
  "clarifications_needed": ["string"]
}"""

@spaces.GPU
def call_llm_for_schema(api_key, problem_statement, system_prompt):
    """
    Tries to query the Google Gemini API using the modern google-genai SDK.
    If no key is supplied, defaults cleanly to sandbox simulation data.
    """
    if not api_key or len(api_key.strip()) < 10:
        p_lower = problem_statement.lower()
        if "sale" in p_lower or "price" in p_lower or "forecast" in p_lower or "demand" in p_lower:
            schema = MOCK_SCHEMAS["retail"]
        else:
            schema = MOCK_SCHEMAS["default"]
        return json.dumps(schema, indent=2), "⚠️ SYSTEM: Running in LOCAL SANDBOX mode (No API Key). Custom data contract simulated."

    try:
        from google import genai
        from google.genai import types
        
        # Initialize Google's GenAI Client
        client = genai.Client(api_key=api_key.strip())
        
        # Build strict JSON Generation config
        config = types.GenerateContentConfig(
            system_instruction=system_prompt,
            response_mime_type="application/json",
            temperature=0.2
        )
        
        # Invoke Gemini 2.5 Flash for high-speed, cost-effective compilation
        response = client.models.generate_content(
            model='gemini-3.5-flash',
            contents=problem_statement,
            config=config
        )
        
        # Clean response string to bypass raw triple backtick blocks if returned
        cleaned_text = response.text.strip()
        if cleaned_text.startswith("```json"):
            cleaned_text = cleaned_text.split("```json", 1)[1].rsplit("```", 1)[0].strip()
        elif cleaned_text.startswith("```"):
            cleaned_text = cleaned_text.split("```", 1)[1].rsplit("```", 1)[0].strip()
            
        return cleaned_text, "✅ Core contract successfully compiled by Gemini 2.5."
    except Exception as e:
        return json.dumps(MOCK_SCHEMAS["default"], indent=2), f"Error querying Gemini API: {str(e)}. Falling back to default mock schema."

@spaces.GPU
def generate_default_python_script(schema_str):
    """
    Generates editable raw Python code matching the JSON contract's inputs/outputs.
    """
    try:
        schema = json.loads(schema_str)
    except Exception:
        return "# Error: Invalid JSON schema generated in Step 1. Please correct it."
    
    script_lines = [
        "import pandas as pd",
        "import random",
        "",
        "def generate_dataset(num_rows=50):",
        "    data = []",
        "    for i in range(num_rows):",
        "        row = {}"
    ]
    
    # Map inputs
    for inp in schema.get("inputs", []):
        name = inp["name"]
        t = inp["type"]
        if t == "categorical":
            cats = inp.get("categories", ["Category A", "Category B"])
            script_lines.append(f"        row['{name}'] = random.choice({cats})")
        elif t == "numeric":
            r = inp.get("range", [0.0, 100.0])
            script_lines.append(f"        row['{name}'] = round(random.uniform({r[0]}, {r[1]}), 2)")
        elif t == "text":
            # FIXED: Double braces {{i+1}} tells Python to treat it as raw text in the output string
            script_lines.append(f"        row['{name}'] = f'Sample text data row {{i+1}}'")
        elif t == "image":
            script_lines.append(f"        row['{name}'] = f'mock_image_path_{{i+1}}.png'")
            
    # Map outputs
    for out in schema.get("outputs", []):
        name = out["name"]
        t = out["type"]
        if t == "categorical":
            cats = out.get("categories", ["Pass", "Fail"])
            script_lines.append(f"        row['{name}'] = random.choice({cats})")
        elif t == "numeric":
            r = out.get("range", [0.0, 1.0])
            script_lines.append(f"        row['{name}'] = round(random.uniform({r[0]}, {r[1]}), 4)")
        elif t == "text":
            # FIXED: Double braces {{i+1}} here too
            script_lines.append(f"        row['{name}'] = f'Target output summary text {{i+1}}'")
        elif t == "image":
            script_lines.append(f"        row['{name}'] = f'mock_processed_image_path_{{i+1}}.png'")

    script_lines.extend([
        "        data.append(row)",
        "    return pd.DataFrame(data)"
    ])
    
    return "\n".join(script_lines)

@spaces.GPU
def execute_custom_script(script_code, num_rows):
    """
    Compiles and executes the user-edited data generation script within a local dictionary.
    Explicitly injects 'random' and 'pandas' to prevent missing module errors during exec().
    """
    try:
        # Pre-populate the execution environment with the required modules
        namespace = {
            "pd": pd,
            "random": random
        }
        
        # Execute the code block inside our prepared environment
        exec(script_code, namespace, namespace)
        
        if "generate_dataset" not in namespace:
            return None, None, "Error: The script must define a function named 'generate_dataset(num_rows)'"
        
        df = namespace["generate_dataset"](int(num_rows))
        csv_filename = "generated_dataset.csv"
        df.to_csv(csv_filename, index=False)
        return df, csv_filename, "✅ Dataset generation executed successfully!"
    except Exception as e:
        return None, None, f"Execution Error: {str(e)}"


# ==========================================
# 3. GRADIO APP INTERFACE LAYOUT
# ==========================================
with gr.Blocks(theme=gr.themes.Soft(), title="Gemini AI Solutions Prototyper") as demo:
    schema_state = gr.State({})
    
    gr.Markdown("# 🚀 Meta-AI Prototyping Sandbox (Powered by Gemini)")
    gr.Markdown("Create a complete AI solution pipeline. Modify, tweak, and approve the structures at every step.")

    with gr.Tabs() as tabs:
        
        # ----------------------------------
        # TAB 1: ARCHITECTURE DESIGN
        # ----------------------------------
        with gr.Tab("Step 1: Architecture Designer", id=0):
            gr.Markdown("### Parse Problem Statement into a JSON Data Contract")
            
            with gr.Row():
                with gr.Column(scale=1):
                    api_key_input = gr.Textbox(
                        label="Google Gemini API Key (Optional)", 
                        placeholder="AIzaSy...", 
                        type="password",
                        info="Leave empty to use local offline simulation mode"
                    )
                    user_problem = gr.Textbox(
                        label="Your AI Problem Statement", 
                        value="Predict housing prices using location metrics, sqft size, and if it is near a transit stop.",
                        lines=4
                    )
                    edit_sys_prompt_btn = gr.Button("⚙️ Show / Edit Architectural System Prompt", size="sm")
                    sys_prompt_box = gr.Textbox(
                        label="Architect System Prompt",
                        value=DEFAULT_SYSTEM_PROMPT,
                        lines=12,
                        visible=False
                    )
                    
                    def toggle_sys_prompt(visible):
                        return gr.update(visible=not visible)
                    edit_sys_prompt_btn.click(toggle_sys_prompt, inputs=[sys_prompt_box], outputs=[sys_prompt_box])

                    generate_btn = gr.Button("Compile Solution Architecture with Gemini", variant="primary")
                    
                with gr.Column(scale=1):
                    engine_status = gr.Markdown("**Status:** Awaiting compilation.")
                    schema_output = gr.Code(
                        label="Generated JSON Data Contract (Editable)", 
                        language="json", 
                        interactive=True,
                        lines=20
                    )
                    
            next_to_step2_btn = gr.Button("Approve Contract & Advance to Data Generation ➡️", variant="secondary")

        # ----------------------------------
        # TAB 2: DATA CREATION
        # ----------------------------------
        with gr.Tab("Step 2: Dummy Dataset Creator", id=1):
            gr.Markdown("### Review and Edit the Custom Generation Script")
            
            with gr.Row():
                with gr.Column(scale=1):
                    row_slider = gr.Slider(minimum=5, maximum=1000, value=50, step=5, label="Number of Rows to Mock")
                    script_editor = gr.Code(
                        label="Data Generation Python Script (Editable)", 
                        language="python", 
                        interactive=True,
                        lines=15
                    )
                    run_script_btn = gr.Button("Execute Script & Build CSV", variant="primary")
                    
                with gr.Column(scale=1):
                    script_status = gr.Markdown("**Status:** Script compiled. Awaiting execution.")
                    data_preview = gr.Dataframe(label="Generated Data Preview", interactive=False)
                    download_link = gr.File(label="Download Generated CSV")

            next_to_step3_btn = gr.Button("Approve Data & Advance to Live Interface ➡️", variant="secondary")

        # ----------------------------------
        # TAB 3: LIVE PROTOTYPE UI
        # ----------------------------------
        with gr.Tab("Step 3 & 4: Live Demo", id=2):
            gr.Markdown("### Interactive Prototype Interface")
            gr.Markdown("The widgets below are generated dynamically using the approved JSON contract in Step 1. Test your model parameters below:")
            
            @gr.render(inputs=schema_state)
            def render_prototype_ui(schema):
                if not schema or "inputs" not in schema:
                    gr.Markdown("### ⚠️ Waiting for Architecture Selection\nPlease generate or copy a valid JSON schema into **Step 1** to activate this view.")
                    return
                
                gr.Markdown(f"### Demo Module: **{schema.get('problem_domain', 'AI Model')}**")
                gr.Markdown(f"**Description:** {schema.get('technical_summary', 'Simulation Module.')}")
                gr.Markdown(f"*Recommended Architecture: `{schema.get('recommended_architecture', 'Vanilla Machine Learning')}`*")
                
                inputs = []
                with gr.Row():
                    # Generate Input Widgets
                    with gr.Column(scale=1, variant="panel"):
                        gr.Markdown("#### Dynamic Parameters (Inputs)")
                        for inp in schema["inputs"]:
                            name = inp["name"]
                            t = inp["type"]
                            desc = inp.get("description", "")
                            
                            if t == "text":
                                inputs.append((name, gr.Textbox(label=name, info=desc)))
                            elif t == "categorical":
                                inputs.append((name, gr.Dropdown(choices=inp.get("categories", ["Option A"]), label=name, info=desc)))
                            elif t == "numeric":
                                r = inp.get("range", [0, 100])
                                inputs.append((name, gr.Slider(minimum=r[0], maximum=r[1], value=(r[0]+r[1])/2, label=name, info=desc)))
                            elif t == "image":
                                inputs.append((name, gr.Image(label=name, type="pil", info=desc)))

                    # Generate Output Display Widgets
                    with gr.Column(scale=1, variant="panel"):
                        gr.Markdown("#### Simulated AI Outputs")
                        outputs = []
                        for out in schema["outputs"]:
                            name = out["name"]
                            t = out["type"]
                            desc = out.get("description", "")
                            
                            if t == "text":
                                outputs.append((name, gr.Textbox(label=name, info=desc, interactive=False)))
                            elif t == "categorical":
                                outputs.append((name, gr.Textbox(label=name, info=desc, interactive=False)))
                            elif t == "numeric":
                                outputs.append((name, gr.Number(label=name, info=desc, interactive=False)))
                            elif t == "image":
                                outputs.append((name, gr.Image(label=name, type="pil", info=desc, interactive=False)))
                
                infer_btn = gr.Button("⚡ Execute Mock Model Inference", variant="primary")
                
                # Dynamic Execution Handler
                def run_inference(*args):
                    input_payload = {inputs[i][0]: args[i] for i in range(len(args))}
                    out_results = []
                    
                    for out in schema["outputs"]:
                        t = out["type"]
                        if t == "categorical":
                            out_results.append(random.choice(out.get("categories", ["N/A"])))
                        elif t == "numeric":
                            r = out.get("range", [0, 100])
                            out_results.append(round(random.uniform(r[0], r[1]), 2))
                        elif t == "text":
                            out_results.append(f"Model inferred successfully based on values: {list(input_payload.values())}")
                        elif t == "image":
                            # Draw a dynamic processing ring in memory
                            img = Image.new("RGB", (300, 300), color=(17, 24, 39)) # Deep charcoal slate
                            draw = ImageDraw.Draw(img)
                            draw.ellipse([100, 100, 200, 200], fill=(16, 185, 129)) # Emerald Green ring
                            out_results.append(img)
                            
                    return out_results

                infer_btn.click(
                    fn=run_inference,
                    inputs=[widget for _, widget in inputs],
                    outputs=[widget for _, widget in outputs]
                )

    # ==========================================
    # INTER-TAB COORDINATION CONTROL FLOW
    # ==========================================
    def step1_action(api_key, problem, prompt):
        raw_json, status_msg = call_llm_for_schema(api_key, problem, prompt)
        try:
            parsed_json = json.loads(raw_json)
        except Exception:
            parsed_json = {}
        script_code = generate_default_python_script(raw_json)
        return raw_json, status_msg, script_code, parsed_json

    generate_btn.click(
        fn=step1_action,
        inputs=[api_key_input, user_problem, sys_prompt_box],
        outputs=[schema_output, engine_status, script_editor, schema_state]
    )

    def advance_to_step2(raw_json):
        try:
            parsed_json = json.loads(raw_json)
        except Exception:
            return gr.update(selected=0), {}, ""
        script_code = generate_default_python_script(raw_json)
        return gr.update(selected=1), parsed_json, script_code

    next_to_step2_btn.click(
        fn=advance_to_step2,
        inputs=[schema_output],
        outputs=[tabs, schema_state, script_editor]
    )

    run_script_btn.click(
        fn=execute_custom_script,
        inputs=[script_editor, row_slider],
        outputs=[data_preview, download_link, script_status]
    )

    def advance_to_step3(raw_json):
        try:
            parsed_json = json.loads(raw_json)
        except Exception:
            parsed_json = {}
        return gr.update(selected=2), parsed_json

    next_to_step3_btn.click(
        fn=advance_to_step3,
        inputs=[schema_output],
        outputs=[tabs, schema_state]
    )

if __name__ == "__main__":
    demo.launch(pwa=True)