import time
from collections.abc import Generator
from typing import Callable, Sequence, Union

import gradio as gr
from gradio_client import Client

from groovy.types import Image, String
from groovy.utils import resolve_space_url_to_id


class Flow:
    """
    The core class in Groovy. A Flow requires three things:
    * A task to run (string or format string)
    * A set of input components (optional)
    * A generator function that accepts a task string and yields 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.Component] | None = None,
        agent_fn: Callable[[str], Generator[Union[String, Image], 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.
            agent_fn: Fancy name for a generator function that accepts a task string and yields an arbitrary number of `gv.String()` or `gv.Image()` responses. If not provided, Groovy includes a built-in browser agent function, which browses the web to complete a task.
        """
        # Import here to speed up the import time of the groovy module

        self.task = task
        self.inputs = inputs or []
        if agent_fn is None:
            try:
                from groovy.agent import browser_agent_fn
            except ImportError:
                raise ImportError(
                    "The browser agent function is not available. Please install the full version of Groovy: `pip install groovy[full]`."
                )
            self.agent_fn = browser_agent_fn
        else:
            self.agent_fn = agent_fn

    def launch(
        self,
        save_recording: bool = True,
        run_immediately: bool = False,
        **kwargs,
    ):
        """
        Launches the flow in a Gradio app.

        Parameters:
            save_recording: If True, collects any images generated by the flow and saves them as a gif in 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.
            kwargs: Additional keyword arguments to pass to the Gradio app's launch() method.
        """
        # Import here to speed up the import time of the groovy module
        from groovy.app import create_app
        from groovy.utils import open_positioned_browser

        self.save_recording = save_recording
        self.run_immediately = run_immediately
        self.app = create_app(
            self,
            self.inputs,
            self.task,
            self.agent_fn,
            self.run_immediately,
            self.save_recording,
        )
        _, self.url, _ = self.app.launch(
            inline=False, inbrowser=False, prevent_thread_lock=True
        )
        open_positioned_browser(self.url, width=500)
        while 1:
            time.sleep(0.001)
        return self.url

    def run(self, *args):
        output = None
        for step in self.agent_fn(self.task, *args):
            output = step
        return output

    def to_json(self) -> dict:
        data = {
            "task": self.task,
            "inputs": [input_comp.get_config() for input_comp in self.inputs],
        }
        return data

    @staticmethod
    def component_from_config(config: dict) -> gr.Component:
        cls = gr.utils.component_or_layout_class(config.pop("name"))
        constructor_args = cls.recover_kwargs(config)
        return cls(**constructor_args)

    @classmethod
    def from_json(
        cls,
        data: dict,
        agent_fn: Callable[[str], Generator[Union[String, Image], None, None]]
        | None = None,
    ):
        return cls(
            task=data["task"],
            inputs=[
                cls.component_from_config(input_config)
                for input_config in data["inputs"]
            ],
            agent_fn=agent_fn,
        )

    @classmethod
    def from_space(
        cls,
        space_url: str,
        agent_fn: Callable[[str], Generator[Union[String, Image], None, None]]
        | None = None,
    ):
        space_id = resolve_space_url_to_id(space_url)
        client = Client(space_id)
        config = client.predict(api_name="/flow_config")
        return cls.from_json(config, agent_fn)
