import json
from abc import abstractmethod
from typing import Any, List, Optional, Union

import requests

from steamship import Block, File, MimeTypes, Task
from steamship.agents.base import BaseTool
from steamship.agents.context.context import AgentContext


class Tool(BaseTool):
    # Working thinking: we don't yet have formalization about whether
    # this is a class-level name, isntance-level name, or
    # instance+context-level name.
    # thought(doug): this should be the planner-facing name (LLM-friendly?)
    name: str

    # Advice, but not hard-enforced:
    # This contains the description, inputs, and outputs.
    ai_description: str
    human_description: str  # Human readable string for logging

    @abstractmethod
    def run(self, tool_input: List[Block], context: AgentContext) -> Union[List[Block], Task[Any]]:
        raise NotImplementedError()

    # This gets called later if you return Task[Any] from run
    def post_process(self, async_task: Task, context: AgentContext) -> List[Block]:
        # nice helpers for making lists of blocks
        pass


class GeneratorTool(Tool):
    """
    A base class for tools that wrap Steamship Generator plugins. Subclass this and implement
    the `accept_output_block` method.
    """

    generator_plugin_handle: str
    generator_plugin_instance_handle: Optional[str] = None
    generator_plugin_config: dict = {}
    merge_blocks: bool = False

    @abstractmethod
    def accept_output_block(self, block: Block) -> bool:
        raise NotImplementedError()

    def run(self, tool_input: List[Block], context: AgentContext) -> Union[List[Block], Task[Any]]:
        generator = context.client.use_plugin(
            plugin_handle=self.generator_plugin_handle,
            instance_handle=self.generator_plugin_instance_handle,
            config=self.generator_plugin_config,
        )

        tasks = []

        if self.merge_blocks:
            block = tool_input[0]
            for extra_block in tool_input[1:]:
                block.text = f"{block.text}\n\n{extra_block.text}"
            tool_input = [block]

        for block in tool_input:
            if isinstance(block, str):
                print(block)
            if not block.is_text():
                continue

            prompt = block.text
            task = generator.generate(text=prompt, append_output_to_file=True)
            tasks.append(task)

        # TODO / REMOVE Synchronous execution is a temporary simplification while we merge code.
        output_blocks = []
        # NB: In an async framework, we need to contend with the fact that we have some Monad-style trickery
        # to contend with. Namely that the below code has taken us from List[Block] as the universal type to
        # List[List[Block]] as the async result that we will have to process.
        #
        # It's unclear to me (ted) if this is something we leave as an exercise for implementors or build in
        # as universally handled.
        for task in tasks:
            task.wait()
            task_blocks = self.post_process(task, context)
            for block in task_blocks:
                output_blocks.append(block)
        # END TODO / REMOVE

        return output_blocks

    def post_process(self, task: Task, context: AgentContext) -> List[Block]:
        """In this case, the Generator returns a GeneratorResponse that has a .blocks method on it"""
        try:
            return task.output.blocks
        except BaseException:
            return json.loads(task.output).get("blocks")


class ImageGeneratorTool(GeneratorTool):
    """
    A base class for tools that wrap Steamship Image Generator plugins.
    """

    def accept_output_block(self, block: Block) -> bool:
        return block.is_image()


class AudioGeneratorTool(GeneratorTool):
    """
    A base class for tools that wrap Steamship Audio Generator plugins.
    """

    def accept_output_block(self, block: Block) -> bool:
        return block.is_audio()


class ScrapeAndBlockifyTool(Tool):
    """
    A base class for tools that wrap Steamship Blockifier plugin which transforms bytes to a set of blocks.
    """

    blockifier_plugin_handle: str
    blockifier_plugin_instance_handle: Optional[str] = None
    blockifier_plugin_config: dict = {}

    @abstractmethod
    def should_blockify(self, block: Block) -> bool:
        raise NotImplementedError()

    def get_mime_type(self):
        return None

    def _scrape(self, url: str, context: AgentContext) -> File:
        response = requests.get(url)
        file = File.create(context.client, content=response.content, mime_type=self.get_mime_type())
        return file

    def run(self, tool_input: List[Block], context: AgentContext) -> Union[List[Block], Task[Any]]:
        tasks = []

        blockifier = context.client.use_plugin(
            plugin_handle=self.blockifier_plugin_handle,
            instance_handle=self.blockifier_plugin_instance_handle,
            config=self.blockifier_plugin_config,
        )

        for input_block in tool_input:
            if not input_block.is_text():
                continue

            url = input_block.text
            file = self._scrape(url, context)
            task = file.blockify(blockifier.handle)
            tasks.append(task)

        # TODO / REMOVE Synchronous execution is a temporary simplification while we merge code.
        output_blocks = []
        # NB: In an async framework, we need to contend with the fact that we have some Monad-style trickery
        # to contend with. Namely that the below code has taken us from List[Block] as the universal type to
        # List[List[Block]] as the async result that we will have to process.
        #
        # It's unclear to me (ted) if this is something we leave as an exercise for implementors or build in
        # as universally handled.
        for task in tasks:
            task.wait()
            task_blocks = self.post_process(task, context)
            for block in task_blocks:
                output_blocks.append(block)
        # END TODO / REMOVE

        return output_blocks

    def post_process(self, task: Task, context: AgentContext) -> List[Block]:
        """In this case, the Blockifier returns a BlockAndTagResponse that has a .file.blocks method on it"""
        try:
            return task.output.file.blocks
        except BaseException:
            return json.loads(task.output).get("file", {}).get("blocks")


class ImageBlockifierTool(ScrapeAndBlockifyTool):
    """
    A base class for tools that wrap Steamship Image Blockifier plugins.
    """

    def get_mime_type(self):
        return MimeTypes.PNG

    def should_blockify(self, block: Block) -> bool:
        return block.is_image()


class AudioBlockifierTool(ScrapeAndBlockifyTool):
    """
    A base class for tools that wrap Steamship Audio Blockifier plugins.
    """

    def get_mime_type(self):
        return MimeTypes.MP3

    def should_blockify(self, block: Block) -> bool:
        return block.is_audio()
