#!/usr/bin/env python3
"""
Streamlit UI for the photo editing pipeline.
Upload an image (or use a file path for DNG), run retrieve → LLM → apply, view result.

Run from project root:
  streamlit run app.py
"""
import sys
from pathlib import Path

_PROJECT_ROOT = Path(__file__).resolve().parent
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(_PROJECT_ROOT))

import numpy as np
import streamlit as st
import streamlit.components.v1 as components

from photo_editor.config import get_settings
from photo_editor.images import dng_to_rgb
from photo_editor.images.estimate_current_recipe import estimate_current_parameters
from photo_editor.pipeline.run import run_pipeline

# Fixed paths so "last result" matches across reruns (upload overwrites same file)
_STREAMLIT_INPUT_JPG_PATH = _PROJECT_ROOT / "_streamlit_input.jpg"
_STREAMLIT_INPUT_PNG_PATH = _PROJECT_ROOT / "_streamlit_input.png"
_STREAMLIT_INPUT_DNG_PATH = _PROJECT_ROOT / "_streamlit_input.dng"
_STREAMLIT_INPUT_HEIC_PATH = _PROJECT_ROOT / "_streamlit_input.heic"
_STREAMLIT_INPUT_HEIF_PATH = _PROJECT_ROOT / "_streamlit_input.heif"
# Use PNG output for UI preview to avoid JPEG quality loss.
_STREAMLIT_OUTPUT_PATH = _PROJECT_ROOT / "streamlit_output.png"
# Reversible toggle: set to False to restore top-1-only expert context.
_USE_MULTI_EXPERT_CONTEXT = True
_MULTI_EXPERT_CONTEXT_TOP_N = 1
_USE_BRIGHTNESS_GUARDRAIL = True


def _load_original_for_display(image_path: Path):
    """Load image for display. Use rawpy for DNG so 'Original' matches pipeline quality."""
    path = Path(image_path)
    if path.suffix.lower() == ".dng":
        rgb = dng_to_rgb(path, output_size=None)  # full resolution, same develop as pipeline
        rgb_u8 = (np.clip(rgb, 0, 1) * 255).astype(np.uint8)
        return rgb_u8
    # JPEG/PNG/HEIC/HEIF: Streamlit/Pillow can show from path (with plugin support).
    return str(path)


def _collapse_sidebar() -> None:
    """Collapse Streamlit sidebar via whichever toggle exists in this version."""
    components.html(
        """
<script>
const doc = window.parent.document;

function collapseIfOpen() {
  // Newer Streamlit versions expose a dedicated collapse button when sidebar is open.
  const closeBtn =
    doc.querySelector('[data-testid="stSidebarCollapseButton"]') ||
    doc.querySelector('button[aria-label="Close sidebar"]');
  if (closeBtn) {
    closeBtn.click();
    return true;
  }
  return false;
}

// Try immediately, then briefly retry in case elements mount after rerun.
if (!collapseIfOpen()) {
  let tries = 0;
  const interval = setInterval(() => {
    tries += 1;
    if (collapseIfOpen() || tries > 20) {
      clearInterval(interval);
    }
  }, 100);
}
</script>
""",
        height=0,
        width=0,
    )



