app.py

import streamlit as st
import uuid
import os
import json
import shutil
import zipfile
from typing import Dict, List

from spec_validator import validate_and_clean_spec
from pipeline import run_pipeline



# =====================================================
# CONFIG
# =====================================================
BASE_RUN_DIR = "runs"
SCHEMA_PATH = "spec_schema.json"

st.set_page_config(
    page_title="AI REST API Generator",
    layout="wide",
)

# =====================================================
# HELPERS
# =====================================================

def create_run_dir():
    run_id = uuid.uuid4().hex[:8]
    run_dir = os.path.join(BASE_RUN_DIR, f"run_{run_id}")
    os.makedirs(run_dir, exist_ok=True)
    return run_id, run_dir


def zip_folder(folder_path: str) -> str:
    zip_path = f"{folder_path}.zip"
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
        for root, _, files in os.walk(folder_path):
            for file in files:
                full_path = os.path.join(root, file)
                arcname = os.path.relpath(full_path, folder_path)
                zipf.write(full_path, arcname)
    return zip_path


def load_schema():
    with open(SCHEMA_PATH, "r", encoding="utf-8") as f:
        return json.load(f)


# =====================================================
# SESSION STATE INIT
# =====================================================

if "spec" not in st.session_state:
    st.session_state.spec = {
        "project_name": "",
        "description": "",
        "database": {},
        "auth": {"type": "jwt"},
        "api_config": {"base_url": "/api/"},
        "apps": {"core": {"models": {}, "apis": {}}},
        "use_ai_models": True,
        "llm": {}
    }

if "models_ui" not in st.session_state:
    st.session_state.models_ui = []

if "logs" not in st.session_state:
    st.session_state.logs = []


# =====================================================
# UI HEADER
# =====================================================

st.title("🤖 AI REST API Generator")
st.markdown(
    """
    **Build Django REST APIs using AI or manual configuration.**  
    🚧 *This website is under development.*
    """
)

st.divider()

# =====================================================
# SECTION 1 — LLM CONFIGURATION
# =====================================================
st.header("🔑 LLM Configuration")

llm_provider = st.selectbox(
    "Select LLM Provider",
    ["Groq", "OpenAI"]
)

api_key = st.text_input(
    f"{llm_provider} API Key",
    type="password",
    help="Your API key is used only for this session and never stored."
)

st.session_state.spec["llm"] = {
    "provider": llm_provider.lower()
}




# =====================================================
# SECTION 2 — PROJECT BASICS
# =====================================================
st.header("📦 Project Information")

col1, col2 = st.columns(2)

with col1:
    st.session_state.spec["project_name"] = st.text_input(
        "Project Name",
        placeholder="my_backend_project"
    )

with col2:
    st.session_state.spec["description"] = st.text_area(
        "Project Description",
        placeholder="Short description of your backend"
    )

# =====================================================
# SECTION 3 — DATABASE CONFIGURATION
# =====================================================
st.header("🗄 Database Configuration")

db_engine = st.selectbox(
    "Database Engine",
    ["sqlite", "postgresql", "mysql"]
)

if db_engine == "sqlite":
    st.session_state.spec["database"] = {
        "engine": "sqlite",
        "name": "db.sqlite3"
    }
else:
    col1, col2 = st.columns(2)
    with col1:
        db_name = st.text_input("Database Name")
        db_user = st.text_input("Database User")
        db_password = st.text_input("Database Password", type="password")
    with col2:
        db_host = st.text_input("Database Host", value="localhost")
        db_port = st.number_input("Database Port", value=5432)

    st.session_state.spec["database"] = {
        "engine": db_engine,
        "name": db_name,
        "user": db_user,
        "password": db_password,
        "host": db_host,
        "port": db_port
    }

# =====================================================
# SECTION 4 — MODEL MODE
# =====================================================
st.header("🧱 Model Generation")

use_ai = st.radio(
    "How do you want to create models?",
    ["AI Generated Models", "Manual Model Builder"]
)

st.session_state.spec["use_ai_models"] = (use_ai == "AI Generated Models")

