import gradio as gr
import numpy as np
import spaces
from fastrtc import WebRTC, get_turn_credentials

# This is the core processing function that will be assigned to the GPU.
# It receives a generator (`frame_stream`) that yields frames from the webcam
# and it yields processed frames back to the output component.
@spaces.GPU
def process_frames_on_gpu(frame_stream):
    """
    This function runs as a persistent process on the GPU.
    It iterates over incoming frames, processes them, and yields the results.
    """
    print("🚀 GPU Frame processing loop started.")
    
    if frame_stream is None:
        print("Input stream is None, ending.")
        return

    # This loop will block until a new frame is available, making it reactive.
    for frame in frame_stream:
        if frame is not None:
            # This is where your GPU-intensive work would happen.
            flipped_frame = np.flip(frame, axis=(0, 1))
            
            # Yield the processed frame to the output stream.
            yield flipped_frame
            
    print("🛑 GPU Frame processing loop finished.")

# --- Gradio UI Layout ---
with gr.Blocks(theme=gr.themes.Soft(), title="FastRTC ZeroGPU Flipper") as demo:
    gr.Markdown("# 🚀 FastRTC Webcam Flipper for ZeroGPU")
    gr.Markdown(
        "*This version uses a separate button to trigger the stream, correctly using Gradio's streaming iterator pattern for ZeroGPU without conflicting with the `fastrtc` component's internal logic.*"
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. Your Webcam Feed (Input)")
            webcam_input = WebRTC(
                label="Webcam Input",
                modality="video", 
                mode="send",
                width=640,
                height=480,
                rtc_configuration=get_turn_credentials(),
            )
            start_button = gr.Button("🚀 Start Processing Stream", variant="primary")
            
        with gr.Column(scale=1):
            gr.Markdown("### 2. Flipped Video (Output)")
            video_output = WebRTC(
                label="Flipped Output Stream",
                modality="video", 
                mode="receive",
                width=640,
                height=480,
                rtc_configuration=get_turn_credentials(),
            )
        
    # By using a button's `click` event to start the stream, we use Gradio's
    # standard streaming API, which correctly provides the iterator to our
    # decorated function without causing a conflict.
    start_button.click(
        fn=process_frames_on_gpu, 
        inputs=[webcam_input],
        outputs=[video_output]
    )

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