# streamlit_app.py — AI Motion Path Planner (Frontend)
# Backend (FastAPI) endpoints:
# /health, /constraints/upload, /constraints/reset, /detect, /plan, /plan_report, /analyze_xyz

import os, io, uuid, base64, json, requests, datetime
from typing import List, Tuple, Optional, Dict, Any
from PIL import Image, ImageDraw, ImageFont
import streamlit as st

# ================= Theme =================
try:
    st._config.set_option("theme.base", "dark")
    st._config.set_option("theme.primaryColor", "#ff7a1a")
    st._config.set_option("theme.backgroundColor", "#0e0f12")
    st._config.set_option("theme.secondaryBackgroundColor", "#16181d")
    st._config.set_option("theme.textColor", "#e6e8ea")
except Exception:
    pass

# ================= Env & Session Defaults =================
def _env(k: str, default: str = "") -> str:
    return (os.getenv(k) or default).strip().strip("'\"")

DEFAULT_API_BASE = _env("MotionPath_AI_API", "http://127.0.0.1:8000").rstrip("/")
HF_TOKEN = _env("MotionPath_AI_TOKEN", "")
HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
TIMEOUT_S = 120

_defaults = {
    "API_BASE": DEFAULT_API_BASE,
    "size_preset": "1280×720",
    "jpeg_q": 90,
    "aspect_mode": "16:9 crop",  # "Original", "16:9 crop", "4:3 crop", "1:1 crop"
    "primary_labels": ["plate", "food"],
    "extra_labels": "",
    "box_pick": "Highest score",  # "Highest score", "Largest area", "Center-most"
    "want_notes": True,
    "variants": 3,
}
for k, v in _defaults.items():
    st.session_state.setdefault(k, v)

# ================= HTTP helpers =================
def _api_base() -> str:
    return (st.session_state.get("API_BASE") or DEFAULT_API_BASE).rstrip("/")

def _post_json(path: str, payload: dict, *, stream: bool=False):
    url = f"{_api_base()}{path}"
    return requests.post(url, json=payload, headers=HEADERS, timeout=TIMEOUT_S, stream=stream)

def _post_file(path: str, field: str, file_bytes: bytes, filename: str):
    url = f"{_api_base()}{path}"
    files = {field: (filename, file_bytes)}
    return requests.post(url, files=files, headers=HEADERS, timeout=TIMEOUT_S)

def _get(path: str):
    url = f"{_api_base()}{path}"
    return requests.get(url, headers=HEADERS, timeout=30)

def _success(msg: str): st.success(msg, icon="✅")
def _warn(msg: str): st.warning(msg, icon="⚠️")
def _err(msg: str): st.error(msg, icon="❌")

# ================= Image utilities =================
_ASPECTS = {
    "16:9 crop": 16/9,
    "4:3 crop": 4/3,
    "1:1 crop": 1/1,
}

def _center_crop_aspect(img: Image.Image, target_ratio: float) -> Image.Image:
    W, H = img.size
    cur_ratio = W / H
    if abs(cur_ratio - target_ratio) < 1e-3:
        return img
    if cur_ratio > target_ratio:
        new_w = int(H * target_ratio)
        x0 = (W - new_w) // 2
        return img.crop((x0, 0, x0 + new_w, H))
    else:
        new_h = int(W / target_ratio)
        y0 = (H - new_h) // 2
        return img.crop((0, y0, W, y0 + new_h))

def _preprocess_image(file, target_w=1280, target_h=720, aspect_mode="16:9 crop") -> Image.Image:
    img = Image.open(file).convert("RGB")
    if aspect_mode in _ASPECTS:
        img = _center_crop_aspect(img, _ASPECTS[aspect_mode])
    img = img.resize((target_w, target_h), Image.Resampling.LANCZOS)
    return img

def _b64_from_pil(img: Image.Image, jpeg_quality=90) -> str:
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=int(jpeg_quality))
    return base64.b64encode(buf.getvalue()).decode("utf-8")

def _b64_from_file(file, target_w=1280, target_h=720, jpeg_quality=90, aspect_mode="16:9 crop") -> Tuple[str, Image.Image]:
    img = _preprocess_image(file, target_w, target_h, aspect_mode)
    return _b64_from_pil(img, jpeg_quality), img

