import gi
import threading
import numpy as np
import socket
import cv2

gi.require_version('Gst', '1.0')
gi.require_version('GstRtspServer', '1.0')
from gi.repository import Gst, GstRtspServer, GLib

def get_local_ip():
    """Retrieve the local IP address of the machine."""
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception as e:
        return str(e)

class VideoStream(GstRtspServer.RTSPMediaFactory):
    def __init__(self, fps, width, height):
        super(VideoStream, self).__init__()
        self.fps = fps
        self.width = width
        self.height = height
        self.frame = np.zeros((height, width, 3), dtype=np.uint8)  # Initial black frame
        self.lock = threading.Lock()

    def update_frame(self, frame):
        """Externally updates the current frame."""
        with self.lock:
            self.frame = cv2.resize(frame, (self.width, self.height))

    def on_need_data(self, src, length):
        """Provides frames to the pipeline when requested."""
        with self.lock:
            frame_rgb = cv2.cvtColor(self.frame, cv2.COLOR_BGR2RGB)
            data = frame_rgb.tobytes()

        buf = Gst.Buffer.new_allocate(None, len(data), None)
        buf.fill(0, data)
        buf.duration = Gst.SECOND // self.fps
        buf.pts = buf.dts = Gst.CLOCK_TIME_NONE
        src.emit("push-buffer", buf)

    def do_create_element(self, url):
        """Creates the GStreamer pipeline for RTSP streaming."""
        pipeline_str = (
            "appsrc name=source is-live=true format=GST_FORMAT_TIME "
            "caps=video/x-raw,format=RGB,width={},height={},framerate={}/1 "
            "! videoconvert ! video/x-raw,format=I420 "
            "! x264enc tune=zerolatency bitrate=500 speed-preset=ultrafast "
            "! h264parse ! rtph264pay config-interval=1 name=pay0 pt=96"
        ).format(self.width, self.height, self.fps)

        pipeline = Gst.parse_launch(pipeline_str)
        src = pipeline.get_by_name("source")
        src.connect("need-data", self.on_need_data)
        return pipeline

class RTSPServer:
    def __init__(self, ip="0.0.0.0", port=8554, mount_point="/stream"):
        Gst.init(None)
        self.server = GstRtspServer.RTSPServer()
        self.server.set_address(ip)
        self.server.set_service(str(port))
        self.mount_point = mount_point

        self.video_stream = None
        self.current_fps = None
        self.current_width = None
        self.current_height = None
        self.lock = threading.Lock()

        self.server.attach(None)

    def create_stream(self, fps, width, height):
        """Create a new RTSP stream with the given resolution and FPS."""
        with self.lock:
            # Check if we need to recreate the stream
            if (self.video_stream and
                self.current_fps == fps and
                self.current_width == width and
                self.current_height == height):
                return  # No need to recreate

            # Remove old stream if exists
            if self.video_stream:
                print("Recreating RTSP stream for new resolution/fps...")
                self.server.get_mount_points().remove_factory(self.mount_point)

            # Create a new stream with the updated settings
            self.video_stream = VideoStream(fps, width, height)
            self.server.get_mount_points().add_factory(self.mount_point, self.video_stream)

            self.current_fps = fps
            self.current_width = width
            self.current_height = height

            print(f"New RTSP stream created: {width}x{height} at {fps}fps")

    def update_frame(self, frame):
        """Push frames to the current streaming channel."""
        if frame is None or frame.size == 0:
            return

        height, width, _ = frame.shape
        fps = 30  # Default FPS

        # Ensure we have a valid stream for the given resolution
        self.create_stream(fps, width, height)

        # Update the current stream with the new frame
        if self.video_stream:
            self.video_stream.update_frame(frame)

# Global RTSP server instance
# rtsp_server = RTSPServer(get_local_ip(), 41231)

def run_server(server):
    """Start the RTSP server loop."""
    print(f"RTSP Server running at rtsp://{server.server.get_address()}:{server.server.get_service()}/stream")
    loop = GLib.MainLoop()
    loop.run()

# def stream_webcam():
#     cap = cv2.VideoCapture("/home/mht/Downloads/bcd2890d71caaf0e095b95c9b525973f61186656-360p.mp4")  # Open webcam
#     while cap.isOpened():
#         ret, frame = cap.read()
#         if ret:
#             rtsp_server.update_frame(frame)  # Send frame to RTSP server

# if __name__ == "__main__":
#     # Start RTSP server in a separate thread
#     threading.Thread(target=run_server, daemon=True).start()
#
#     # Stream webcam frames
#     stream_webcam()