#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2024/6/27
@Author  : mashenquan
@File    : pic2txt.py
"""
import json
from pathlib import Path
from typing import List

from tenacity import retry, stop_after_attempt, wait_random_exponential

from metagpt.actions import Action
from metagpt.logs import logger
from metagpt.tools.tool_registry import register_tool
from metagpt.utils.common import encode_image, general_after_log, to_markdown_code_block


@register_tool(include_functions=["run"])
class Pic2Txt(Action):
    """Pic2Txt deal with the following situations:
    Given some pictures depicting user requirements alongside contextual description, write out the intact textual user requirements.
    """

    async def run(
        self,
        *,
        image_paths: List[str],
        textual_user_requirement: str = "",
        legacy_output: str = "",
        evaluation_conclusion: str = "",
        additional_technical_requirements: str = "",
    ) -> str:
        """
        Given some pictures depicting user requirements alongside contextual description, write out the intact textual user requirements

        Args:
            image_paths (List[str]): A list of file paths to the input image(s) depicting user requirements.
            textual_user_requirement (str, optional): Textual user requirement that alongside the given images, if any.
            legacy_output (str, optional): The intact textual user requirements generated by you last time, if any.
            evaluation_conclusion (str, optional): Conclusion or evaluation based on the processed requirements.
            additional_technical_requirements (str, optional): Any supplementary technical details relevant to the process.

        Returns:
            str: Textual representation of user requirements extracted from the provided image(s).

        Raises:
            ValueError: If image_paths list is empty.
            OSError: If there is an issue accessing or reading the image files.

        Example:
            >>> images = ["requirements/pic/1.png", "requirements/pic/2.png", "requirements/pic/3.png"]
            >>> textual_user_requirements = "User requirement paragraph 1 ..., ![](1.png). paragraph 2...![](2.png)..."
            >>> action = Pic2Txt()
            >>> intact_textual_user_requirements = await action.run(image_paths=images, textual_user_requirement=textual_user_requirements)
            >>> print(intact_textual_user_requirements)
            "User requirement paragraph 1 ..., ![...](1.png) This picture describes... paragraph 2...![...](2.png)..."

        """
        descriptions = {}
        for i in image_paths:
            filename = Path(i)
            base64_image = encode_image(filename)
            rsp = await self._pic2txt(
                "Generate a paragraph of text based on the content of the image, the language of the text is consistent with the language in the image.",
                base64_image=base64_image,
            )
            descriptions[filename.name] = rsp

        prompt = PROMPT.format(
            textual_user_requirement=textual_user_requirement,
            acknowledge=to_markdown_code_block(val=json.dumps(descriptions), type_="json"),
            legacy_output=to_markdown_code_block(val=legacy_output),
            evaluation_conclusion=evaluation_conclusion,
            additional_technical_requirements=to_markdown_code_block(val=additional_technical_requirements),
        )
        return await self._write(prompt)

    @retry(
        wait=wait_random_exponential(min=1, max=20),
        stop=stop_after_attempt(6),
        after=general_after_log(logger),
    )
    async def _write(self, prompt: str) -> str:
        rsp = await self.llm.aask(prompt)
        return rsp

    @retry(
        wait=wait_random_exponential(min=1, max=20),
        stop=stop_after_attempt(6),
        after=general_after_log(logger),
    )
    async def _pic2txt(self, prompt: str, base64_image: str) -> str:
        rsp = await self.llm.aask(prompt, images=base64_image)
        return rsp


PROMPT = """
## Textual User Requirements
{textual_user_requirement}

## Acknowledge
{acknowledge}

## Legacy Outputs
{legacy_output}

## Evaluation Conclusion
{evaluation_conclusion}

## Additional Technical Requirements
{additional_technical_requirements}

---
You are a tool that generates an intact textual user requirements given a few of textual fragments of user requirements and some fragments of UI pictures.
The content of "Textual User Requirements" provides a few of textual fragments of user requirements;
The content of "Acknowledge" provides the descriptions of pictures used in "Textual User Requirements";
"Legacy Outputs" contains the intact textual user requirements generated by you last time, which you can improve by addressing the issues raised in "Evaluation Conclusion";
"Additional Technical Requirements" specifies the additional technical requirements that the generated textual user requirements must meet;
You need to merge the text content of the corresponding image in the "Acknowledge" into the "Textual User Requirements" to generate a complete, natural and coherent description of the user requirements;
Return the intact textual user requirements according to the given fragments of the user requirement of "Textual User Requirements" and the UI pictures;
"""
