import gradio as gr
import time

class DebounceManager:
    def __init__(self, debounce_time: float, tick_time: float, loading_text: str):
        """
        Manages debounce logic and UI updates.

        Args:
            debounce_time (float): The time in seconds to wait before triggering the action.
            tick_time (float): The interval in seconds for the timer tick.
            loading_text (str): The text to display while waiting.
        """
        self.debounce_time = debounce_time
        self.tick_time = tick_time
        self.loading_text = loading_text

    def _generate_progress_html(self, progress_percent: int, remaining_time: float = None) -> str:
        """
        Generates the HTML string for the progress bar.
        
        Args:
            progress_percent (int): Current progress percentage (0-100).
            remaining_time (float, optional): Remaining time in seconds. If None, uses debounce_time.
        
        Returns:
            str: HTML string for the progress bar.
        """
        if remaining_time is None:
            display_time = self.debounce_time
        else:
            display_time = remaining_time
            
        return f"<div style='height: 20px; display: flex; align-items: center;'><progress value='{progress_percent}' max='100' style='width: 100px; margin-right: 10px;'></progress> <span>{self.loading_text} {display_time:.1f}s</span></div>"

    def create_ui(self):
        """
        Creates the necessary UI components for the debounce mechanism.

        Returns:
            tuple: (debounce_state, debounce_timer, debounce_progress)
        """
        # State to store: last_change timestamp, active status, and payload (context)
        debounce_state = gr.State({"last_change": 0, "active": False, "payload": None})
        debounce_timer = gr.Timer(self.tick_time, active=False)
        # Use the new method for initial value
        initial_progress_html = self._generate_progress_html(0) 
        debounce_progress = gr.HTML(value=initial_progress_html, visible=True, elem_classes=["no-transition"])
        return debounce_state, debounce_timer, debounce_progress

    def reset(self, *args):
        """
        Resets the debounce timer. Call this when the monitored input changes.
        Passes through any arguments as the 'payload' to be stored in state.
        """
        # Store all arguments as payload
        payload = args if len(args) > 1 else (args[0] if args else None)
        
        # Use the new method for progress HTML
        progress_html = self._generate_progress_html(0) # Start at 0%

        return {
            "last_change": time.time(),
            "active": True,
            "payload": payload
        }, gr.update(active=True), gr.update(visible=True, value=progress_html)

    def tick(self, debounce_state, trigger_fn):
        """
        Called on every timer tick. checks if debounce time has passed.

        Args:
            debounce_state (dict): The current debounce state.
            trigger_fn (callable): The function to execute when debounce completes. 
                                   It should accept the stored 'payload' as arguments.

        Returns:
             tuple: Updates for (debounce_progress, debounce_state, debounce_timer) + result of trigger_fn
        """
        # 1. If not active, do nothing
        if not debounce_state["active"]:
            # Return empty updates for UI components, and a dummy update for the trigger output
            return gr.update(), debounce_state, gr.update(), gr.update()

        elapsed = time.time() - debounce_state["last_change"]

        # 2. Check if time is up
        if elapsed >= self.debounce_time:
            # Execute the trigger function with the stored payload
            payload = debounce_state["payload"]
            if isinstance(payload, tuple):
                result = trigger_fn(*payload)
            else:
                result = trigger_fn(payload)
            
            # Reset state to inactive
            new_state = {"last_change": 0, "active": False, "payload": None}
            
            # Return: Hide Progress, Update State, Stop Timer, Trigger Result
            return gr.update(value=self._generate_progress_html(0)), new_state, gr.update(active=False), result
        
        else:
            # 3. Update Progress
            progress_percent = int((elapsed / self.debounce_time) * 100)
            remaining = self.debounce_time - elapsed
            
            # Use the new method for progress HTML
            progress_html = self._generate_progress_html(progress_percent, remaining)
            
            # Return: Update Progress, Keep State, Keep Timer, Dummy Update for Result
            return gr.update(value=progress_html, visible=True), debounce_state, gr.update(), gr.update()