import sys, os, time  # time kept for perf logging in inference calls
sys.path.append(os.path.dirname(os.path.dirname(__file__)))   # import inference/

import streamlit as st
from PIL import Image
from inference.logger import get_logger

log = get_logger("app")

st.set_page_config(page_title="Yugensys Vision AI — Manufacturing", layout="wide")

from inference.detect import detect_assets, _model as get_poc1_model
from inference.classify import classify_qc, get_efficientnet, get_yolo

# ---------------------------------------------------------------------------
# Sidebar + main UI — models load lazily on-demand via @st.cache_resource
# ---------------------------------------------------------------------------
st.sidebar.title("Yugensys Vision AI")
screen = st.sidebar.radio(
    "Select POC",
    ["Asset Identification", "Quality Inspection (Auto Fallback)"],
)

# ---------------------------------------------------------------------------
# POC screens
# ---------------------------------------------------------------------------
if screen == "Asset Identification":
    st.header("Manufacturing Asset Identification")

    conf_pct = st.slider(
        "Detection Confidence Threshold",
        min_value=0, max_value=100, value=25, step=1,
    )
    conf = conf_pct / 100.0

    up = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"], key="poc1")
    if up:
        img = Image.open(up).convert("RGB")
        with st.spinner("Detecting..."):
            t0 = time.perf_counter()
            annotated, dets = detect_assets(img, conf=conf)
            log.info("detect_assets — %.3fs  detections=%d", time.perf_counter() - t0, len(dets))
        c1, c2 = st.columns(2)
        c1.image(img, caption="Input", use_column_width=True)
        c2.image(annotated, caption="Detections", use_column_width=True)
        if dets:
            st.success(f"**{dets[0]['name']}** — Confidence: {dets[0]['confidence']}%")
            st.table(dets)
        else:
            st.warning("No assets detected. Try another image.")

elif screen == "Quality Inspection (Auto Fallback)":
    st.header("Visual Quality Inspection")

    threshold_pct = st.slider(
        "EfficientNet Confidence Threshold (Triggers YOLO Fallback)",
        min_value=0, max_value=100, value=80, step=1,
    )
    threshold = threshold_pct / 100.0

    up = st.file_uploader("Upload a product image", type=["jpg", "jpeg", "png"], key="poc2_auto")
    if up:
        img = Image.open(up).convert("RGB")
        with st.spinner("Inspecting..."):
            t0 = time.perf_counter()
            res = classify_qc(img, threshold=threshold)
            log.info("classify_qc — %.3fs  source=%s  status=%s  conf=%.1f%%",
                     time.perf_counter() - t0, res.get("source"), res.get("status"), res.get("confidence", 0))
        st.image(img, caption="Input", width=400)

        if res.get("warning"):
            st.warning(res["warning"])

        if res["status"] == "QC Approved":
            st.success(f"QC Status: **{res['status']}** · Defect: **{res['defect']}** · Confidence: {res['confidence']}%")
        else:
            st.error(f"QC Status: **{res['status']}** · Defect: **{res['defect']}** · Confidence: {res['confidence']}%")
        st.info(f"Model Source: {res.get('source', 'Unknown')}")