import json, json_repair
from typing import List, Dict
from haystack import Pipeline, component
from haystack.components.builders import PromptBuilder
from haystack_integrations.components.generators.google_ai import GoogleAIGeminiGenerator

template = """
    次のテキストに基づいて、JSON 形式で 回答を生成して下さい。
    テキストの「明確さ」を様々な点から評価し、100点満点で点数を採点してください。
    次にその点数とした理由の解説と、どの辺が曖昧だったかを説明してください。
    テキストをより明確にする上で修正するべき箇所をピックアップし、その箇所が何行目の何文字目であるかと、修正前のテキストと修正後のテキストをまとめて提案して下さい。

    以上のことをマークダウンや説明なしで、JSON のみで回答して下さい。
    次のようなJSONフォーマット形式に必ず従ってください。

    {
        "score" : 採点した点数(数値) ,
        "description": "その点数になった理由",
        "suggestions":[
            {
                "position": [何行目(数値) , 何文字目(数値)],
                "before": "修正前のテキスト",
                "after":  "修正後のテキスト"
            }, ...
        ]
    }

    テキスト: {{text}}
"""

@component
class ResultParser:
    @component.output_types(parsed=Dict)
    def run(self, replies: List[str]):
        reply = replies[0]
        first_index = min(reply.find("{"), reply.find("["))
        last_index = max(reply.rfind("}"), reply.rfind("]")) + 1
        json_portion = reply[first_index:last_index]

        try:
            parsed = json.loads(json_portion)
        except json.JSONDecodeError:
            parsed = json_repair.loads(json_portion)
        if isinstance(parsed, list):
            parsed = parsed[0]
        return {"parsed": parsed}

prompt_builder =  PromptBuilder(template=template)
llm = GoogleAIGeminiGenerator(model="models/gemini-1.5-flash")
parser = ResultParser()

pipe = Pipeline()
pipe.add_component("prompt_builder",prompt_builder)
pipe.add_component("llm",llm)
pipe.add_component("parser",parser)
pipe.connect("prompt_builder", "llm")
pipe.connect("llm.replies", "parser.replies")