app.py

import os
import uuid
import shutil
import traceback
import json

import gradio as gr
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware

from visualization import process_wireframe
from CodingModule import HTMLGenerator

# -----------------------------------------------------------------------------
# FASTAPI BACKEND
# -----------------------------------------------------------------------------
api = FastAPI()

api.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

TEMP_DIR = "./temp"
OUTPUT_DIR = "./output"

os.makedirs(TEMP_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)


@api.get("/")
def health_check():
    return {"status": "ok"}


@api.post("/process-wireframe")
async def process_wireframe_api(image: UploadFile = File(...)):
    """
    Firebase / Programmatic endpoint
    Returns JSON + HTML content directly.
    """
    file_id = str(uuid.uuid4())
    temp_path = os.path.join(TEMP_DIR, f"{file_id}_{image.filename}")

    try:
        with open(temp_path, "wb") as f:
            shutil.copyfileobj(image.file, f)

        results = process_wireframe(
            image_path=temp_path,
            save_json=True,
            save_html=False,
            show_visualization=False
        )

        if not results or not results.get('normalized_elements'):
            return JSONResponse(
                status_code=400,
                content={"error": "No elements detected"}
            )

        json_path = results.get("json_path")

        # Load JSON content
        json_content = None
        if json_path and os.path.exists(json_path):
            with open(json_path, "r") as f:
                json_content = json.load(f)

        # Generate styled HTML
        html_path = None
        try:
            html_generator = HTMLGenerator(json_path)
            html_filename = f"{file_id}_styled.html"
            html_path = os.path.join(OUTPUT_DIR, html_filename)
            html_generator.generate_html(html_path)
        except:
            html_path = results.get("html_path")

        html_content = None
        if html_path and os.path.exists(html_path):
            with open(html_path, "r", encoding="utf-8") as f:
                html_content = f.read()

        elements = (
            results.get("normalized_elements")
            or json_content.get("normalized_elements")
            or json_content.get("elements")
            or []
        )

        return {
            "success": True,
            "json_data": json_content,
            "html_content": html_content,
            "total_elements": len(elements),
            "normalized_elements": elements
        }

    except Exception as e:
        traceback.print_exc()
        return JSONResponse(status_code=500, content={"error": str(e)})

    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)


@api.get("/view-html/{filename}")
async def serve_output_file(filename: str):
    """
    Serve stored HTML/JSON files.
    """
    file_path = os.path.join(OUTPUT_DIR, filename)
    if not os.path.exists(file_path):
        return JSONResponse(status_code=404, content={"error": "File not found"})

    if filename.endswith(".html"):
        return FileResponse(file_path, media_type="text/html")
    elif filename.endswith(".json"):
        return FileResponse(file_path, media_type="application/json")
    return FileResponse(file_path)


# -----------------------------------------------------------------------------
# GRADIO UI LOGIC
# -----------------------------------------------------------------------------
def gradio_process(image):
    if image is None:
        return "Please upload an image", None, None

    file_id = str(uuid.uuid4())
    temp_path = os.path.join(TEMP_DIR, f"{file_id}.png")

    try:
        image.save(temp_path)
    except Exception as e:
        return f"Error saving: {str(e)}", None, None

    try:
        results = process_wireframe(
            image_path=temp_path,
            save_json=True,
            save_html=False,
            show_visualization=False
        )

        if not results.get("normalized_elements"):
            return "No elements detected", None, None

        json_path = results["json_path"]

        # Generate styled HTML
        try:
            html_generator = HTMLGenerator(json_path)
            html_filename = f"{file_id}_styled.html"
            html_path = os.path.join(OUTPUT_DIR, html_filename)
            html_generator.generate_html(html_path)
        except:
            traceback.print_exc()
            html_path = results.get("html_path")

        num = len(results["normalized_elements"])
        status = f"Detected {num} elements\n"
        status += f"JSON: {os.path.basename(json_path)}\n"
        status += f"HTML: {os.path.basename(html_path)}"

        return status, json_path, html_path

    except Exception as e:
        traceback.print_exc()
        return f"Error: {str(e)}", None, None

    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)


# -----------------------------------------------------------------------------
# GRADIO UI (updated for Gradio 6 compliance)
# -----------------------------------------------------------------------------
with gr.Blocks(title="Wireframe Layout Normalizer") as demo:

    gr.Markdown("# Wireframe Layout Normalizer")
    gr.Markdown("Upload a wireframe image to extract and normalize UI layout elements.")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Wireframe Image")
            process_btn = gr.Button("Process Wireframe")

        with gr.Column():
            status_output = gr.Textbox(label="Status")
            json_output = gr.File(label="JSON Output")
            html_output = gr.File(label="HTML Preview")

    process_btn.click(
        fn=gradio_process,
        inputs=image_input,
        outputs=[status_output, json_output, html_output],
    )


# -----------------------------------------------------------------------------
# FINAL APP MOUNT (fixes reload + /new spam)
# -----------------------------------------------------------------------------
app = gr.mount_gradio_app(
    api,
    demo,
    path="/gradio",
    theme=gr.themes.Soft()
)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)