# ================= Overlay drawing =================
def _draw_boxes(
    img: Image.Image,
    boxes: List[Tuple[float, float, float, float]],
    labels: Optional[List[str]] = None,
    scores: Optional[List[float]] = None,
    primary_idx: Optional[int] = None,
) -> Image.Image:
    im = img.copy()
    draw = ImageDraw.Draw(im)
    W, H = im.size
    try:
        font = ImageFont.load_default()
    except Exception:
        font = None

    def _text_w(t: str) -> int:
        try:
            return int(draw.textlength(t, font=font))
        except Exception:
            return 7 * len(t)

    for i, b in enumerate(boxes):
        x1 = int(b[0] * W); y1 = int(b[1] * H); x2 = int(b[2] * W); y2 = int(b[3] * H)
        color = (255, 180, 60) if i != primary_idx else (80, 220, 120)
        width = 3 if i != primary_idx else 5
        draw.rectangle([x1, y1, x2, y2], outline=color, width=width)
        tag = ""
        if labels and i < len(labels) and labels[i]:
            tag = labels[i]
        if scores and i < len(scores) and scores[i] is not None:
            tag = f"{tag} {scores[i]*100:.1f}%" if tag else f"{scores[i]*100:.1f}%"
        if tag:
            tw = _text_w(tag)
            draw.rectangle([x1, max(0, y1-18), x1 + tw + 10, y1], fill=color)
            draw.text((x1 + 5, y1 - 16), tag, fill=(0, 0, 0), font=font)
    return im

def _pick_primary_box(
    boxes: List[Tuple[float, float, float, float]],
    scores: List[float],
    strategy: str = "Highest score",
) -> Optional[int]:
    if not boxes:
        return None
    if strategy == "Largest area":
        areas = [(b[2]-b[0]) * (b[3]-b[1]) for b in boxes]
        return int(max(range(len(boxes)), key=lambda i: areas[i]))
    if strategy == "Center-most":
        def center_dist(b):
            cx = 0.5 * (b[0]+b[2]); cy = 0.5 * (b[1]+b[3])
            return (cx-0.5)**2 + (cy-0.5)**2
        return int(min(range(len(boxes)), key=lambda i: center_dist(boxes[i])))
    return int(max(range(len(boxes)), key=lambda i: scores[i] if scores and i < len(scores) else -1e9))

def _detect_boxes_backend(image_b64: str, labels: List[str]) -> Optional[dict]:
    payload = {"image_b64": image_b64, "detect_query": ", ".join(labels)}
    for endpoint in ("/detect", "/detect_preview"):
        try:
            r = _post_json(endpoint, payload, stream=False)
            if r.status_code == 200:
                data = r.json()
                if "boxes" in data:
                    return data
        except Exception:
            pass
    return None

# ================= Sidebar (backend + constraints) =================
with st.sidebar:
    st.title("Backend")
    st.caption("Enter a private backend URL if needed.")
    url_in = st.text_input(
        "Backend URL",
        value=st.session_state["API_BASE"],
        placeholder="https://<org>-backend.hf.space"
    )
    colb1, colb2 = st.columns([1,1])
    if colb1.button("Use URL"):
        st.session_state["API_BASE"] = url_in.strip().rstrip("/")
        _success(f"Using {st.session_state['API_BASE']}")
        st.rerun()
    if colb2.button("Health"):
        try:
            r = _get("/health")
            if r.status_code == 200:
                _success("Backend reachable.")
                st.json(r.json(), expanded=False)
            else:
                _err(f"{r.status_code}: {r.text[:300]}")
        except Exception as e:
            _err(str(e))

    st.markdown("---")
    st.subheader("Rig Limits (export.txt)")
    st.caption("Upload Flair export.txt to apply pan/tilt/roll and track limits.")
    export_file = st.file_uploader("Upload export.txt", type=["txt"], key="exp_txt")
    cola, colb = st.columns(2)
    with cola:
        if st.button("Load constraints"):
            if not export_file:
                _err("Select export.txt first.")
            else:
                r = _post_file("/constraints/upload", "file", export_file.read(), export_file.name)
                if r.status_code == 200:
                    _success("Constraints loaded.")
                    st.json(r.json(), expanded=False)
                else:
                    _err(f"{r.status_code}: {r.text[:300]}")
    with colb:
        if st.button("Reset constraints"):
            r = _get("/constraints/reset")
            if r.status_code in (200, 204):
                _success("Constraints cleared.")
            else:
                _err(f"{r.status_code}: {r.text[:300]}")

