import os
import cv2
import numpy as np
import time 
import gradio as gr
import tempfile
ENVIRONMENT = "pi"
try:
  from sensor import Sensor
except ImportError:
  Sensor = lambda : None
  ENVIRONMENT = "dev"

from agent import Agent


sensor = Sensor()
pano_agent = Agent(sensor = sensor)

if ENVIRONMENT == "dev":
  def work_once(self, state=[]):
    if len(state) < 2:
      return np.random.normal(loc=128, scale=30, size=(300,400,3)).clip(0,255).astype(np.uint8)
    prev_image = state[-2]
    current_image = state[-1]
    coco_bboxes = self.predict(current_image, prev_image)
    artifacts = self.handle_prediction(current_image, prev_image, coco_bboxes,0)
    return artifacts
  pano_agent.work_once = work_once.__get__(pano_agent)
  with gr.Blocks() as demo:
    with gr.Row():
      threshold_slider = gr.Slider(
          minimum=0.0, 
          maximum=1.0, 
          value=pano_agent.bird_detector.threshold, 
          step=0.01, 
          label="Threshold"
      )
      def change_detection_threshold(threshold):
          pano_agent.bird_detector.threshold = threshold
          
      threshold_slider.change(change_detection_threshold,
                              inputs=threshold_slider,
                              outputs=None)
      percentile_slider = gr.Slider(
          minimum=0, 
          maximum=100, 
          value=pano_agent.bird_detector.percentile, 
          step=1, 
          label="Percentile"
      )
      def change_detection_percentile(percentile):
          pano_agent.bird_detector.percentile = percentile
      percentile_slider.change(change_detection_percentile,
                              inputs=percentile_slider,
                              outputs=None)
    with gr.Row():
      with gr.Tab("Upload Video"):
        gr.Markdown("### Video Bird Detector (Demo Version)")
        with gr.Row():
          with gr.Column():
            video_input = gr.Video(label="Upload Video")
          with gr.Column():
            gr.Examples(
              examples=[["nyagatare_2025_11_24_video.mp4"]],
              inputs=video_input,
            )
            process_button = gr.Button("Process Video")

        with gr.Row():
          video_output = gr.Video(label="Processed Output")

        def process_video_dummy(video):
            """
            Processes the input video by extracting frames at 2fps, running pano_agent.work_once
            on each pair of consecutive frames, and returning the processed video.
            """

            # video is generally a path in Gradio, but can sometimes be a numpy array.
            # Assume it's a file path.
            cap = cv2.VideoCapture(video)
            if not cap.isOpened():
                raise RuntimeError(f"Cannot open video {video}")

            # Get video properties
            input_fps = cap.get(cv2.CAP_PROP_FPS)
            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

            if frame_count < 0:
              input_fps = 20.0

            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

            new_fps = 2.0
            frame_interval = int(round(input_fps / new_fps)) if input_fps > new_fps else 1

            frames = []
            success, frame = cap.read()
            idx = 0
            while success:
                if idx % frame_interval == 0:
                    frames.append(frame)
                success, frame = cap.read()
                idx += 1
            cap.release()

            if len(frames) < 2:
                # Not enough frames, just return the original video
                return video
            output_frames = []
            prev_frame = frames[0]
            for current_frame in frames[1:]:
                out_frame = pano_agent.work_once([prev_frame, current_frame])
                output_frames.append(out_frame if out_frame is not None else current_frame)
                prev_frame = current_frame

            # Write result to a temporary file
            temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
            temp_path = temp_file.name
            temp_file.close()

            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            writer = cv2.VideoWriter(temp_path, fourcc, new_fps, (output_frames[0].shape[1], output_frames[0].shape[0]))
            for f in output_frames:
                # If output frame is grayscale, convert to BGR
                if len(f.shape) == 2:
                    f = cv2.cvtColor(f, cv2.COLOR_GRAY2BGR)
                writer.write(f)
            writer.release()
            return temp_path
        
        process_button.click(
            process_video_dummy,
            inputs=video_input,
            outputs=video_output,
        )

      with gr.Tab("Real-Time (Webcam Only)"):
        with gr.Row():
          # start_agent_btn = gr.Button("Submit")
          # stop_button = gr.Button("Stop Process", variant="stop")
          text_info = gr.Markdown("### Welcome to the Bird Detector Demo!")
        with gr.Row():
          input_img = gr.Image(label="webcam input", sources=["webcam"], type="numpy")
          dummy_img = gr.Image(visible=False)
          img_state = gr.State([])
          agent_view = gr.Image(label="Detection output", type="numpy")
          input_img.stream(
            lambda img, state: ((state[-1:] + [img]), 
                                pano_agent.work_once(state[-1:] + [img])),
            inputs=[input_img, img_state],
            outputs=[img_state, agent_view],
          )
      
if ENVIRONMENT == "pi":
  with gr.Blocks() as demo:
    with gr.Row():
      start_agent_btn = gr.Button("Submit")
      stop_button = gr.Button("Stop Process", variant="stop")
    with gr.Row():
      agent_view = gr.Image()
      agent_work = start_agent_btn.click(pano_agent.work,
                            inputs = None,
                            outputs = agent_view
                            )
      stop_button.click(
        fn=None,  # No function to execute on stop, just cancel
        inputs=None,
        outputs=None,
        cancels=[agent_work]
      )

demo.launch(
    # debug = True,
    server_name="0.0.0.0",
)