import os
import cv2
import numpy as np
from datetime import datetime

# ── 🛠️ 0. THE "DUCT TAPE" MONKEYPATCHES ──────────────────────────────────────
# Python 3.13 + Gradio 4.44.0 + Hugging Face Spaces requires these overrides.

# Patch A: Hugging Face Hub (Fixes: ImportError: cannot import name 'HfFolder')
import huggingface_hub
if not hasattr(huggingface_hub, "HfFolder"):
    class MockHfFolder:
        @staticmethod
        def get_token():
            return os.environ.get("HF_TOKEN")
    huggingface_hub.HfFolder = MockHfFolder

# Patch B: Gradio Client Schema Parser (Fixes: TypeError: argument of type 'bool' is not iterable)
try:
    import gradio_client.utils
    
    if hasattr(gradio_client.utils, "get_type"):
        old_get_type = gradio_client.utils.get_type
        def patched_get_type(schema):
            if isinstance(schema, bool):
                return "Any"
            return old_get_type(schema)
        gradio_client.utils.get_type = patched_get_type

    if hasattr(gradio_client.utils, "_json_schema_to_python_type"):
        old_internal_parser = gradio_client.utils._json_schema_to_python_type
        def patched_internal_parser(schema, defs=None):
            if isinstance(schema, bool):
                return "Any"
            return old_internal_parser(schema, defs)
        gradio_client.utils._json_schema_to_python_type = patched_internal_parser
except Exception:
    pass
# ─────────────────────────────────────────────────────────────────────────────

import gradio as gr
import modal

# ── Logging Setup ───────────────────────────────────────────────────────────
init_logs = []

def log_system_event(message: str):
    """Formats logs with a timestamp and syncs to stdout and UI."""
    timestamp = datetime.now().strftime("%H:%M:%S")
    formatted_log = f"[{timestamp}] {message}"
    print(formatted_log)
    init_logs.append(formatted_log)

# ── Modal Backend Connection ────────────────────────────────────────────────
log_system_event("Initializing connection to Modal remote infrastructure...")

try:
    VoxelModelCls = modal.Cls.from_name("flux-klein-voxel-backend", "VoxelModel")
    voxel_backend = VoxelModelCls().process_frame
    log_system_event("✅ Success: Bound directly to Modal Class deployment ('VoxelModel').")
except Exception as e:
    log_system_event(f"⚠️ Modal Class lookup failed ({e}). Attempting fallback function router...")
    try:
        voxel_backend = modal.Function.from_name("flux-klein-voxel-backend", "demo_stream_frame")
        log_system_event("✅ Success: Hooked into fallback standalone Modal Function.")
    except Exception as ex:
        log_system_event(f"❌ Critical: All remote Modal endpoints are unreachable. Error: {ex}")
        voxel_backend = None

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


# ── 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 ONLY raw image bytes to the Modal worker."""
    success, encoded = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 85])
    if not success:
        return frame
    try:
        # Pushing ONLY the camera feed byte data across the network
        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 ─────────────────────────────────────────────────────
frame_counter = 0

def process_video_stream(frame: np.ndarray, mode: str, is_running: bool) -> np.ndarray:
    """
    Accepts incoming frame from the webcam, pipeline settings, and execution state.
    """
    global frame_counter
    if frame is None:
        return None

    # If user hasn't toggled "Start Processing", safely pass raw video back as preview
    if not is_running:
        return frame

    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

    # Runtime console indicator: logs transmission health to terminal every 15 frames
    frame_counter += 1
    if frame_counter % 15 == 0:
        print(f"🚀 [LIVE PIPELINE] Transmitting frames. Dispatched {frame_counter} payloads to Modal.")

    # Process via the 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:
    # State tracking engine variable
    is_running = gr.State(value=False)
    
    gr.Markdown("# ⛏️ Minecraft Spatial Voxel Filter")
    
    with gr.Row():
        # Configuration Settings & Log Panel (Left Column)
        with gr.Column(scale=1):
            gr.Markdown(
                f"### ⚡ Backend Connection Status\n"
                f"Status: {status_color} **{status_text}**"
            )
            
            # Live Diagnostic Log Display View
            ui_logs = gr.Textbox(
                value="\n".join(init_logs),
                label="💻 System Initialization Logs",
                lines=4,
                max_lines=5,
                interactive=False,
            )
            
            mode_dropdown = gr.Dropdown(
                choices=["Minecraft Filter", "Streaming Demo"],
                value="Minecraft Filter",
                label="🎯 Pipeline Mode",
                interactive=True,
            )
            
            with gr.Row():
                start_btn = gr.Button("🚀 Start Processing", variant="primary")
                stop_btn = gr.Button("🛑 Stop", variant="secondary")
            
        # Video Capture Viewports (Right Column)
        with gr.Column(scale=2):
            input_stream = gr.Image(sources=["webcam"], streaming=True, label="Live Webcam Input")
            output_stream = gr.Image(interactive=False, label="Voxel Output Viewport")

    # Wire up button interface trigger mappings to flip our State flag
    start_btn.click(fn=lambda: True, inputs=None, outputs=is_running)
    stop_btn.click(fn=lambda: False, inputs=None, outputs=is_running)

    # Core engine transmission loop linked up with ONLY the stream, mode, and run-state
    input_stream.stream(
        fn=process_video_stream,
        inputs=[input_stream, mode_dropdown, is_running],
        outputs=[output_stream],
        trigger_mode="always_last", # Drops intermediate backlog frames when backend is busy
        concurrency_limit=1         # Ensures only one frame flies over the network at a time
    )

if __name__ == "__main__":
    # Bound to 0.0.0.0 to bypass Hugging Face Spaces localhost proxy issues
    demo.launch(server_name="0.0.0.0", server_port=7860)