# ================= Main =================
st.title("AI Motion Path Planner for Flair")
st.caption("Describe the move in English. Generate a Flair .xyz path, apply limits, and get kinematic peaks.")

with st.expander("How it works"):
    st.markdown(
        """
1) Upload an image (what the camera sees).  
2) Describe the move: “Orbit 90° in 4s, jib up 0.6 m, radius 0.5 m, start 180°, 25 fps.”  
3) Use **Generate + Report** for variants and limit checks.  
4) Use **Generate Path** to download the `.xyz`.
        """
    )

# ---------- Image optimization ----------
with st.expander("Image optimization"):
    st.session_state["size_preset"] = st.selectbox(
        "Target size",
        ["960×540", "1280×720", "1920×1080"],
        index=["960×540", "1280×720", "1920×1080"].index(st.session_state["size_preset"])
    )
    st.session_state["jpeg_q"] = st.slider("JPEG quality", 70, 95, int(st.session_state["jpeg_q"]))
    st.session_state["aspect_mode"] = st.selectbox(
        "Frame aspect",
        ["Original", "16:9 crop", "4:3 crop", "1:1 crop"],
        index=["Original", "16:9 crop", "4:3 crop", "1:1 crop"].index(st.session_state["aspect_mode"])
    )

def _target_hw():
    preset = st.session_state.get("size_preset", "1280×720")
    if "960×540" in preset: return 960, 540
    if "1920×1080" in preset: return 1920, 1080
    return 1280, 720

# ---------- Inputs ----------
col_img, col_text = st.columns([1, 1])
with col_img:
    img_file = st.file_uploader("Camera image (JPG/PNG)", type=["jpg","jpeg","png"], key="scene")
with col_text:
    presets = [
        "Orbit 90° in 4 seconds, jib up 0.6 m, radius 0.5 m, start 180°, 25 fps.",
        "3.5s: orbit -60°, jib down 0.2 m; radius 0.45 m; start 180°; 25 fps.",
        "Two segments: 2s orbit 40°, then 2s orbit 50° while jib up 0.3 m; radius 0.5 m; 25 fps."
    ]
    preset_pick = st.selectbox("Presets (optional)", options=["(none)"] + presets, index=0)
    instr = st.text_area(
        "Describe the move",
        value=(preset_pick if preset_pick != "(none)" else ""),
        height=120,
        placeholder="Example: Orbit 90° in 4s, jib up 0.6 m, radius 0.5 m, start 180°, 25 fps."
    )
    st.markdown("**Subject labels**")
    all_suggestions = ["plate", "food", "burger", "person", "bottle", "product"]
    st.session_state["primary_labels"] = st.multiselect(
        "Primary labels",
        options=sorted(set(all_suggestions + st.session_state.get("primary_labels", []))),
        default=st.session_state["primary_labels"],
        key="primary_labels_widget",
    )
    st.session_state["extra_labels"] = st.text_input(
        "Additional labels (comma-separated)",
        value=st.session_state["extra_labels"],
        placeholder="e.g., chicken, arugula, sauce"
    )
    st.session_state["box_pick"] = st.selectbox(
        "Primary box strategy",
        ["Highest score", "Largest area", "Center-most"],
        index=["Highest score", "Largest area", "Center-most"].index(st.session_state["box_pick"])
    )
    st.session_state["want_notes"] = st.checkbox("Show AI Notes in report", value=bool(st.session_state["want_notes"]))
    st.session_state["variants"] = st.slider("How many variants", 1, 5, int(st.session_state["variants"]))

def _gather_labels() -> List[str]:
    labels = list(st.session_state.get("primary_labels") or [])
    extra = [x.strip() for x in (st.session_state.get("extra_labels") or "").split(",") if x.strip()]
    return [l for l in (labels + extra) if l][:12]

