import os
import cv2
import numpy as np
import gradio as gr
import modal
from fastrtc import WebRTC, get_cloudflare_turn_credentials

# ── Modal Backend ───────────────────────────────────────────────────────────
try:
    VoxelModelCls = modal.Cls.from_name("flux-klein-voxel-backend", "VoxelModel")
    voxel_backend = VoxelModelCls().process_frame
    print("✅ Modal Cls connected.")
except Exception as e:
    try:
        voxel_backend = modal.Function.from_name("flux-klein-voxel-backend", "demo_stream_frame")
        print("✅ Modal Function connected.")
    except Exception as ex:
        print(f"❌ Modal backend offline: {ex}")
        voxel_backend = None

status_color = "🟢" if voxel_backend is not None else "🔴"
status_text  = "Connected" if voxel_backend is not None else "Offline"


# ── WebRTC Configuration Fallback ───────────────────────────────────────────
def get_safe_rtc_configuration():
    """
    Prevents startup crashes if Cloudflare or HF tokens are missing from the environment.
    Falls back gracefully to standard public Google STUN routing if name resolution drops out.
    """
    has_hf = bool(os.environ.get("HF_TOKEN"))
    has_cf = bool(os.environ.get("CLOUDFLARE_TURN_KEY_ID") and os.environ.get("CLOUDFLARE_TURN_KEY_API_TOKEN"))
    
    if has_hf or has_cf:
        try:
            return get_cloudflare_turn_credentials()
        except Exception as e:
            print(f"⚠️ Failed to fetch Cloudflare TURN credentials: {e}")
            
    print("ℹ️ Missing credentials for Cloudflare TURN. Falling back to public Google STUN server.")
    return {"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}


# ── Helpers ─────────────────────────────────────────────────────────────────
def _offline_frame(frame: np.ndarray, message: str) -> np.ndarray:
    """Generates a styled placeholder frame when backend is offline."""
    h, w = (frame.shape[:2] if frame is not None else (480, 640))
    out = np.zeros((h, w, 3), dtype=np.uint8)
    cv2.putText(out, message, (max(10, w // 8), h // 2),
                cv2.FONT_HERSHEY_DUPLEX, 0.8, (0, 0, 220), 2)
    return out


def _run_voxel_backend(frame: np.ndarray) -> np.ndarray:
    """Encodes and ships raw image bytes to the Modal worker without extra parameters."""
    success, encoded = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 85])
    if not success:
        return frame
    try:
        processed_bytes = voxel_backend.remote(encoded.tobytes())
        result = cv2.imdecode(np.frombuffer(processed_bytes, dtype=np.uint8), cv2.IMREAD_COLOR)
        return result if result is not None else frame
    except Exception as err:
        out = frame.copy()
        cv2.putText(out, f"Modal error: {str(err)[:45]}", (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 255), 1)
        return out


# ── Core Stream Handler ─────────────────────────────────────────────────────
def process_video_stream(frame: np.ndarray, mode: str) -> np.ndarray:
    """
    Unified real-time handler mapping directly to FastRTC's stream pipeline.
    Accepts incoming frame from the WebRTC component and the mode dropdown selection.
    """
    if frame is None:
        return None

    if voxel_backend is None:
        err_frame = _offline_frame(frame, "Modal Backend Offline")
        if mode == "Streaming Demo":
            return np.hstack([frame, err_frame])
        return err_frame

    # Process via the lightweight single-image pipeline
    processed = _run_voxel_backend(frame)

    # Mode A: Full view rendering
    if mode == "Minecraft Filter":
        return processed

    # Mode B: Side-by-side split rendering
    elif mode == "Streaming Demo":
        if processed.shape[0] != frame.shape[0]:
            processed = cv2.resize(processed, (frame.shape[1], frame.shape[0]))
            
        raw_labeled = frame.copy()
        cv2.putText(raw_labeled, "RAW", (10, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2)
        cv2.putText(processed, "MINECRAFT", (10, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2)
        return np.hstack([raw_labeled, processed])

    return processed


# ── Custom Gradio Interface Layout ──────────────────────────────────────────
with gr.Blocks(title="⛏️ Minecraft Spatial Voxel Filter") as demo:
    gr.Markdown("# ⛏️ Minecraft Spatial Voxel Filter")
    
    with gr.Row():
        # Configuration Settings & Status (Left Side Box)
        with gr.Column(scale=1):
            status_md = gr.Markdown(
                f"### ⚡ Backend Connection Status\n"
                f"Status: {status_color} **{status_text}**\n\n"
                f"Switching between pipeline modes instantly updates your active viewport feed."
            )
            
            mode_dropdown = gr.Dropdown(
                choices=["Minecraft Filter", "Streaming Demo"],
                value="Minecraft Filter",
                label="🎯 Pipeline Mode",
                interactive=True,
            )
            
        # Contained Video Box Viewport (Right Side Box)
        with gr.Column(scale=2):
            webrtc_stream = WebRTC(
                label="Live Filter Stream",
                modality="video",
                mode="send-receive",
                rtc_configuration=get_safe_rtc_configuration(),
            )

    # Explicitly attach the handler to the custom WebRTC component
    webrtc_stream.stream(
        fn=process_video_stream,
        inputs=[webrtc_stream, mode_dropdown],
        outputs=[webrtc_stream],
        time_limit=150,
        concurrency_limit=4,
    )

if __name__ == "__main__":
    demo.launch()