# =====================================================
# SECTION 5 — MANUAL MODEL BUILDER
# =====================================================
if not st.session_state.spec["use_ai_models"]:
    st.subheader("🛠 Manual Model Builder")

    if st.button("➕ Add Model"):
        st.session_state.models_ui.append({
            "name": "",
            "fields": []
        })

    model_names = []

    for mi, model in enumerate(st.session_state.models_ui):
        with st.expander(f"Model {mi + 1}", expanded=True):
            model["name"] = st.text_input(
                "Model Name",
                model["name"],
                key=f"model_{mi}"
            )

            if model["name"]:
                model_names.append(model["name"])

            if st.button("➕ Add Column", key=f"add_col_{mi}"):
                model["fields"].append({
                    "name": "",
                    "type": "CharField",
                    "primary_key": False,
                    "unique": False,
                    "null": False,
                    "relation": None,
                    "on_delete": "CASCADE"
                })

            for fi, field in enumerate(model["fields"]):
                cols = st.columns(7)

                field["name"] = cols[0].text_input(
                    "Column",
                    field["name"],
                    key=f"fname_{mi}_{fi}"
                )

                field["type"] = cols[1].selectbox(
                    "Type",
                    [
                        "CharField",
                        "IntegerField",
                        "UUIDField",
                        "BooleanField",
                        "DateField",
                        "OneToOne",
                        "OneToMany"
                    ],
                    key=f"ftype_{mi}_{fi}"
                )

                field["primary_key"] = cols[2].checkbox(
                    "PK",
                    key=f"fpk_{mi}_{fi}"
                )
                field["unique"] = cols[3].checkbox(
                    "Unique",
                    key=f"funq_{mi}_{fi}"
                )
                field["null"] = cols[4].checkbox(
                    "Null",
                    key=f"fnull_{mi}_{fi}"
                )

                if field["type"] in ["OneToOne", "OneToMany"]:
                    field["relation"] = cols[5].selectbox(
                        "Reference Model",
                        model_names,
                        key=f"frel_{mi}_{fi}"
                    )
                    field["on_delete"] = cols[6].selectbox(
                        "On Delete",
                        ["CASCADE", "SET_NULL", "PROTECT"],
                        key=f"fdel_{mi}_{fi}"
                    )

    # Convert UI → spec
    models_spec = {}
    apis_spec = {}

    for model in st.session_state.models_ui:
        if not model["name"]:
            continue

        fields_spec = {}

        for f in model["fields"]:
            if not f["name"]:
                continue

            if f["type"] == "OneToOne":
                fields_spec[f["name"]] = {
                    "type": "OneToOneField",
                    "to": f["relation"],
                    "on_delete": f["on_delete"],
                    "null": f["null"],
                    "unique": True
                }
            elif f["type"] == "OneToMany":
                fields_spec[f["name"]] = {
                    "type": "ForeignKey",
                    "to": f["relation"],
                    "on_delete": f["on_delete"],
                    "null": f["null"]
                }
            else:
                fields_spec[f["name"]] = {
                    "type": f["type"],
                    "primary_key": f["primary_key"],
                    "unique": f["unique"],
                    "null": f["null"]
                }

        models_spec[model["name"]] = {"fields": fields_spec}
        apis_spec[model["name"]] = ["list", "create", "retrieve", "update", "delete"]

    st.session_state.spec["apps"]["core"]["models"] = models_spec
    st.session_state.spec["apps"]["core"]["apis"] = apis_spec

# =====================================================
# SECTION 6 — GENERATION
# =====================================================
st.header("🚀 Generate Project")



# UI elements (define once, above button)
progress_bar = st.progress(0)

log_box = st.empty()

def log(msg: str):
    st.session_state.logs.append(msg)
    log_box.code(
        "\n".join(st.session_state.logs),
        language="text"
    )

def ui_log(msg):
    st.session_state.logs.append(msg)
    log_box.write("\n".join(st.session_state.logs))

def ui_progress(value):
    progress_bar.progress(value)


if st.button("Generate Backend"):
    # ----------------------------
    # 1️⃣ Basic validation
    # ----------------------------
    if not api_key:
        st.error("API key is required to generate the project.")
        st.stop()

    st.session_state.logs.clear()
    progress_bar.progress(0)

    # ----------------------------
    # 2️⃣ Load schema & validate UI spec
    # ----------------------------
    schema = load_schema()

    valid, cleaned_spec, errors, warnings = validate_and_clean_spec(
        st.session_state.spec,
        schema
    )

    if not valid:
        st.error("Specification validation failed.")
        st.json(errors)
        st.stop()

    # ----------------------------
    # 3️⃣ Create isolated run directory
    # ----------------------------
    run_id, run_dir = create_run_dir()
    ui_log(f"🆔 Run ID: {run_id}")
    ui_log("🚀 Starting generation pipeline...")

    # ----------------------------
    # 4️⃣ Run backend pipeline
    # ----------------------------
    try:
        run_pipeline(
            spec=cleaned_spec,
            run_dir=run_dir,
            llm_provider=llm_provider.lower(),
            api_key=api_key,
            log_callback=ui_log,
            progress_callback=ui_progress
        )

        ui_log("📦 Zipping project...")
        zip_path = zip_folder(run_dir)

        ui_log("✅ Project ready for download")

        with open(zip_path, "rb") as f:
            st.download_button(
                "⬇ Download Generated Project",
                f,
                file_name=f"{cleaned_spec['project_name']}.zip"
            )

    # ----------------------------
    # 5️⃣ Error handling
    # ----------------------------
    except Exception as e:
        ui_log("❌ Pipeline failed")
        ui_log(str(e))
        st.error(f"Generation failed at step: {e}")


    # ----------------------------
    # 6️⃣ Cleanup (safe)
    # ----------------------------
    finally:
        shutil.rmtree(run_dir, ignore_errors=True)


# =====================================================
# FOOTER
# =====================================================
st.divider()
st.markdown(
    """
    📧 **Contact:** harshadhole04@gmail.com  
    © AI REST API Generator — Under Development
    """
)