import json
import logging

from writer.abstract import register_abstract_template
from writer.blocks.base_block import WriterBlock
from writer.ss_types import AbstractTemplate


class WriterAskGraphQuestion(WriterBlock):
    @classmethod
    def register(cls, type: str):
        super(WriterAskGraphQuestion, cls).register(type)
        register_abstract_template(type, AbstractTemplate(
            baseType="blueprints_node",
            writer={
                "name": "Ask graph question",
                "description": "Asks a natural language question using one or more knowledge graphs and puts the result into a state variable.",
                "category": "Writer",
                "fields": {
                    "question": {
                        "name": "Question",
                        "type": "Text",
                        "desc": "The natural language question to ask.",
                        "control": "Textarea",
                    },
                    "useStreaming": {
                        "name": "Use streaming",
                        "type": "Boolean",
                        "default": "yes",
                        "validator": {
                            "type": "boolean",
                        },
                    },
                    "stateElement": {
                        "name": "Link Variable",
                        "type": "Binding",
                        "desc": "Set the variable here and use it across your agent.",
                    },
                    "graphIds": {
                        "name": "Graph Ids",
                        "type": "Graph Ids",
                        "desc": "IDs of the graphs to query.",
                        "default": "",
                        "validator": {
                            "type": "string",
                        }
                    },
                    "subqueries": {
                        "name": "Use subqueries",
                        "type": "Boolean",
                        "desc": "Enables LLM to ask follow-up questions to the knowledge graph. This improves answers, but may be slower.",
                        "default": "yes",
                        "validator": {
                            "type": "boolean",
                        },
                    },
                    "graphCitations": {
                        "name": "Add inline graph citations",
                        "type": "Boolean",
                        "desc": "Shows what specific graph sources were used to answer the question.",
                        "default": "no",
                        "validator": {
                            "type": "boolean",
                        },
                    }
                },
                "outs": {
                    "success": {
                        "name": "Success",
                        "description": "Successfully streamed the answer.",
                        "style": "success"
                    },
                    "error": {
                        "name": "Error",
                        "description": "There was an error answering the question.",
                        "style": "error"
                    }
                }
            }
        ))

    def run(self):
        try:
            client = self.writer_sdk_client

            graph_ids = self._get_field(
                "graphIds", as_json=True, required=True)
            use_streaming = self._get_field(
                "useStreaming", False, "yes") == "yes"
            if isinstance(graph_ids, str):
                graph_ids = [graph_ids]
            elif not isinstance(graph_ids, list):
                raise ValueError(
                    "graphIds must be a string or a list of strings")
            if len(graph_ids) == 0:
                raise ValueError("graphIds must not be empty")

            question = self._get_field("question", required=True)
            state_element = self._get_field("stateElement", required=False)
            if not state_element and use_streaming:
                raise ValueError(
                    "A state element must be provided when using streaming.")
            subqueries = self._get_field(
                "subqueries", default_field_value="yes") == "yes"
            graph_citations = self._get_field(
                "graphCitations", default_field_value="no") == "yes"

            response = client.graphs.question(
                graph_ids=graph_ids,
                question=question,
                stream=use_streaming,
                subqueries=subqueries,
                query_config= {
                    "inline_citations": graph_citations
                }
            )

            self.result = self._parse_response(response, state_element, use_streaming, graph_citations)
            if state_element:
                self._set_state(state_element, self.result)
            self.outcome = "success"

        except BaseException as e:
            self.outcome = "error"
            raise e
            


    def _parse_response(self, response, state_element, use_streaming: bool, graph_citations: bool):
        if not use_streaming:
            if graph_citations:
                return {"answer": response.answer, "citations": response.sources or []}
            return response.answer

        answer = ""
        citations = []

        for chunk in response:
            try:
                delta_answer = chunk.model_extra.get("answer", "")
                answer += delta_answer

                if graph_citations:
                    delta_sources = chunk.model_extra.get("sources", "")
                    citations.extend(delta_sources)
                    self._set_state(state_element, {"answer": answer, "citations": citations})
                else:
                    self._set_state(state_element, answer)

            except json.JSONDecodeError:
                logging.error("Could not parse stream chunk from graph.question")
        
        if graph_citations:
            return {"answer": answer, "citations": citations}
        return answer
        