from haystack import Document, Pipeline,component
from haystack.components.embedders import SentenceTransformersDocumentEmbedder,SentenceTransformersTextEmbedder
from haystack_integrations.document_stores.astra import AstraDocumentStore
from haystack.components.writers import DocumentWriter
from haystack.document_stores.types import DuplicatePolicy
from haystack.components.builders import PromptBuilder
from haystack_integrations.components.generators.google_ai import GoogleAIGeminiGenerator
from haystack_integrations.components.retrievers.astra import AstraEmbeddingRetriever
import json, json_repair
from typing import List, Dict
import os

document_store = AstraDocumentStore(
    collection_name='company_contents',
    duplicates_policy=DuplicatePolicy.OVERWRITE
)

retrieve_template = '''
    次の日本語の文章の明確さを判定したい
    テキスト: {{text}}
'''

template = """
    どのように日本語の文章を記載するとよいかまとめた自社の資料があります。
    この資料を参考にして、以下のテキストの「明確さ」を様々な点から評価し、100点満点で点数を採点してください。
    次にその点数とした理由の解説と、どの辺が曖昧だったかを説明してください。
    テキストをより明確にする上で修正するべき箇所をピックアップし、その箇所が何行目の何文字目であるかと、修正前のテキストと修正後のテキストをまとめて提案して下さい。
    以上のことをマークダウンや説明なしで、JSON のみで回答して下さい。
    次のようなJSONフォーマット形式に必ず従ってください。
    {
        "score" : 採点した点数(数値) ,
        "description": "その点数になった理由",
        "suggestions":[
            {
                "position": [何行目(数値) , 何文字目(数値)],
                "before": "修正前のテキスト",
                "after":  "修正後のテキスト"
            }, ...
        ]
    }
    資料：
      {% for document in documents %}
        {{ document.content }}
      {% endfor %}
    テキスト: {{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}

retrieve_prompt_builder = PromptBuilder(template=retrieve_template)
text_embedder =  SentenceTransformersTextEmbedder(model=os.environ['embedding_model'])
embedding_retriever = AstraEmbeddingRetriever(document_store=document_store)
prompt_builder =  PromptBuilder(template=template)
llm = GoogleAIGeminiGenerator(model=os.environ["gemini_model"])
parser = ResultParser()
text_embedder.warm_up()

pipe = Pipeline()
pipe.add_component("retrieve_prompt_builder",retrieve_prompt_builder)
pipe.add_component("text_embedder",text_embedder)
pipe.add_component("embedding_retriever",embedding_retriever)
pipe.add_component("prompt_builder",prompt_builder)
pipe.add_component("llm",llm)
pipe.add_component("parser",parser)
pipe.connect("retrieve_prompt_builder.prompt", "text_embedder.text")
pipe.connect("text_embedder", "embedding_retriever")
pipe.connect("embedding_retriever.documents", "prompt_builder.documents")
pipe.connect("prompt_builder", "llm")
pipe.connect("llm.replies", "parser.replies")