import os
from collections.abc import Generator
from typing import TYPE_CHECKING, Callable, Sequence, Union

if TYPE_CHECKING:
    import gradio as gr


class Flow:
    """
    The core class in Groovy. A Flow consists of:
    * A task to run (string or format string)
    * A set of input components (optional)
    * A function that accepts a task string and yields an arbitrary number of `str`, `PIL.Image`, or `gr.ChatMessage` responses.

    After instantiating a Flow, you can call .launch() to launch the flow in a Gradio app.
    """

    def __init__(
        self,
        task: str,
        inputs: Sequence[gr.components.Component] | None = None,
        stream_fn: Callable[[str], Generator[Union[str, gr.ChatMessage], None, None]]
        | None = None,
    ):
        """
        Parameters:
            task: The task to run. Can be a regular string or a format string, in which case the input components' values will be passed to it.
            inputs: The input components whose values will be passed to the task, if it's a format string.
            stream_fn: The generator function that accepts a task string and yields an arbitrary number of `str`, `PIL.Image`, or `gr.ChatMessage` responses. If not provided, the default streamer (which browses the web to complete a task) will be used.
        """
        # Import here to speed up the import time of the groovy module
        from groovy.agent import browser_agent_streamer

        self.task = task
        self.inputs = inputs or []
        self.stream_fn = stream_fn or browser_agent_streamer

    def launch(self, artifacts_dir: str | None = None, run_immediately: bool = False):
        """
        Launches the flow in a Gradio app.

        Parameters:
            artifacts_dir: The directory to save artifacts (screenshots, etc.) to. If None, uses the working directory.
            run_immediately: Whether to run the task immediately at .launch() or whether to wait for the user to click the "Run" button in the Gradio app.
        """
        # Import here to speed up the import time of the groovy module
        from groovy.app import create_app

        self.artifacts_dir = artifacts_dir or os.getcwd()
        self.app = create_app(
            self,
            self.inputs,
            self.task,
            self.stream_fn,
            run_immediately,
            self.artifacts_dir,
        )
        _, self.url, _ = self.app.launch(inline=False, inbrowser=True)
        return self.url
