import cv2
import numpy as np
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager
from selenium.webdriver.chrome.options import Options
import time
from flask import Response
import cv2
import time

def start_html_stream(output_frame, lock):
    def generate():
        while True:
            with lock:
                if output_frame[0] is None:
                    # TODO: Investigate why if remove this print or time.sleep (but not both at the same time), the stream does not work
                    # time.sleep(0.5)
                    print(f"{time.time()} - Frame is None.")
                    continue
                (flag, encoded_image) = cv2.imencode(".jpg", output_frame[0])
                if not flag:
                    continue
            yield(b'--frame\r\n' b'Content-Type: image/jpeg\r\n\r\n' + bytearray(encoded_image) + b'\r\n')
    return Response(generate(), mimetype="multipart/x-mixed-replace; boundary=frame")

def capture_webpage(driver):
    screenshot = driver.get_screenshot_as_png()
    image = np.frombuffer(screenshot, dtype=np.uint8)
    image = cv2.imdecode(image, cv2.IMREAD_COLOR)
    
    height, width, _ = image.shape
    new_width = 1920
    new_height = int(new_width * height / width)
    
    image = cv2.resize(image, (new_width, new_height))
    
    if new_height > 1080:
        start_y = (new_height - 1080) // 2
        image = image[start_y:start_y + 1080, :]
    
    return image

def get_webpage_frames(stream_url):
    chrome_options = Options()
    chrome_options.add_argument("--headless")
    chrome_options.add_argument("--window-size=1920x1080")
    
    driver = webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=chrome_options)
    driver.get(stream_url)
    
    frames_per_second = 10
    frame_interval = 1 / frames_per_second

    try:
        while True:
            image = capture_webpage(driver)
            yield image
            time.sleep(frame_interval)
    finally:
        driver.quit()

def get_mjpeg_frames(stream_url):
    cap = cv2.VideoCapture(stream_url)
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        yield frame
    
    cap.release()