def main() -> None:
    st.set_page_config(page_title="LumiGrade AI", page_icon="📷", layout="wide")

    if "sidebar_collapsed" not in st.session_state:
        st.session_state["sidebar_collapsed"] = False
    if "run_pipeline_requested" not in st.session_state:
        st.session_state["run_pipeline_requested"] = False
    if "selected_image_path" not in st.session_state:
        st.session_state["selected_image_path"] = ""
    if "is_processing" not in st.session_state:
        st.session_state["is_processing"] = False
    if "refresh_after_run" not in st.session_state:
        st.session_state["refresh_after_run"] = False

    collapse_css = ""
    if st.session_state["sidebar_collapsed"]:
        collapse_css = """
/* Force-hide sidebar after run is triggered; avoids version-specific JS toggles. */
[data-testid="stSidebar"] {
  display: none !important;
}
"""

    st.markdown(
        """
<style>
/* Keep custom styling minimal; rely on Streamlit theme config for core colors. */
[data-testid="stSidebar"] {
  background: #0b1220 !important;
  border-right: 1px solid rgba(148, 163, 184, 0.28);
}
[data-testid="stSidebar"] > div:first-child {
  background: #0b1220 !important;
}
[data-testid="stSidebar"] [data-testid="stVerticalBlock"] > div {
  box-shadow: inset -1px 0 0 rgba(148, 163, 184, 0.12);
}
.muted { color: #a8b3c7; font-size: 0.95rem; }
.section-title { font-size: 1.15rem; font-weight: 700; margin-bottom: 0.35rem; }
.action-card {
  border: 1px solid rgba(148, 163, 184, 0.22);
  border-radius: 10px;
  padding: 0.7rem 0.85rem;
  margin-bottom: 0.45rem;
  background: rgba(30, 41, 59, 0.28);
}
.json-box {
  border: 1px solid rgba(148, 163, 184, 0.2);
  border-radius: 10px;
  padding: 0.5rem 0.65rem;
  background: rgba(30, 41, 59, 0.2);
}
.loading-wrap {
  border: 1px solid rgba(148, 163, 184, 0.28);
  border-radius: 12px;
  padding: 0.9rem 1rem;
  background: rgba(30, 41, 59, 0.32);
  margin: 0.4rem 0 0.8rem 0;
}
.loading-head {
  display: flex;
  align-items: center;
  gap: 0.6rem;
  margin-bottom: 0.55rem;
  font-weight: 600;
}
.loader-spinner {
  width: 16px;
  height: 16px;
  border: 2px solid rgba(148, 163, 184, 0.25);
  border-top-color: #60A5FA;
  border-radius: 50%;
  animation: spin 0.8s linear infinite;
}
@keyframes spin {
  to { transform: rotate(360deg); }
}
.step-line {
  padding: 0.22rem 0;
  font-size: 0.94rem;
}

/* Move main title area slightly up */
[data-testid="stAppViewContainer"] .main .block-container {
  padding-top: 00.6rem !important;
}
h1 {
  margin-top: -0.25rem !important;
}

/* Push sidebar inputs a bit lower under the title */
[data-testid="stSidebar"] [data-testid="stSidebarContent"] {
  padding-top: 0 !important;
}
"""
        + collapse_css
        + """
</style>
""",
        unsafe_allow_html=True,
    )

    st.title("📷 LumiGrade AI")
    st.caption("Upload an image to get expert-informed edit recommendations and an instant enhanced result.")
    if st.session_state["sidebar_collapsed"]:
        if st.button("⚙️ Show Inputs", key="show_inputs_btn", disabled=st.session_state["is_processing"]):
            st.session_state["sidebar_collapsed"] = False
            st.rerun()

    # Config check
    s = get_settings()
    if not s.azure_search_configured():
        st.error("Azure AI Search not configured. Set AZURE_SEARCH_ENDPOINT and AZURE_SEARCH_KEY in .env")
        st.stop()
    if not s.azure_openai_configured():
        st.error("Azure OpenAI not configured. Set AZURE_OPENAI_* in .env")
        st.stop()

    # External editing API toggle has been removed from the UI for simplicity.
    # If you want to use the external API again, you can reintroduce a sidebar
    # control and wire it to this flag.
    use_editing_api = False

    image_path = Path(st.session_state["selected_image_path"]) if st.session_state["selected_image_path"] else None

    with st.sidebar:
        # Reliable spacing so only the Pipeline Inputs card moves down.
        st.markdown('<div style="height: 8.1rem;"></div>', unsafe_allow_html=True)
        with st.container(border=True):
            st.markdown('<div class="section-title">Pipeline Inputs</div>', unsafe_allow_html=True)
            uploaded = st.file_uploader(
                "Upload JPEG, PNG, DNG, HEIC, or HEIF",
                type=["jpg", "jpeg", "png", "dng", "heic", "heif"],
                help="Upload JPEG/PNG/DNG/HEIC/HEIF to run the edit recommendation pipeline.",
            )
            if uploaded is not None:
                suffix = Path(uploaded.name).suffix.lower()
                if suffix == ".dng":
                    target = _STREAMLIT_INPUT_DNG_PATH
                elif suffix == ".heic":
                    target = _STREAMLIT_INPUT_HEIC_PATH
                elif suffix == ".heif":
                    target = _STREAMLIT_INPUT_HEIF_PATH
                elif suffix == ".png":
                    target = _STREAMLIT_INPUT_PNG_PATH
                else:
                    target = _STREAMLIT_INPUT_JPG_PATH
                target.write_bytes(uploaded.getvalue())
                image_path = target
                st.session_state["selected_image_path"] = str(target)

            run_clicked = st.button(
                "✨ Generate Edit Recommendations",
                type="primary",
                use_container_width=True,
                disabled=st.session_state["is_processing"],
            )
            status = st.empty()
            if image_path is None:
                status.info("Provide an image to run.")

            if run_clicked and image_path is not None:
                # Mark busy before rerun so any control rendered on next pass is disabled.
                st.session_state["is_processing"] = True
                st.session_state["sidebar_collapsed"] = True
                st.session_state["run_pipeline_requested"] = True
                st.rerun()
    should_run_pipeline = st.session_state.pop("run_pipeline_requested", False)
    if should_run_pipeline and st.session_state["selected_image_path"]:
        image_path = Path(st.session_state["selected_image_path"])
    if should_run_pipeline and image_path is not None:
        st.session_state["is_processing"] = True
        _collapse_sidebar()
        loading_box = st.empty()

        def _render_loading(current_stage: str, state: str = "running") -> None:
            stage_order = ["retrieving", "consulting", "applying"]
            stage_labels = {
                "retrieving": "Analyzing similar expert edits",
                "consulting": "Generating personalized recommendations",
                "applying": "Rendering your enhanced preview",
            }
            current_idx = stage_order.index(current_stage) if current_stage in stage_order else 0

            if state == "done":
                title = "Done"
                spinner_html = ""
            elif state == "failed":
                title = "Pipeline failed"
                spinner_html = ""
            else:
                title = "Running pipeline"
                spinner_html = '<span class="loader-spinner"></span>'

            lines = []
            for i, key in enumerate(stage_order):
                if state == "done":
                    icon = "✅"
                elif state == "failed" and i > current_idx:
                    icon = "⏳"
                else:
                    icon = "✅" if i < current_idx else ("🔄" if i == current_idx and state == "running" else "⏳")
                lines.append(f'<div class="step-line">{icon} {stage_labels[key]}</div>')

            loading_box.markdown(
                f"""
<div class="loading-wrap">
  <div class="loading-head">{spinner_html}<span>{title}</span></div>
  {''.join(lines)}
</div>
""",
                unsafe_allow_html=True,
            )

        _render_loading("retrieving", "running")
        try:
            current_params = estimate_current_parameters(image_path)
            result = run_pipeline(
                image_path,
                _STREAMLIT_OUTPUT_PATH,
                top_k=50,
                top_n=1,
                use_editing_api=use_editing_api,
                use_multi_expert_context=_USE_MULTI_EXPERT_CONTEXT,
                context_top_n=_MULTI_EXPERT_CONTEXT_TOP_N,
                use_brightness_guardrail=_USE_BRIGHTNESS_GUARDRAIL,
                progress_callback=lambda stage: _render_loading(stage, "running"),
            )
            if result.get("success"):
                st.session_state["pipeline_result"] = result
                st.session_state["pipeline_output_path"] = _STREAMLIT_OUTPUT_PATH
                st.session_state["pipeline_input_path"] = str(image_path)
                st.session_state["pipeline_current_params"] = current_params
                status.success("Done!")
                _render_loading("applying", "done")
            else:
                st.session_state.pop("pipeline_result", None)
                st.session_state.pop("pipeline_output_path", None)
                st.session_state.pop("pipeline_input_path", None)
                st.session_state.pop("pipeline_current_params", None)
                status.error("Editing step failed.")
                _render_loading("applying", "failed")
        except Exception as e:
            status.error("Pipeline failed.")
            st.exception(e)
            st.session_state.pop("pipeline_result", None)
            st.session_state.pop("pipeline_output_path", None)
            st.session_state.pop("pipeline_input_path", None)
            st.session_state.pop("pipeline_current_params", None)
            _render_loading("consulting", "failed")
        finally:
            st.session_state["is_processing"] = False
            # Button states are computed at render time; rerun once so controls
            # immediately reflect processing completion (re-enable Show Inputs).
            st.session_state["refresh_after_run"] = True

    if st.session_state.get("refresh_after_run"):
        st.session_state["refresh_after_run"] = False
        st.rerun()

    display_input_path = image_path
    if display_input_path is None and st.session_state.get("pipeline_input_path"):
        display_input_path = Path(st.session_state["pipeline_input_path"])

    with st.container(border=True):
        st.subheader("Results Dashboard")
        st.markdown("### 📊 Pipeline Analysis & Recommendations")

        # Show result if available
        if (
            display_input_path is not None
            and st.session_state.get("pipeline_result")
            and st.session_state.get("pipeline_input_path") == str(display_input_path)
        ):
            result = st.session_state["pipeline_result"]
            out_path = st.session_state["pipeline_output_path"]

            if out_path.exists():
                summary = result.get("summary", "")
                suggested = result.get("suggested_edits", {})
                expert_id = result.get("expert_image_id", "")
                current_params = st.session_state.get("pipeline_current_params") or {}

                with st.expander("AI Analysis Summary", expanded=True):
                    st.markdown(summary)

                with st.expander("Parameters: Details", expanded=True):
                    st.markdown("#### Parameters: Current vs Suggested vs Delta")
                    keys = [
                        "exposure",
                        "contrast",
                        "highlights",
                        "shadows",
                        "whites",
                        "blacks",
                        "temperature",
                        "tint",
                        "vibrance",
                        "saturation",
                    ]
                    rows = []
                    for k in keys:
                        cur = current_params.get(k, None)
                        sug = suggested.get(k, None)
                        try:
                            cur_f = float(cur) if cur is not None else None
                        except Exception:
                            cur_f = None
                        try:
                            sug_f = float(sug) if sug is not None else None
                        except Exception:
                            sug_f = None
                        delta = (sug_f - cur_f) if (sug_f is not None and cur_f is not None) else None
                        rows.append(
                            {
                                "parameter": k,
                                "current_estimated": cur_f,
                                "suggested": sug_f,
                                "delta": delta,
                            }
                        )
                    st.dataframe(rows, use_container_width=True, hide_index=True)
                    st.caption('“Current” values are estimated from pixels (not true Lightroom sliders).')
            else:
                st.info("Run the pipeline to populate results.")
        else:
            st.info("Run the pipeline from the left pane to view analysis and recommendations.")

    # Keep this full-width and at the bottom, per request.
    if (
        display_input_path is not None
        and st.session_state.get("pipeline_result")
        and st.session_state.get("pipeline_input_path") == str(display_input_path)
    ):
        result = st.session_state["pipeline_result"]
        out_path = st.session_state["pipeline_output_path"]
        if out_path.exists():
            st.markdown("---")
            st.subheader("Original vs Result")
            col_orig, col_result = st.columns(2)
            with col_orig:
                st.image(_load_original_for_display(display_input_path), caption="Original", use_container_width=True)
            with col_result:
                st.image(str(out_path), caption="Edited", use_container_width=True)


if __name__ == "__main__":
    main()