"""
Frame display functionality for the Behavioral Labeler.

This module provides a GUI window to display video frames and associated
labeling information.
"""

import customtkinter as ctk
from PIL import Image
import numpy as np
import cv2

from ..src import utils

import logging
logger = logging.getLogger(__name__)

class FrameDisplayWindow:
    """
    A CustomTkinter Toplevel window for displaying video frames and labeling controls.

    This class creates a two-panel layout: one for the video frame and one for
    displaying real-time information such as frame number, behavior statuses,
    and keyboard shortcuts. It communicates user input back to the main
    application controller via callbacks.
    """

    # Interaction constants
    PAN_STEP = 20
    ZOOM_IN_FACTOR = 1.1
    ZOOM_OUT_FACTOR = 0.9
    MAX_SCALE = 3.0
    MIN_SCALE = 0.1

    def __init__(self, master: ctk.CTk, video_name: str, total_frames: int,
                 behavior_info: dict, operant_keys: dict, fixed_control_keys: dict,
                 on_key_press: callable, on_close: callable):
        """
        Initializes the Frame Display Window.

        Args:
            master (ctk.CTk): The root tkinter object.
            video_name (str): The name of the video file being labeled.
            total_frames (int): The total number of frames in the video.
            behavior_info (dict): A dictionary containing details about each behavior.
            operant_keys (dict): A dictionary of keys for navigation and actions.
            fixed_control_keys (dict): A dictionary of fixed application control keys.
            on_key_press (callable): A callback function to be invoked on a key press event.
                                     It receives the pressed key character as an argument.
            on_close (callable): A callback function to be invoked when the window is closed.
        """
        self.master = master
        self.on_key_press_callback = on_key_press
        self.on_close_callback = on_close
        
        # Generate behavior colors for consistent highlighting
        self.behavior_colors = utils.generate_behavior_colors(list(behavior_info.keys()))
        
        # Frame scaling factor (1.0 = original size)
        self.frame_scale = 1.0
        # Frame panning offsets in pixels (positive x -> right, positive y -> down)
        self.pan_x = 0
        self.pan_y = 0
        
        # Store current frame data for immediate redraw
        self.current_frame_data = None

        # Create the toplevel window
        self.window = ctk.CTkToplevel(master)
        self.window.title(f"Labeling: {video_name}")
        self.window.geometry("1200x700")
        self.window.minsize(1000, 600)

        # Configure the main grid layout
        self.window.grid_columnconfigure(0, weight=1) # Video panel - flexible
        self.window.grid_columnconfigure(1, weight=0) # Info panel - fixed width
        self.window.grid_rowconfigure(0, weight=1)

        # --- Video Frame Panel ---
        self.video_frame_panel = ctk.CTkFrame(self.window, fg_color="black")
        self.video_frame_panel.grid(row=0, column=0, sticky="nsew", padx=5, pady=5)
        self.video_frame_panel.grid_propagate(False)
        
        # Create a container frame with minimal padding for the video label
        self.video_container = ctk.CTkFrame(self.video_frame_panel, fg_color="transparent")
        self.video_container.pack(fill="both", expand=True, padx=2, pady=2)
        
        self.video_label = ctk.CTkLabel(self.video_container, text="", anchor="center")
        # Center the image within the container; pan will be applied via x/y offsets
        self.video_label.place(relx=0.5, rely=0.5, anchor="center")

        # --- Info Panel ---
        self.info_panel = ctk.CTkFrame(self.window)
        self.info_panel.grid(row=0, column=1, sticky="nsew", padx=(0, 10), pady=10)
        self.info_panel.grid_propagate(False)  # Prevent panel from shrinking
        self.info_panel.configure(width=350)  # Set fixed width for info panel
        self._create_info_panel(video_name, total_frames, behavior_info, operant_keys, fixed_control_keys)

        # Bind events
        self.window.bind("<KeyPress>", self._handle_key_press)
        self.window.protocol("WM_DELETE_WINDOW", self.on_close_callback)
        
        # Make window stay on top and grab focus
        self.window.attributes("-topmost", True)
        self.window.lift()
        self.window.focus_force()

        logger.info("FrameDisplayWindow initialized.")

    def _create_info_panel(self, video_name: str, total_frames: int,
                           behavior_info: dict, operant_keys: dict, fixed_control_keys: dict):
        """Creates and populates the widgets for the information panel."""
        self.info_panel.grid_columnconfigure(0, weight=1)
        
        # --- Header ---
        header_frame = ctk.CTkFrame(self.info_panel, fg_color="transparent")
        header_frame.pack(fill="x", padx=15, pady=10)
        ctk.CTkLabel(header_frame, text="RAINSTORM Behavioral Labeler", font=ctk.CTkFont(size=16, weight="bold")).pack()
        ctk.CTkLabel(header_frame, text="https://github.com/sdhers/Rainstorm", font=ctk.CTkFont(size=10)).pack()
        
        # --- Session Info ---
        session_frame = ctk.CTkFrame(self.info_panel)
        session_frame.pack(fill="x", padx=10, pady=10)
        session_frame.grid_columnconfigure(0, weight=1)
        ctk.CTkLabel(session_frame, text="Session Info", font=ctk.CTkFont(weight="bold")).grid(row=0, column=0, pady=(5,10), sticky="w", padx=10)
        
        self.video_name_label = ctk.CTkLabel(session_frame, text=f"Video: {video_name}", wraplength=250, justify="left")
        self.video_name_label.grid(row=1, column=0, sticky="w", padx=10)
        
        self.frame_count_label = ctk.CTkLabel(session_frame, text=f"Frame: 1/{total_frames}")
        self.frame_count_label.grid(row=2, column=0, pady=(0,5), sticky="w", padx=10)
        
        # --- Controls Info ---
        controls_frame = ctk.CTkFrame(self.info_panel)
        controls_frame.pack(fill="x", padx=10, pady=10)
        controls_frame.grid_columnconfigure(0, weight=1)
        ctk.CTkLabel(controls_frame, text="Controls", font=ctk.CTkFont(weight="bold")).grid(row=0, column=0, pady=(5,10), sticky="w", padx=10)
        
        nav_text = f"Next: '{operant_keys['next']}' | Prev: '{operant_keys['prev']}' | FFW: '{operant_keys['ffw']}'"
        ctk.CTkLabel(controls_frame, text=nav_text).grid(row=1, column=0, sticky="w", padx=10)
        
        action_text = f"Erase Label: '{operant_keys['erase']}' | OpenTimeline: '{fixed_control_keys['go_to']}'"
        ctk.CTkLabel(controls_frame, text=action_text).grid(row=2, column=0, sticky="w", padx=10)
        
        resize_text = "Frame Size: '+' zoom in | '-' zoom out"
        ctk.CTkLabel(controls_frame, text=resize_text).grid(row=3, column=0, pady=(0,2), sticky="w", padx=10)
        pan_text = "Arrow keys: pan frame"
        ctk.CTkLabel(controls_frame, text=pan_text).grid(row=4, column=0, pady=(0,5), sticky="w", padx=10)

        # --- Behaviors Panel ---
        ctk.CTkLabel(self.info_panel, text="Behaviors", font=ctk.CTkFont(size=14, weight="bold")).pack(pady=(10,5))
        
        # Create header row with subtle background
        header_frame = ctk.CTkFrame(self.info_panel, fg_color="#1E1E1E", corner_radius=8)
        header_frame.pack(fill="x", padx=10, pady=(0, 5))
        header_frame.grid_columnconfigure(0, weight=1)
        header_frame.grid_columnconfigure(1, weight=0)
        header_frame.grid_columnconfigure(2, weight=0)
        
        ctk.CTkLabel(header_frame, text="Behavior", font=ctk.CTkFont(weight="bold"), text_color="#FFFFFF").grid(row=0, column=0, sticky="w", padx=(8,0), pady=4)
        ctk.CTkLabel(header_frame, text="Key", font=ctk.CTkFont(weight="bold"), width=40, text_color="#FFFFFF").grid(row=0, column=1, pady=4)
        ctk.CTkLabel(header_frame, text="Events", font=ctk.CTkFont(weight="bold"), width=60, text_color="#FFFFFF").grid(row=0, column=2, padx=(0,8), pady=4)
        
        behaviors_scroll_frame = ctk.CTkScrollableFrame(self.info_panel)
        behaviors_scroll_frame.pack(fill="both", expand=True, padx=10, pady=(0, 10))
        behaviors_scroll_frame.grid_columnconfigure(0, weight=1)

        self.behavior_labels = {}
        for behavior_name, info in behavior_info.items():
            # Create individual cell frames for each column
            row_frame = ctk.CTkFrame(behaviors_scroll_frame, fg_color="transparent")
            row_frame.pack(fill="x", pady=1)
            row_frame.grid_columnconfigure(0, weight=1)
            row_frame.grid_columnconfigure(1, weight=0)
            row_frame.grid_columnconfigure(2, weight=0)
            
            # Behavior name cell
            name_cell = ctk.CTkFrame(row_frame, fg_color="#2B2B2B", corner_radius=8)
            name_cell.grid(row=0, column=0, sticky="ew", padx=(0, 2))
            name_label = ctk.CTkLabel(name_cell, text=f"{behavior_name}", anchor="w", font=ctk.CTkFont(size=12))
            name_label.pack(fill="x", padx=8, pady=4)
            
            # Key cell
            key_cell = ctk.CTkFrame(row_frame, fg_color="#2B2B2B", corner_radius=8)
            key_cell.grid(row=0, column=1, padx=1)
            key_label = ctk.CTkLabel(key_cell, text=f"{info['key']}", width=25, font=ctk.CTkFont(size=12))
            key_label.pack(padx=8, pady=4)

            # Events cell
            events_cell = ctk.CTkFrame(row_frame, fg_color="#2B2B2B", corner_radius=8)
            events_cell.grid(row=0, column=2, padx=(2, 0))
            sum_label = ctk.CTkLabel(events_cell, text=f"{info['sum']}", width=30, font=ctk.CTkFont(size=12))
            sum_label.pack(padx=8, pady=4)

            self.behavior_labels[behavior_name] = {'name': name_label, 'sum': sum_label, 'cell': name_cell}

    def _handle_key_press(self, event):
        """
        Internal handler for key press events. It calls the registered callback.
        """
        # Handle panning with arrow keys (maintain centering baseline and add offsets)
        if event.keysym in ("Left", "Right", "Up", "Down"):
            step = self.PAN_STEP
            if event.keysym == "Left":
                self.pan_x -= step
            elif event.keysym == "Right":
                self.pan_x += step
            elif event.keysym == "Up":
                self.pan_y -= step
            elif event.keysym == "Down":
                self.pan_y += step
            self._reposition_image()
            return
        # Handle frame resizing with - and + keys
        if event.char == '-' or event.char == '_':  # Handle both - and shift+-
            self._decrease_frame_size()
            return
        elif event.char == '+' or event.char == '=':  # Handle both + and shift+=
            self._increase_frame_size()
            return
        
        # We ignore modifier keys and only pass character keys
        if event.char:
            logger.debug(f"Key pressed: '{event.char}'")
            self.on_key_press_callback(event.char)

    def _increase_frame_size(self):
        """
        Increase the frame size by 10%.
        """
        self.frame_scale = min(self.frame_scale * self.ZOOM_IN_FACTOR, self.MAX_SCALE)
        logger.info(f"Frame scale increased to: {self.frame_scale:.2f}")
        # Clear cached frame size to force recalculation
        if hasattr(self, '_frame_size'):
            delattr(self, '_frame_size')
        # Immediately redraw current frame if available
        self._redraw_current_frame()

    def _decrease_frame_size(self):
        """
        Decrease the frame size by 10%.
        """
        self.frame_scale = max(self.frame_scale * self.ZOOM_OUT_FACTOR, self.MIN_SCALE)
        logger.info(f"Frame scale decreased to: {self.frame_scale:.2f}")
        # Clear cached frame size to force recalculation
        if hasattr(self, '_frame_size'):
            delattr(self, '_frame_size')
        # Immediately redraw current frame if available
        self._redraw_current_frame()

    def _redraw_current_frame(self):
        """
        Redraw the current frame with the new scale factor.
        """
        if self.current_frame_data is not None:
            # Force recalculation of frame size on redraw
            if hasattr(self, '_frame_size'):
                delattr(self, '_frame_size')
            self._apply_image_with_calculated_size(self.current_frame_data, recalc_size=True)

    def _reposition_image(self):
        """Apply pan offsets relative to centered anchor."""
        if not hasattr(self, '_frame_size'):
            return
        # With anchor at the center (relx=0.5, rely=0.5), just offset by pan values
        self.video_label.place_configure(x=self.pan_x, y=self.pan_y)

    def update_display(self, frame: np.ndarray, frame_number: int, total_frames: int, behavior_info: dict):
        """
        Updates the window with a new frame and new information.

        This method is called by the main controller to refresh the UI.

        Args:
            frame (np.ndarray): The new video frame to display (in BGR format from OpenCV).
            frame_number (int): The current frame number (0-indexed).
            total_frames (int): The total number of frames in the video.
            behavior_info (dict): The updated dictionary with behavior sums and current statuses.
        """
        # --- Update Video Frame ---
        if frame is not None:
            # Convert BGR (OpenCV) to RGB (PIL)
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            pil_image = Image.fromarray(rgb_frame)
            
            # Store current frame data for immediate redraw on resize
            self.current_frame_data = pil_image
            
            # Apply image using helper (calculates size on first use)
            self._apply_image_with_calculated_size(pil_image, recalc_size=not hasattr(self, '_frame_size'))
            
        # --- Update Info Panel ---
        self.frame_count_label.configure(text=f"Frame: {frame_number + 1}/{total_frames}")

        for behavior_name, info in behavior_info.items():
            if behavior_name in self.behavior_labels:
                labels = self.behavior_labels[behavior_name]
                
                # Update sum
                labels['sum'].configure(text=f"{info['sum']}")
                
                # Update cell background color based on current status
                is_active = info.get('current_behavior', 0) == 1
                if is_active:
                    # Highlight the entire cell with the behavior's assigned color
                    behavior_color = self.behavior_colors.get(behavior_name, "#FF6B6B")
                    labels['cell'].configure(fg_color=behavior_color)
                    labels['name'].configure(
                        text_color="#FFFFFF", # White text on colored background
                        font=ctk.CTkFont(size=12, weight="bold")
                    )
                else:
                    # Return to default gray background
                    labels['cell'].configure(fg_color="#2B2B2B")
                    labels['name'].configure(
                        font=ctk.CTkFont(size=12, weight="normal")
                    )
        logger.debug(f"Display updated for frame {frame_number + 1}.")

    def _apply_image_with_calculated_size(self, pil_image: Image.Image, recalc_size: bool):
        """Calculate appropriate size for current panel and apply image.

        If recalc_size is True, compute and store a new self._frame_size based on the
        container size, image aspect, and current frame_scale. Otherwise, reuse the stored size.
        """
        # Ensure geometry is updated
        self.window.update_idletasks()
        if recalc_size or not hasattr(self, '_frame_size'):
            # Get panel dimensions
            panel_width = self.video_container.winfo_width()
            panel_height = self.video_container.winfo_height()

            # Fallback if panel is too small
            if panel_width <= 1 or panel_height <= 1:
                panel_width = 800
                panel_height = 600

            # Calculate base size to fit panel while maintaining aspect ratio
            img_aspect = pil_image.width / pil_image.height
            panel_aspect = panel_width / panel_height

            if img_aspect > panel_aspect:
                base_width = panel_width
                base_height = int(base_width / img_aspect)
            else:
                base_height = panel_height
                base_width = int(base_height * img_aspect)

            # Apply scaling factor
            new_width = int(base_width * self.frame_scale)
            new_height = int(base_height * self.frame_scale)

            # Store the calculated size - recalculate when scale changes
            self._frame_size = (new_width, new_height)
            logger.info(f"Frame size calculated: {new_width}x{new_height} (scale: {self.frame_scale:.2f})")

        # Always use the stored size
        ctk_image = ctk.CTkImage(light_image=pil_image, size=self._frame_size)
        self.video_label.configure(image=ctk_image)
        # Reposition based on centering + pan offsets
        self._reposition_image()

    def close(self):
        """Destroys the window."""
        if self.window.winfo_exists():
            self.window.destroy()
        logger.info("FrameDisplayWindow closed.")
