"""Gradio UI — image analysis, video processing, and live webcam streaming."""
import gradio as gr
import cv2
import numpy as np
from PIL import Image as PILImage
from pipeline import process_image, process_video
from core import detector, depth, pose, risk_engine
from core.depth import SlopeSmoother
from core.detector import ObstacleTracker
from core.guidance import GuidanceEngine
from renderers import overlay as overlay_renderer

# ─── Image handler ───

def handle_image(input_image):
    if input_image is None:
        return None, "Upload an image."
    img_bgr = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
    rendered, risk_dict, guidance = process_image(img_bgr, mode="overlay")

    rpt = f"## Risk: **{risk_dict['risk']}** (Score: {risk_dict['score']})\n\n"
    rpt += f"**Terrain:** {risk_dict['terrain']} ({risk_dict['terrain_slope']:.1f}°)\n\n"
    gs = risk_dict.get('gait_summary', {})
    if gs:
        rpt += "| Metric | Value |\n|---|---|\n"
        for k, v in gs.items():
            rpt += f"| {k.title()} | {v:.1f}° |\n"
    rpt += "\n**Risk Factors:**\n"
    rpt += "\n".join(f"- ⚠️ {r}" for r in risk_dict['reasons']) if risk_dict['reasons'] else "- ✅ None"
    rpt += f"\n\n**Voice:** {guidance['voice']}"

    return PILImage.fromarray(cv2.cvtColor(rendered, cv2.COLOR_BGR2RGB)), rpt


# ─── Video handler ───

def handle_video(video, mode):
    if video is None:
        return None
    return process_video(video, mode=mode)


# ─── Live webcam handler (Gradio streaming) ───

# Persistent state for live stream
_live_state = {
    'guide': None,
    'smoother': None,
    'tracker': None,
    'depth': None,
    'depth_mini': None,
    'counter': 0,
}


def _reset_live():
    _live_state['guide'] = GuidanceEngine()
    _live_state['smoother'] = SlopeSmoother()
    _live_state['tracker'] = ObstacleTracker()
    _live_state['depth'] = None
    _live_state['depth_mini'] = None
    _live_state['counter'] = 0


def handle_webcam_frame(frame):
    """Process a single webcam frame from Gradio's streaming input."""
    if frame is None:
        return None, ""

    # Init state on first frame
    if _live_state['guide'] is None:
        _reset_live()

    img_bgr = cv2.cvtColor(np.array(frame), cv2.COLOR_RGB2BGR)
    h, w = img_bgr.shape[:2]
    rgb = np.array(frame)

    obstacles = detector.detect(img_bgr, track=True)
    new_obs, closing_obs, _ = _live_state['tracker'].update(obstacles)
    gait, landmarks, foot_y = pose.analyze(rgb, w, h)

    _live_state['counter'] += 1
    if _live_state['counter'] % 5 == 0 or _live_state['depth'] is None:
        _live_state['depth'] = depth.estimate_depth(rgb, h, w)
        raw_s, raw_d, raw_t, _ = depth.estimate_slope(_live_state['depth'], h, w, foot_y)
        _live_state['smoother'].update(raw_s, raw_d, raw_t)
        _live_state['depth_mini'] = overlay_renderer.render_depth_mini(_live_state['depth'], w, h)

    sm = _live_state['smoother']
    risk_dict = risk_engine.assess(sm.angle, sm.direction, gait, len(obstacles))
    guidance = _live_state['guide'].compute(
        sm.angle, sm.direction, sm.terrain, obstacles,
        slope_trend=sm.trend, new_obstacles=new_obs, closing_obstacles=closing_obs)

    rendered = overlay_renderer.render(
        img_bgr, obstacles, gait, landmarks, risk_dict, guidance, _live_state['depth_mini'])

    status = (f"**{risk_dict['risk']}** | Slope: {sm.direction} {sm.angle:.0f}° "
              f"[{sm.trend}] | Obs: {len(obstacles)} | {guidance['voice'][:80]}")

    return PILImage.fromarray(cv2.cvtColor(rendered, cv2.COLOR_BGR2RGB)), status


# ─── Build UI ───

with gr.Blocks(title="Navigation Assist") as demo:
    gr.Markdown("# 🦯 Vision-Based Navigation Assistance")

    with gr.Tab("📷 Image"):
        with gr.Row():
            img_in = gr.Image(type="pil", label="Upload Image")
            img_out = gr.Image(label="Analysis")
        report = gr.Markdown()
        gr.Button("🔍 Analyze", variant="primary").click(
            fn=handle_image, inputs=img_in, outputs=[img_out, report])

    with gr.Tab("🎥 Video"):
        with gr.Row():
            vid_in = gr.Video(label="Upload Video")
            vid_out = gr.Video(label="Output")
        vid_mode = gr.Radio(["overlay", "blind_nav"], value="overlay", label="Render Mode")
        gr.Button("🔍 Process", variant="primary").click(
            fn=handle_video, inputs=[vid_in, vid_mode], outputs=vid_out)

    with gr.Tab("📹 Live Camera"):
        gr.Markdown("Enable your webcam below. Each frame is processed in real-time.")
        with gr.Row():
            cam_in = gr.Image(sources=["webcam"], streaming=True, label="Webcam")
            cam_out = gr.Image(label="Live Analysis")
        live_status = gr.Markdown("Waiting for camera...")
        cam_in.stream(fn=handle_webcam_frame, inputs=cam_in, outputs=[cam_out, live_status])
        gr.Button("🔄 Reset State").click(fn=lambda: (_reset_live(), None, "Reset."),
                                           outputs=[cam_out, live_status])

    gr.Markdown("---\n**Models:** YOLOv8n · BlazePose · Depth Anything · Rule-based risk fusion")

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)