import streamlit as st
import json
import time
from PIL import Image
import os
import sys
import shutil
import gdown
from io import BytesIO

# ==================================
# SETUP
# ==================================

print("🚀 Streamlit App Starting...")

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# Setup Paths
UPLOAD_DIR = "/tmp/uploads/"
MODEL_DIR = os.path.join(BASE_DIR, "rcnn_model", "scripts")
JSON_DIR = "/tmp/results/"
OUTPUT_DIR = "/tmp/output/"
SAMPLE_DIR = os.path.join(BASE_DIR, "rcnn_model", "sample")
logo_path = os.path.join(BASE_DIR, "public", "logo.png")
model_path = os.path.join(OUTPUT_DIR, "model_final.pth")

# Google Drive file download link
GOOGLE_DRIVE_FILE_ID = "1yr64AOgaYZPTcQzG6cxG6lWBENHR9qjW"
GDRIVE_URL = f"https://drive.google.com/uc?id={GOOGLE_DRIVE_FILE_ID}"

# Create necessary folders
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(JSON_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)

# ==================================
# DOWNLOAD MODEL IF MISSING
# ==================================

if not os.path.exists(model_path):
    print("🚀 Model file not found! Downloading from Google Drive...")
    try:
        gdown.download(GDRIVE_URL, model_path, quiet=False)
        print("✅ Model downloaded successfully.")
    except Exception as e:
        print(f"❌ Failed to download model: {e}")

# ==================================
# IMPORT MODEL RUNNER
# ==================================

sys.path.append(MODEL_DIR)
from rcnn_model.scripts.rcnn_run import main, write_config

# ==================================
# PAGE CONFIG
# ==================================

st.set_page_config(
    page_title="2D Floorplan Vectorizer",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# ==================================
# HEADER
# ==================================

st.image(logo_path, width=250)
st.markdown("<div class='header-title'>2D Floorplan Vectorizer</div>", unsafe_allow_html=True)

# ==================================
# FILE UPLOAD SECTION
# ==================================

st.subheader("Upload your Floorplan Image")
uploaded_file = st.file_uploader("Choose an image", type=["png", "jpg", "jpeg"])

# Initialize session state
if "processing_complete" not in st.session_state:
    st.session_state.processing_complete = False
if "json_output" not in st.session_state:
    st.session_state.json_output = None

# ==================================
# IMAGE + JSON Layout
# ==================================

col1, col2 = st.columns([1, 2])

# ==================================
# MAIN LOGIC
# ==================================

if uploaded_file is not None:
    print("📤 File Uploaded:", uploaded_file.name)

    image_bytes = uploaded_file.read()
    img = Image.open(BytesIO(image_bytes)).convert("RGB")

    uploaded_path = os.path.join(UPLOAD_DIR, uploaded_file.name)
    with open(uploaded_path, "wb") as f:
        f.write(uploaded_file.getbuffer())
    print("✅ Uploaded file saved at:", uploaded_path)

    with col1:
        st.markdown("<div class='upload-container'>", unsafe_allow_html=True)
        st.image(Image.open(uploaded_path), caption="Uploaded Image", use_container_width=True)
        st.markdown("</div>", unsafe_allow_html=True)

    with col2:
        if not st.session_state.processing_complete:
            status_placeholder = st.empty()
            status_placeholder.info("⏳ Model is processing the uploaded image...")
            progress_bar = st.progress(0)
            status_text = st.empty()

            # === 🔥 Model Run Here ===
            input_image = uploaded_path
            output_json_name = uploaded_file.name.replace(".png", "_result.json").replace(".jpg", "_result.json").replace(".jpeg", "_result.json")
            output_image_name = uploaded_file.name.replace(".png", "_result.png").replace(".jpg", "_result.png").replace(".jpeg", "_result.png")

            output_json_path = os.path.join(JSON_DIR, output_json_name)
            output_image_path = os.path.join(JSON_DIR, output_image_name)

            cfg = write_config()
            print("⚙️ Model config created. Running model...")

            # Simulate progress
            for i in range(1, 30):
                time.sleep(0.01)
                progress_bar.progress(i)
                status_text.text(f"Preprocessing: {i}%")

            # Run model
            main(cfg, input_image, output_json_path, output_image_path)
            print("✅ Model run complete.")

            while not os.path.exists(output_json_path):
                print("Waiting for JSON output...")
                time.sleep(0.5)

            for i in range(30, 100):
                time.sleep(0.01)
                progress_bar.progress(i)
                status_text.text(f"Postprocessing: {i}%")

            progress_bar.empty()
            status_text.text("✅ Processing Complete!")
            status_placeholder.success("✅ Model finished and JSON is ready!")

            # Read generated JSON
            if os.path.exists(output_json_path):
                with open(output_json_path, "r") as jf:
                    st.session_state.json_output = json.load(jf)
                    print("📄 JSON Output Loaded Successfully.")
            else:
                st.session_state.json_output = {"error": "JSON output not generated."}
                print("❌ JSON output missing.")

            st.session_state.processing_complete = True

        # ==================================
        # DISPLAY OUTPUTS
        # ==================================

        out_col1, out_col2 = st.columns(2)

        with out_col1:
            if os.path.exists(output_image_path):
                with open(output_image_path, "rb") as img_file:
                    image = Image.open(img_file)
                    st.image(image, caption="🖼 Output Vectorized Image", use_container_width=True)

                    img_file.seek(0)
                    st.download_button(
                        label="Download Output Image",
                        data=img_file,
                        file_name="floorplan_output.png",
                        mime="image/png"
                    )

            if os.path.exists(output_json_path):
                json_str = json.dumps(st.session_state.json_output, indent=4)
                st.download_button(
                    label="Download JSON",
                    data=json_str,
                    file_name="floorplan_output.json",
                    mime="application/json"
                )

        with out_col2:
            st.markdown("<div class='json-container'>", unsafe_allow_html=True)
            st.json(st.session_state.json_output)
            st.markdown("</div>", unsafe_allow_html=True)

else:
    st.warning("⚠️ No image uploaded yet.")
    st.session_state.processing_complete = False