from __future__ import annotations

import time
from abc import abstractmethod
from typing import Any, List, Tuple, Union

import gradio_client as grc
import huggingface_hub
from gradio_client.client import Job
from gradio_client.utils import QueueError

try:
    import langchain as lc

    LANGCHAIN_INSTALLED = True
except (ModuleNotFoundError, ImportError):
    LANGCHAIN_INSTALLED = False


class GradioTool:
    def __init__(
        self,
        name: str,
        description: str,
        src: str,
        hf_token: str | None = None,
        duplicate: bool = True,
    ) -> None:
        self.name = name
        self.description = description
        if hf_token and self._is_space(src) and duplicate:
            self.client = grc.Client.duplicate(from_id=src, hf_token=hf_token)
            self.src = self.client.space_id
        else:
            self.src = src
            self.client = grc.Client(self.src, hf_token=hf_token)
        self.description += (
            "If the output is the string QUEUE_FULL you "
            "should try this action again and never stop until you "
            "don't see this message."
        )
        self._block = None

    @staticmethod
    def _is_space(src: str) -> bool:
        try:
            huggingface_hub.get_space_runtime(src)
            return True
        except huggingface_hub.hf_api.RepositoryNotFoundError:
            return False

    @abstractmethod
    def create_job(self, query: str) -> Job:
        pass

    @abstractmethod
    def postprocess(self, output: Union[Tuple[Any], Any]) -> str:
        pass

    def run(self, query: str):
        job = self.create_job(query)
        while not job.done():
            status = job.status()
            print(f"\nJob Status: {str(status.code)} eta: {status.eta}")
            time.sleep(30)
        try:
            output = self.postprocess(job.result())
        except QueueError:
            output = "QUEUE_FULL"
        return output

    # Optional gradio functionalities
    def _block_input(self, gr) -> List["gr.components.Component"]:
        return [gr.Textbox()]

    def _block_output(self, gr) -> List["gr.components.Component"]:
        return [gr.Textbox()]

    def block_input(self) -> List["gr.components.Component"]:
        try:
            import gradio as gr

            GRADIO_INSTALLED = True
        except (ModuleNotFoundError, ImportError):
            GRADIO_INSTALLED = False
        if not GRADIO_INSTALLED:
            raise ModuleNotFoundError("gradio must be installed to call block_input")
        else:
            return self._block_input(gr)

    def block_output(self) -> List["gr.components.Component"]:
        try:
            import gradio as gr

            GRADIO_INSTALLED = True
        except (ModuleNotFoundError, ImportError):
            GRADIO_INSTALLED = False
        if not GRADIO_INSTALLED:
            raise ModuleNotFoundError("gradio must be installed to call block_output")
        else:
            return self._block_output(gr)

    def block(self):
        """Get the gradio Blocks of this tool for visualization."""
        try:
            import gradio as gr
        except (ModuleNotFoundError, ImportError):
            raise ModuleNotFoundError("gradio must be installed to call block")
        if not self._block:
            self._block = gr.load(name=self.src, src="spaces")
        return self._block

    # Optional langchain functionalities
    @property
    def langchain(self) -> "langchain.agents.Tool":  # type: ignore
        if not LANGCHAIN_INSTALLED:
            raise ModuleNotFoundError(
                "langchain must be installed to access langchain tool"
            )

        return lc.agents.Tool(  # type: ignore
            name=self.name, func=self.run, description=self.description
        )

    def __repr__(self) -> str:
        return f"GradioTool(name={self.name}, src={self.src})"

langchain.agents.Tool
[((3706, 3781), 'langchain.agents.Tool', 'lc.agents.Tool', ([], {'name': 'self.name', 'func': 'self.run', 'description': 'self.description'}), '(name=self.name, func=self.run, description=self.description)\n', (3720, 3781), True, 'import langchain as lc\n'), ((742, 794), 'gradio_client.Client.duplicate', 'grc.Client.duplicate', ([], {'from_id': 'src', 'hf_token': 'hf_token'}), '(from_id=src, hf_token=hf_token)\n', (762, 794), True, 'import gradio_client as grc\n'), ((906, 945), 'gradio_client.Client', 'grc.Client', (['self.src'], {'hf_token': 'hf_token'}), '(self.src, hf_token=hf_token)\n', (916, 945), True, 'import gradio_client as grc\n'), ((1259, 1297), 'huggingface_hub.get_space_runtime', 'huggingface_hub.get_space_runtime', (['src'], {}), '(src)\n', (1292, 1297), False, 'import huggingface_hub\n'), ((1807, 1821), 'time.sleep', 'time.sleep', (['(30)'], {}), '(30)\n', (1817, 1821), False, 'import time\n'), ((2092, 2104), 'gradio.Textbox', 'gr.Textbox', ([], {}), '()\n', (2102, 2104), True, 'import gradio as gr\n'), ((2191, 2203), 'gradio.Textbox', 'gr.Textbox', ([], {}), '()\n', (2201, 2203), True, 'import gradio as gr\n'), ((3342, 3378), 'gradio.load', 'gr.load', ([], {'name': 'self.src', 'src': '"""spaces"""'}), "(name=self.src, src='spaces')\n", (3349, 3378), True, 'import gradio as gr\n')]