# ---------- Detection preview card ----------
det_card = st.container()
if img_file:
    tw, th = _target_hw()
    img_proc = _preprocess_image(
        img_file, tw, th,
        aspect_mode=st.session_state["aspect_mode"],
    )
    b64_img = _b64_from_pil(img_proc, int(st.session_state["jpeg_q"]))
    with det_card:
        st.subheader("Preview")
        st.caption("Subject placement and detection overlay.")
        boxes_img = None
        try:
            labels = _gather_labels() or ["subject"]
            data = _detect_boxes_backend(b64_img, labels)
            if data and isinstance(data.get("boxes"), list) and len(data["boxes"]) > 0:
                boxes = data["boxes"]
                lbls = data.get("labels", ["object"] * len(boxes))
                scrs = data.get("scores", [0.0] * len(boxes))
                pick = _pick_primary_box(boxes, scrs, st.session_state["box_pick"])
                vis = _draw_boxes(img_proc, boxes, lbls, scrs, primary_idx=pick)
                boxes_img = vis
        except Exception as e:
            _warn(f"Detection preview issue: {e}")

        max_w = 900
        if boxes_img is not None:
            st.image(boxes_img, caption="Detection overlay", width=max_w)
        else:
            st.image(img_proc, caption="Image (no detection overlay available)", width=max_w)

# ---------- Actions ----------
c1, c2 = st.columns(2)
btn_generate = c1.button("Generate Path (.xyz)", type="primary", use_container_width=True)
btn_report   = c2.button("Generate + Report (variants)", use_container_width=True)

def _require_inputs() -> bool:
    if not img_file:
        _err("Upload a camera image.")
        return False
    if not instr or not instr.strip():
        _err("Describe the move.")
        return False
    return True

def _payload_prebuilt(include_variants: bool) -> dict:
    tw, th = _target_hw()
    img_b64, _img = _b64_from_file(
        img_file, tw, th,
        jpeg_quality=int(st.session_state["jpeg_q"]),
        aspect_mode=st.session_state["aspect_mode"],
    )
    labels = _gather_labels() or ["subject"]
    payload = {
        "instruction": instr.strip(),
        "image_b64": img_b64,
        "detect_query": ", ".join(labels),
        "want_notes": bool(st.session_state["want_notes"]),
    }
    if include_variants:
        payload["variants"] = int(st.session_state["variants"])
    return payload

# Generate-only (.xyz)
if btn_generate and _require_inputs():
    with st.spinner("Generating motion path…"):
        r = _post_json("/plan", _payload_prebuilt(include_variants=False), stream=True)
    if r.status_code == 200:
        xyz_bytes = r.raw.read() if hasattr(r, "raw") else r.content
        fname = f"motion_{uuid.uuid4().hex[:8]}.xyz"
        _success("Path generated.")
        st.download_button("Download .xyz", data=xyz_bytes, file_name=fname, mime="text/plain", use_container_width=True)
        with st.expander("Preview (first lines)"):
            try:
                st.code("\n".join(xyz_bytes.decode("utf-8", errors="ignore").splitlines()[:60]), language="text")
            except Exception:
                pass
    else:
        try:
            _err(f"{r.status_code}: {r.text[:600]}")
        except Exception:
            _err(f"{r.status_code}: response error")

