import os
import nbformat
import gradio as gr
from fastapi import FastAPI
from nbconvert.preprocessors import ExecutePreprocessor
from fastapi.responses import JSONResponse

# =========================
# Setup
# =========================

TEMP_FOLDER = "temp"
os.makedirs(TEMP_FOLDER, exist_ok=True)

# FastAPI App
app = FastAPI()

# =========================
# Notebook Execution
# =========================


def execute_notebook(notebook_path):
    try:
        with open(notebook_path, "r", encoding="utf-8") as f:
            nb = nbformat.read(f, as_version=4)

        ep = ExecutePreprocessor(
            timeout=1800,
            kernel_name="python3"
        )

        ep.preprocess(nb, {'metadata': {'path': TEMP_FOLDER}})

        executed_file = os.path.join(
            TEMP_FOLDER,
            f"executed_{os.path.basename(notebook_path)}"
        )

        with open(executed_file, "w", encoding="utf-8") as f:
            nbformat.write(nb, f)

        return {
            "success": True,
            "output": executed_file
        }

    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }


# =========================
# Gradio Upload Logic
# =========================


def process_notebooks(notebook1, notebook2):
    results = []

    if notebook1:
        result1 = execute_notebook(notebook1.name)
        results.append(f"Notebook 1: {result1}")

    if notebook2:
        result2 = execute_notebook(notebook2.name)
        results.append(f"Notebook 2: {result2}")

    return "\n\n".join(results)


# =========================
# API Endpoints for n8n
# =========================

@app.get("/health")
def health_check():
    return {
        "status": "running",
        "service": "movie recommendation pipeline"
    }


@app.post("/run-pipeline")
def run_pipeline():
    notebook1_path = "notebook1.ipynb"
    notebook2_path = "notebook2.ipynb"

    results = {}

    if os.path.exists(notebook1_path):
        results["notebook1"] = execute_notebook(notebook1_path)

    if os.path.exists(notebook2_path):
        results["notebook2"] = execute_notebook(notebook2_path)

    return JSONResponse(content={
        "status": "completed",
        "results": results
    })


# =========================
# Gradio UI
# =========================

with gr.Blocks() as demo:
    gr.Markdown("# Movie Recommendation Notebook Pipeline")

    notebook1 = gr.File(
        label="Upload Notebook 1 (.ipynb)",
        file_types=[".ipynb"]
    )

    notebook2 = gr.File(
        label="Upload