import json
import time
import zipfile
import tempfile
from io import BytesIO
from concurrent.futures import ThreadPoolExecutor

import streamlit as st
import roboflow
import pandas as pd
import matplotlib.pyplot as plt

from shapely.geometry import Polygon
from PIL import Image

import gspread
from google.oauth2.service_account import Credentials
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseUpload


# =========================
# Page config
# =========================
st.set_page_config(page_title="Scratch Assay Segmentation", layout="wide")

APP_VERSION = "2.4"


# =========================
# Roboflow init
# =========================
API_KEY = st.secrets["roboflow_api_key"]
rf = roboflow.Roboflow(api_key=API_KEY)
project = rf.workspace(st.secrets["roboflow_workspace"]).project(st.secrets["roboflow_project"])
model = project.version(st.secrets["roboflow_version"]).model
model.confidence = 80
model.overlap = 25
dpi_value = 300


# =========================
# Google Drive + Sheets (Service Account)  ✅ FIX
# =========================
SCOPES = [
    "https://www.googleapis.com/auth/drive",
    "https://www.googleapis.com/auth/spreadsheets",
]

sa_info = json.loads(st.secrets["gcp_service_account"])
credentials = Credentials.from_service_account_info(sa_info, scopes=SCOPES)

drive_service = build("drive", "v3", credentials=credentials)
sheets_client = gspread.authorize(credentials)
sheet = sheets_client.open_by_url(st.secrets["feedback_sheet_url"]).sheet1


# =========================
# Helpers
# =========================
def calculate_polygon_area(points):
    polygon = Polygon([(p["x"], p["y"]) for p in points])
    return polygon.area


def safe_predict(image_path):
    for _ in range(3):
        try:
            return model.predict(image_path)
        except Exception:
            time.sleep(1)
    return None


def resize_image(image):
    return image.resize((640, 640))


def upload_to_drive(image_bytes, filename, folder_id):
    media = MediaIoBaseUpload(image_bytes, mimetype="image/png")
    drive_service.files().create(
        body={"name": filename, "parents": [folder_id]},
        media_body=media,
        fields="id",
    ).execute()


def find_or_create_folder(folder_name, parent=None):
    query = f"name='{folder_name}' and mimeType='application/vnd.google-apps.folder' and trashed=false"
    if parent:
        query += f" and '{parent}' in parents"

    results = drive_service.files().list(
        q=query,
        spaces="drive",
        fields="files(id, name)"
    ).execute()

    folders = results.get("files", [])
    if folders:
        return folders[0]["id"]

    file_metadata = {"name": folder_name, "mimeType": "application/vnd.google-apps.folder"}
    if parent:
        file_metadata["parents"] = [parent]

    file = drive_service.files().create(body=file_metadata, fields="id").execute()
    return file.get("id")


def get_image_bytes(image):
    buf = BytesIO()
    image.save(buf, format="PNG")
    buf.seek(0)
    return buf


def process_image(uploaded_file, fov_um=None, pixel_size_um=None):
    """
    - Runs inference through Roboflow
    - Computes area in px² and, if calibrated, in µm²
    - Returns buffers for Original / Segmented overlay / Polygon
    """
    try:
        safe_name = uploaded_file.name.replace(" ", "_")
        image = Image.open(uploaded_file).convert("RGB")

        width_px, _ = image.size

        effective_pixel_size_um = None
        if pixel_size_um is not None and pixel_size_um > 0:
            effective_pixel_size_um = pixel_size_um
        elif fov_um is not None and fov_um > 0:
            effective_pixel_size_um = fov_um / float(width_px)

        with tempfile.NamedTemporaryFile(suffix=".png", delete=True) as temp_file:
            image.save(temp_file.name)
            prediction = safe_predict(temp_file.name)

            if not prediction:
                return {
                    "Imagem": safe_name,
                    "Área Segmentada (px²)": None,
                    "Área Segmentada (µm²)": None,
                    "SemSegmentacao": True,
                    "Exibir": image,
                    "Original": get_image_bytes(image),
                    "Segmentada": None,
                    "Poligono": None,
                }

            prediction_data = prediction.json()

        if not prediction_data.get("predictions"):
            return {
                "Imagem": safe_name,
                "Área Segmentada (px²)": None,
                "Área Segmentada (µm²)": None,
                "SemSegmentacao": True,
                "Exibir": image,
                "Original": get_image_bytes(image),
                "Segmentada": None,
                "Poligono": None,
            }

        points = prediction_data["predictions"][0]["points"]

        area_px2 = calculate_polygon_area(points)
        area_um2 = None
        if effective_pixel_size_um is not None:
            area_um2 = area_px2 * (effective_pixel_size_um ** 2)

        x = [p["x"] for p in points] + [points[0]["x"]]
        y = [p["y"] for p in points] + [points[0]["y"]]

        original_buffer = get_image_bytes(image)

        # Segmented overlay
        segmented_buffer = BytesIO()
        fig, ax = plt.subplots(figsize=(6, 6), dpi=dpi_value)
        ax.imshow(image)
        ax.plot(x, y, color="red", linewidth=2)
        ax.axis("off")
        plt.savefig(segmented_buffer, format="png", bbox_inches="tight", pad_inches=0)
        plt.close()

        # Polygon only
        polygon_buffer = BytesIO()
        fig2, ax2 = plt.subplots(figsize=(6, 6), dpi=dpi_value)
        ax2.plot(x, y, "r-", linewidth=2)
        ax2.scatter(x, y, color="red", s=5)
        ax2.set_title("Polygon contour")
        ax2.grid(True)
        plt.savefig(polygon_buffer, format="png", bbox_inches="tight")
        plt.close()

        return {
            "Imagem": safe_name,
            "Área Segmentada (px²)": area_px2,
            "Área Segmentada (µm²)": area_um2,
            "Original": original_buffer,
            "Segmentada": segmented_buffer,
            "Poligono": polygon_buffer,
            "Exibir": image,
            "SemSegmentacao": False,
        }

    except Exception:
        return None


