import os, sys, zipfile, shutil, subprocess
import gradio as gr

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ZIP_PATH = os.path.join(BASE_DIR, "engineering_drawing_code.zip")
REPO_DIR = os.path.join(BASE_DIR, "engineering_drawing_repo")

def extract_two_outputs(uploaded_image_path):
    # 1) Unzip the extractor repo if needed
    if not os.path.isdir(REPO_DIR):
        with zipfile.ZipFile(ZIP_PATH, "r") as z:
            z.extractall(REPO_DIR)
    # 2) Clean previous outputs
    out_excel = os.path.join(BASE_DIR, "drawingInfo.xlsx")
    out_png   = os.path.join(BASE_DIR, "extracted", "drawing01.png")
    if os.path.exists(out_excel):
        os.remove(out_excel)
    if os.path.isdir(os.path.join(BASE_DIR, "extracted")):
        shutil.rmtree(os.path.join(BASE_DIR, "extracted"))
    # 3) Prepare images folder
    images_folder = os.path.join(REPO_DIR, "images")
    if os.path.isdir(images_folder):
        shutil.rmtree(images_folder)
    os.makedirs(images_folder, exist_ok=True)
    shutil.copyfile(uploaded_image_path,
                    os.path.join(images_folder, "01.png"))
    # 4) Run the OCR/extraction script
    subprocess.run(
        [sys.executable, "mainExtractionOCR.py"],
        cwd=REPO_DIR,
        check=True,
    )
    # 5) Return the ballooned drawing + Excel
    return out_png, out_excel

# Build Gradio UI
demo = gr.Blocks()
with demo:
    gr.Markdown("## Engineering-Drawing Information Extractor")
    with gr.Row():
        inp = gr.Image(type="filepath", label="Upload drawing")
    with gr.Row():
        balloon = gr.Image(type="filepath", label="Ballooned Drawing")
        excel   = gr.File(label="Download Extracted Data (Excel)")
    btn = gr.Button("Extract")
    btn.click(fn=extract_two_outputs, inputs=inp, outputs=[balloon, excel])

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