# Generate + Report (variants)
if btn_report and _require_inputs():
    with st.spinner("Generating paths and report…"):
        r = _post_json("/plan_report", _payload_prebuilt(include_variants=True), stream=False)
    if r.status_code == 200:
        data = r.json()
        _success("Report ready.")

        # Quick summary header
        hdr = st.container()
        with hdr:
            cols = st.columns(3)
            cols[0].metric("Variants", f"{len(data.get('variants', []))}")
            cols[1].metric("Param source", f"{data.get('param_source','unknown')}")
            ck = data.get("constraints_keys") or []
            cols[2].metric("Constraints parsed", f"{len(ck)} keys")

        # Save full JSON report
        with st.expander("Report JSON"):
            s = json.dumps(data, ensure_ascii=False, indent=2)
            st.code(s, language="json")
            st.download_button(
                "Download report.json",
                data=s.encode("utf-8"),
                file_name=f"report_{datetime.datetime.utcnow().strftime('%Y%m%dT%H%M%SZ')}.json",
                mime="application/json",
                use_container_width=True
            )

        # Detection summary
        with st.expander("Detection summary"):
            st.json(data.get("detection_summary", {}), expanded=False)

        # Primary block (if backend returns)
        primary = data.get("primary") or {}
        if primary:
            st.subheader("Primary (auto-tuned)")
            pc1, pc2, pc3 = st.columns(3)
            pc1.metric("Duration (s)", f"{primary.get('duration_s',0.0):.2f}")
            pc2.metric("Orbit measured (deg)", f"{primary.get('orbit_measured',0.0):.1f}")
            pc3.metric("Constraints", "Passed ✅" if primary.get("constraints_passed", False) else "Failed ❌")
            if "autotune" in primary:
                with st.expander("Auto-tune details"):
                    st.write(primary.get("autotune", {}))
            if "ratios" in primary:
                with st.expander("Limit ratios (measured / limit)"):
                    st.json(primary.get("ratios", {}), expanded=False)

            st.markdown("---")

        # Variants
        variants = data.get("variants", []) or []
        if variants:
            st.subheader("Variants")
            for v in variants:
                title = f"Variant {int(v.get('variant', 0))} — {v.get('duration_s', 0):.2f}s @ {int(v.get('fps',0))} fps"
                with st.expander(title):
                    cols = st.columns(3)
                    with cols[0]:
                        st.metric("Orbit requested (deg)", f"{v.get('orbit_requested', 0):.1f}")
                        st.metric("Orbit measured (deg)", f"{v.get('orbit_measured', 0):.1f}")
                    with cols[1]:
                        st.metric("Jib requested (m)", f"{v.get('jib_requested_m', 0):.3f}")
                        st.metric("Jib measured (m)", f"{v.get('jib_measured_m', 0):.3f}")
                    with cols[2]:
                        st.metric("Constraints", "Passed ✅" if v.get("constraints_passed", False) else "Failed ❌")
                        if v.get("final_params"):
                            st.caption("Final params used by variant:")
                            st.code(json.dumps(v["final_params"], indent=2), language="json")

                    kin = v.get("kinematics", {}) or {}
                    pan = kin.get("pan", {}) or {}
                    tilt = kin.get("tilt", {}) or {}
                    st.write("**Kinematics — Pan**")
                    cpk = st.columns(4)
                    cpk[0].metric("Max Speed (°/s)", f"{pan.get('max_dps', 0.0):.2f}")
                    cpk[1].metric("Max Acc (°/s²)", f"{pan.get('max_dps2', 0.0):.2f}")
                    cpk[2].metric("Max Jerk (°/s³)", f"{pan.get('max_dps3', 0.0):.2f}")
                    cpk[3].metric("Avg Speed (°/s)", f"{pan.get('avg_dps', 0.0):.2f}")

                    st.write("**Kinematics — Tilt**")
                    ctk = st.columns(4)
                    ctk[0].metric("Max Speed (°/s)", f"{tilt.get('max_dps', 0.0):.2f}")
                    ctk[1].metric("Max Acc (°/s²)", f"{tilt.get('max_dps2', 0.0):.2f}")
                    ctk[2].metric("Max Jerk (°/s³)", f"{tilt.get('max_dps3', 0.0):.2f}")
                    ctk[3].metric("Avg Speed (°/s)", f"{tilt.get('avg_dps', 0.0):.2f}")

                    if v.get("ratios"):
                        with st.expander("Limit ratios"):
                            st.json(v["ratios"], expanded=False)

                    if not v.get("constraints_passed", False):
                        st.error("Violated:")
                        st.write(v.get("violated", []))
                        if v.get("suggestions"):
                            st.info("Suggestions:")
                            st.write(v.get("suggestions", []))

        # AI Notes
        ai_notes = data.get("ai_notes", []) or []
        if ai_notes:
            st.subheader("AI Notes")
            for n in ai_notes:
                st.write(f"• {n}")
        else:
            st.caption("AI Notes unavailable.")

        st.caption("Use “Generate Path (.xyz)” for a downloadable .xyz. The report shows limit checks and variants.")
    else:
        _err(f"{r.status_code}: {r.text[:600]}")

# ---------- Analyze existing .xyz ----------
st.markdown("---")
st.subheader("Analyze a .xyz file")
st.caption("Upload a Flair .xyz to get duration, FPS, orbit, jib, and pan/tilt kinematic peaks.")
xyz_up = st.file_uploader("Upload .xyz", type=["xyz"], key="xyz_file")
if xyz_up and st.button("Analyze"):
    with st.spinner("Analyzing…"):
        r = _post_file("/analyze_xyz", "file", xyz_up.read(), xyz_up.name)
    if r.status_code == 200:
        _success("Analysis complete.")
        st.json(r.json(), expanded=False)
    else:
        _err(f"{r.status_code}: {r.text[:600]}")

# ---------- Footer ----------
st.markdown("---")
st.caption(
    "Set env vars for this Space:\n"
    "• MotionPath_AI_API = https://<your-private-backend>.hf.space\n"
    "• MotionPath_AI_TOKEN = <same token as backend>\n"
    "Shows concise AI Notes (not chain-of-thought)."
)