def save_feedback(result, avaliacao, observacao):
    """
    - Appends feedback to Google Sheet
    - Uploads images to Drive for curation (Acceptable/Bad/No segmentation)
    """
    image_name = result["Imagem"]

    # 1) Sheet (same structure as you had)
    sheet.append_row([image_name, avaliacao, observacao])

    # 2) Drive curation (same logic as before)
    if avaliacao in ["Acceptable", "Bad", "No segmentation"]:
        sufixo = (
            "aceitavel" if avaliacao == "Acceptable"
            else "ruim" if avaliacao == "Bad"
            else "sem_segmentacao"
        )

        parent_folder = find_or_create_folder("Feedback Segmentacoes")
        subfolder = find_or_create_folder(image_name.replace(".png", ""), parent_folder)

        # Original
        resized_original = resize_image(result["Exibir"])
        buf = BytesIO()
        resized_original.save(buf, format="PNG")
        buf.seek(0)
        upload_to_drive(buf, f"original_{sufixo}.png", subfolder)

        # Segmented + Polygon only if segmentation exists
        if avaliacao != "No segmentation" and result.get("Segmentada") and result.get("Poligono"):
            resized_segmented = resize_image(Image.open(BytesIO(result["Segmentada"].getvalue())))
            resized_polygon = resize_image(Image.open(BytesIO(result["Poligono"].getvalue())))

            for img_obj, nome in zip([resized_segmented, resized_polygon], ["segmentada", "poligono"]):
                buf = BytesIO()
                img_obj.save(buf, format="PNG")
                buf.seek(0)
                upload_to_drive(buf, f"{nome}_{sufixo}.png", subfolder)


def render_metrics(result):
    area_px2 = result["Área Segmentada (px²)"]
    area_um2 = result["Área Segmentada (µm²)"]

    st.markdown("**Segmented area**")
    if area_px2 is not None:
        st.markdown(f"- {area_px2:.2f} px²")
    if area_um2 is not None:
        st.markdown(f"- {area_um2:.2f} µm²")


def render_feedback_block(result, prefix_key=""):
    st.markdown("#### Segmentation quality feedback")
    st.caption("User evaluation used for future model refinement.")

    avaliacao = st.radio(
        "Segmentation quality assessment:",
        ["Great", "Acceptable", "Bad", "No segmentation"],
        horizontal=True,
        key=f"{prefix_key}radio_{result['Imagem']}",
    )
    observacao = st.text_area(
        "Observations (optional):",
        key=f"{prefix_key}obs_{result['Imagem']}",
    )
    if st.button("Save feedback", key=f"{prefix_key}btn_{result['Imagem']}"):
        save_feedback(result, avaliacao, observacao)
        st.success("Feedback saved successfully.")


# =========================
# UI
# =========================
st.title("Scratch Assay Segmentation Tool")
st.caption(f"Version {APP_VERSION} · Deep learning–based wound closure segmentation")
st.markdown("---")

# Input block
st.markdown("### Input")
upload_option = st.radio("Choose upload type:", ["Single image", "Image folder"], horizontal=True)

# Advanced settings (collapsed)
with st.expander("⚙️ Advanced Settings", expanded=False):
    model.confidence = st.slider("Model confidence (%)", 20, 100, 80)
    st.markdown(
        "### Physical calibration (optional)\n"
        "Provide the physical scale for conversion from pixel area to physical units (µm²). "
        "If left empty, results will be reported only in pixels²."
    )
    c1, c2 = st.columns(2)
    fov_um = c1.number_input(
        "Field of view width (µm)",
        min_value=0.0,
        value=0.0,
        step=1.0,
        help="Physical width of the image field, in micrometers.",
    )
    pixel_size_um = c2.number_input(
        "Pixel size (µm / pixel)",
        min_value=0.0,
        value=0.0,
        step=0.01,
        help="If provided, this overrides the FOV-based calibration.",
    )

