import json
from typing import Dict, List

from haystack import Pipeline, component
from haystack.components.builders import PromptBuilder
from haystack.components.generators import OpenAIGenerator
from haystack.utils import Secret

@component()
class QueryMetadataExtractor:

    def __init__(self):
        prompt = """
          Bạn là một phần của hệ thống thông tin xử lý các truy vấn của người dùng.
          Cho một truy vấn từ người dùng, bạn trích xuất thông tin khớp với danh sách các trường metadata đã cho.
          Thông tin được trích xuất từ truy vấn phải khớp với ngữ nghĩa liên quan đến các trường metadata đã cho.
          Thông tin bạn trích xuất từ truy vấn sẽ được sử dụng như bộ lọc để thu hẹp không gian tìm kiếm
          khi truy vấn một chỉ mục.
          Chỉ bao gồm giá trị của metadata đã trích xuất mà không bao gồm tên của trường metadata.
          Thông tin đã trích xuất trong 'Extracted metadata' phải trả về dưới dạng cấu trúc JSON hợp lệ.
          ###
          Ví dụ 1:
          Truy vấn: "Luật được Quốc hội ban hành năm 2014?"
          Metadata fields: {"publisher", "publish_year", "document_type"}
          Extracted metadata fields: {"publisher": "Quốc hội", "publish_year": 2014, "document_type":"Luật"}
          ###
          Ví dụ 2:
          Truy vấn: "Nghị định được Chính phủ ban hành năm 2021?"
          Trường metadata: {"publisher", "publish_year", "document_type"}
          Extracted metadata fields: {"publisher": "Chính phủ", "publish_year": 2021, "document_type":"Nghị định"}
          ###
          Ví dụ 3:
          Truy vấn: "{{query}}"
          Trường metadata: "{{metadata_fields}}"
          Extracted metadata fields:
          """
        generator = OpenAIGenerator(
            api_key=Secret.from_env_var("OCTOAI_TOKEN"),
            api_base_url="https://text.octoai.run/v1",
            model="meta-llama-3-70b-instruct",
            generation_kwargs = {"max_tokens": 512,"response_format": { "type": "json_object" }}
        )
        self.pipeline = Pipeline()
        self.pipeline.add_component(name="builder", instance=PromptBuilder(prompt))
        self.pipeline.add_component(name="llm", instance=generator)
        self.pipeline.connect("builder", "llm")

    @component.output_types(filters=Dict[str, str])
    def run(self, query: str, metadata_fields: List[str]):
        result = self.pipeline.run({'builder': {'query': query, 'metadata_fields': metadata_fields}})
        
        import json
        first_reply = result['llm']['replies'][0]
        parsed_data = json.loads(first_reply)
        metadata = parsed_data['answer']
        metadata = metadata.replace("'",'"')
        data = json.loads(metadata)
        # this can be done with specific data structures and in a more sophisticated way
        filters = []
        for key, value in data.items():
            field = f"meta.{key}"
            filters.append({f"field": field, "operator": "==", "value": value})
        print(filters)
        return {"filters": {"operator": "AND", "conditions": filters}}