import logging
import flet as ft
from js_execution import set_camera_frame_callback
from flet_extensions import eval_js
import asyncio
from typing import List, Dict

logger = logging.getLogger(__name__)

class Frame:
    def __init__(self, data: str, timestamp: float):
        self.data = data
        self.timestamp = timestamp

class CameraImage(ft.Container):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.stream_active = False
        self.video_image1 = ft.Image(
            src="",
            width=self.width,
            height=self.height,
            fit=ft.ImageFit.CONTAIN,
            visible=False,
        )
        self.video_image2 = ft.Image(
            src="",
            width=self.width,
            height=self.height,
            fit=ft.ImageFit.CONTAIN,
            visible=False,
        )
        self.current_video_image = self.video_image1
        self.video_stack = ft.Stack([self.video_image1, self.video_image2])
        self.content = self.video_stack
        set_camera_frame_callback(self.handle_camera_frame)
        self.captured_frames: Dict[int, Frame] = {}
        self.stream_start_time = None

    async def start_camera_stream_async(self) -> str:
        try:
            result = await eval_js(self.page, """
            window.cameraStreamActive = false;
            window.cameraWebSocket = null;

            async function setupCameraStream() {
                console.log("setupCameraStream called");
                try {
                    const stream = await navigator.mediaDevices.getUserMedia({ video: true });
                    console.log("Camera access granted");
                    const video = document.createElement('video');
                    video.srcObject = stream;
                    await video.play();

                    const canvas = document.createElement('canvas');
                    canvas.width = video.videoWidth / 2;
                    canvas.height = video.videoHeight / 2;
                    const ctx = canvas.getContext('2d');

                    window.cameraWebSocket = new WebSocket(`ws://${window.location.host}/ws/camera`);
                    
                    function sendFrame() {
                        if (window.cameraWebSocket && window.cameraWebSocket.readyState === WebSocket.OPEN) {
                            ctx.drawImage(video, 0, 0, video.videoWidth, video.videoHeight, 0, 0, canvas.width, canvas.height);
                            window.cameraWebSocket.send(canvas.toDataURL('image/jpeg', 0.2)); // Added quality parameter
                            window.cameraStreamActive && setTimeout(sendFrame, 1000 / 60);
                        }
                    }

                    window.cameraWebSocket.onopen = () => {
                        console.log("WebSocket connection opened");
                        window.cameraStreamActive = true;
                        sendFrame();
                    };

                    window.cameraWebSocket.onclose = () => {
                        console.log("WebSocket connection closed");
                        window.cameraStreamActive = false;
                        stream.getTracks().forEach(track => track.stop());
                    };
                    
                    return "Camera stream setup complete";
                } catch (error) {
                    console.error("Error accessing camera:", error);
                    throw error;
                }
            }

            return await setupCameraStream();
            """)
            logger.info(f"start_camera_stream_async result: {result}")
            self.stream_active = True
            self.stream_start_time = asyncio.get_event_loop().time()
            self.captured_frames.clear()
            return result
        except Exception as e:
            logger.error(f"Error in start_camera_stream_async: {str(e)}")
            return f"Error: {str(e)}"

    async def stop_camera_stream_async(self) -> str:
        try:
            result = await eval_js(self.page, """
            if (window.cameraWebSocket) {
                window.cameraStreamActive = false;
                window.cameraWebSocket.close();
                window.cameraWebSocket = null;
            }
            return 'Camera stream stopped';
            """)
            logger.info(f"stop_camera_stream_async result: {result}")
            self.stream_active = False
            self.stream_start_time = None
            return result
        except Exception as e:
            logger.error(f"Error in stop_camera_stream_async: {str(e)}")
            return f"Error: {str(e)}"

    async def handle_camera_frame(self, frame_data: str):
        if self.stream_active:
            await self.update_webcam_image(frame_data)
            self.capture_timed_frames(frame_data)

    def capture_timed_frames(self, frame_data: str):
        current_time = asyncio.get_event_loop().time()
        if self.stream_start_time is None:
            self.stream_start_time = current_time

        elapsed_time = current_time - self.stream_start_time

        for seconds in [1, 2, 3]:
            if elapsed_time >= seconds and seconds not in self.captured_frames:
                self.captured_frames[seconds] = Frame(frame_data, current_time)

    def get_captured_frames(self) -> List[Frame]:
        return list(self.captured_frames.values())

    async def update_webcam_image(self, frame_data: str) -> None:
        if not frame_data:
            print("Empty data!")
            return
        
        # Frame data is already in Data URI format
        new_src = frame_data 

        self.video_image1.visible = True
        self.video_image2.visible = True

        # Determine which image to update
        if self.current_video_image == self.video_image1:
            image_to_update = self.video_image2
            image_to_hide = self.video_image1
        else:
            image_to_update = self.video_image1
            image_to_hide = self.video_image2

        # Update the hidden image
        image_to_update.src = new_src
        image_to_update.update()
        
        # Wait for the browser to load the new image
        await asyncio.sleep(0.02)  # Adjust this delay if needed
        
        # Update opacity
        image_to_update.opacity = 1
        image_to_hide.opacity = 0
        self.current_video_image = image_to_update

        # Final update to apply opacity changes
        self.video_stack.update()