results = []

# =========================
# Single image
# =========================
if upload_option == "Single image":
    uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg", "tiff"])

    if uploaded_file:
        st.markdown("---")
        st.markdown("### Result")

        result = process_image(uploaded_file, fov_um=fov_um, pixel_size_um=pixel_size_um)
        if result:
            results.append(result)

            st.markdown(f"#### {result['Imagem']}")

            if result["SemSegmentacao"]:
                st.image(result["Exibir"], caption="Original", use_container_width=True)
                st.warning("No segmentation was detected for this image.")
            else:
                col1, col2, col3 = st.columns(3)
                with col1:
                    st.image(result["Exibir"], caption="Original", use_container_width=True)
                with col2:
                    st.image(result["Segmentada"], caption="Segmentation", use_container_width=True)
                with col3:
                    st.image(result["Poligono"], caption="Polygon", use_container_width=True)

                render_metrics(result)

                st.markdown("### Export")
                st.download_button(
                    "Download segmented overlay (PNG)",
                    data=result["Segmentada"],
                    file_name=f"segmented_{result['Imagem']}.png",
                    mime="image/png",
                )

            st.markdown("---")
            render_feedback_block(result, prefix_key="single_")

# =========================
# Folder
# =========================
elif upload_option == "Image folder":
    uploaded_files = st.file_uploader(
        "Upload multiple images",
        type=["png", "jpg", "jpeg", "tiff"],
        accept_multiple_files=True,
    )

    if uploaded_files:
        st.markdown("---")
        st.markdown("### Processing")

        def process_wrapper(f):
            return process_image(f, fov_um=fov_um, pixel_size_um=pixel_size_um)

        with ThreadPoolExecutor(max_workers=4) as executor:
            processed = list(executor.map(process_wrapper, uploaded_files))

        falhas = [f.name for f, r in zip(uploaded_files, processed) if r and r.get("SemSegmentacao")]
        if falhas:
            st.warning(f"{len(falhas)} image(s) with no segmentation detected:\n\n- " + "\n- ".join(falhas))

        zip_images_buffer = BytesIO()
        with zipfile.ZipFile(zip_images_buffer, "w") as zip_file:
            for idx, result in enumerate(processed, start=1):
                if not result:
                    continue

                results.append(result)

                st.markdown("---")
                st.markdown(f"### Result {idx} · {result['Imagem']}")

                if result["SemSegmentacao"]:
                    st.image(result["Exibir"], caption="Original", use_container_width=True)
                    st.warning("No segmentation was detected for this image.")
                else:
                    col1, col2, col3 = st.columns(3)
                    with col1:
                        st.image(result["Exibir"], caption="Original", use_container_width=True)
                    with col2:
                        st.image(result["Segmentada"], caption="Segmentation", use_container_width=True)
                    with col3:
                        st.image(result["Poligono"], caption="Polygon", use_container_width=True)

                    render_metrics(result)

                    zip_file.writestr(f"segmentada_{result['Imagem']}.png", result["Segmentada"].getvalue())
                    zip_file.writestr(f"poligono_{result['Imagem']}.png", result["Poligono"].getvalue())

                render_feedback_block(result, prefix_key="folder_")

        zip_images_buffer.seek(0)

        # Summary table + exports
        if results:
            st.markdown("---")
            st.markdown("### Quantitative results")

            df = pd.DataFrame(
                [
                    {
                        "Image": r["Imagem"],
                        "Segmented Area (px²)": (
                            r["Área Segmentada (px²)"]
                            if (not r["SemSegmentacao"] and r["Área Segmentada (px²)"] is not None)
                            else "No Segmentation"
                        ),
                        "Segmented Area (µm²)": (
                            f"{r['Área Segmentada (µm²)']:.2f}"
                            if (not r["SemSegmentacao"] and r["Área Segmentada (µm²)"] is not None)
                            else ""
                        ),
                    }
                    for r in results
                ]
            )

            st.dataframe(df, use_container_width=True)

            excel_buffer = BytesIO()
            df.to_excel(excel_buffer, index=False)
            excel_buffer.seek(0)

            st.markdown("### Export results")
            c1, c2 = st.columns(2)
            with c1:
                st.download_button(
                    "Download table (Excel)",
                    data=excel_buffer,
                    file_name="segmentation_results.xlsx",
                    mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                    use_container_width=True,
                )
            with c2:
                st.download_button(
                    "Download segmented images (ZIP)",
                    data=zip_images_buffer,
                    file_name="segmented_images.zip",
                    mime="application/zip",
                    use_container_width=True,
                )