#!/usr/bin/env python3

from trustgraph.clients.prompt_client import PromptClient
import json
import textwrap
import ibis
import textwrap
import time
import dataclasses

class Input:
    state: dict
    plan_state: str
    plan: str
    iteration: int

class Output:
    state: dict
    thought: str
    next_state: str
    plan: str
    iteration: int
    

def output(text, prefix="> ", width=78):

    out = textwrap.wrap(
        text, initial_indent=prefix, subsequent_indent=prefix,
        width=width
    )
    print("\n".join(out))

pulsar_host = "pulsar://localhost:6650"

prompt_client = PromptClient(pulsar_host=pulsar_host)

class State:
    pass

class ActionWikipedia:

    def __init__(self, prompt_client):
        self.prompt_client = prompt_client

    def act(self, state, topic):

        pr = ibis.Template(
            "Tell me what you think Wikipedia would say about {{topic}}"
        )

        input = pr.render({
            "topic": topic,
        })

        resp = self.prompt_client.request(
            "question",
            {
                "question": input
            }
        )

        state.previous.append(
            f"Question: {input.strip()}"
        )

        state.previous.append(
            f"Answer: {resp.strip()}"
        )

        return state

class ActionCalculate:

    def __init__(self, prompt_client):
        self.prompt_client = prompt_client

    def act(self, state, calc):

        pr = ibis.Template(
            "Compute this: {{computation}}.  Just answer the question without explanation."
        )

        input = pr.render({
            "computation": calc,
        })

        resp = self.prompt_client.request(
            "question",
            {
                "question": input
            }
        )

        print(resp)

        state.previous.append(
            "Calculate: " + calc.strip()
        )

        state.previous.append(
            "Answer: " + resp.strip()
        )

        return state

class ActionAnswer:

    def __init__(self, prompt_client):
        self.prompt_client = prompt_client

    def act(self, state, question):

        pr = ibis.Template(
            "Answer this: {{question}}"
        )

        input = pr.render({
            "question": question,
        })

        resp = self.prompt_client.request(
            "question",
            {
                "question": input
            }
        )

        state.facts.append(
            "Answer to " + input + ":\n" +
            resp
        )

        return state

tools = {
    "calculate": {
        "description": "Takes a numeric computation and calculates the answer",
        "implementation": ActionCalculate,
    },
    "cats-knowledge-store": {
        "description": "Answer questions on Mark's cats using a knowledge store",
        "implementation": ActionCalculate,
    },
    # "wikipedia": {
    #     "description": "Takes a query and looks it up on Wikipedia",
    #     "implementation": ActionWikipedia,
    # },
    # "answer": {
    #     "description": "Take a simple question and use the LLM to provide the answer",
    #     "implementation": ActionAnswer,
    # },
    # "space-shuttle-rag": {
    #     "description": "Take a question about space shuttles and answer it using the TrustGraph GraphRAG service",
    #     "implementation": ActionAnswer,
    # }
}

class AgentManager:

    determine_prompt = ibis.Template("""
You have access to the following tools:
{% for id, tool in tools.items() %}- tool-name: {{id}}
  tool-description: {{tool.description}}
{% endfor %}
You operate in a loop.  For each iteration of the loop, you take the
question and learnt knowledge.  If the knowledge is enough to answer
the question, you will response with an answer.  If the knowledge is not
enough, you will response with an action to be invoked to acquire more
knowledge.

Your output must ALWAYS be a well-formed JSON object.  Output ONLY a JSON
object with no explanatory text or markup formatting.  The JSON object
can have the following fields:
- thought: Mandatory.  Your explanation of the current step, including the
  goal and the reason for choosing this action.  A string.
- action: An optional field, one of the tool-names above if you are using
  an action.
- action-input: An optional field, input to the selected tool if you are using
  an action.
- failure: a boolean.  Set to true if you cannot proceed.
- final-answer: Your answer to the question as a string.

When you know the final answer to the original question, emit a 'final-answer'
JSON object with the following fields: thought and final-answer.

When you don't know the answer but know an action to move towards the
solution, emit an 'action' JSON object with the following fields:
thought, action, action-input

When you are stuck and aren't able to move torwards a solution, emit a
'failure' JSON object the following fields: thought, failure

{{previous}}""")

    def __init__(self, tools, prompt_client):

        self.tools = tools
        self.prompt_client = prompt_client

    def determine(self, state, thought=None):

        input = __class__.determine_prompt.render({
            "previous": "\n".join(state.previous),
            "tools": self.tools,
        })

        print(input)

        resp = self.prompt_client.request(
            "question",
            {
                "question": input
            }
        )

        resp = resp.replace("```json", "")
        resp = resp.replace("```", "")
        print(resp)
    #     print(resp)
        resp = json.loads(resp)
    #     print(json.dumps(resp, indent=4))
        return resp

    def invoke(self, q, thought):

        state = State()
        state.question = q
        state.previous = [
            f"Question: {q}"
        ]

        while True:

            resp = self.determine(state)

            print(resp)

            if "thought" in resp:
                if thought:
                    thought(resp["thought"])
                state.previous.append(f"Thought: {resp['thought'].strip()}")

            if "failure" in resp:
                if resp["failure"]:
                    print("Failed")
                    return "Failed"

            if "final-answer" in resp:
                return resp["final-answer"]

            if "action" not in resp:
                raise RuntimeError("Didn't get final-answer/action response")

            print(resp)

            action = resp["action"]

            print(action)

            if action not in self.tools:
                raise RuntimeError(f"Tool {action} not known")

            if "implementation" not in self.tools[action]:
                raise RuntimeError(f"Tool {action} not implemented")

            if "action-input" in resp:
                input = resp["action-input"]
            else:
                input = None

            print("Action:", action)
            print("Action input:", input)
            print()

            impl_class = self.tools[action]["implementation"]
            impl = impl_class(self.prompt_client)
            state = impl.act(state, input)

        time.sleep(2)
               
def thought(t):
    output("\U0001f914... " + t, prefix="? ")
    print()

q = """Take the square root of 1600.  That number is the number of a US
President.  Who is that president and how do they relate to the space
shuttle programme?
"""

q = "How many eggs do I have if I have 3 eggs and Fred gives me some more eggs.  The number of eggs he gives me is the square root of 64."

q = "Find out the behavioral nature of Mark's cats"

am = AgentManager(tools, prompt_client)

output("Q: " + q)
print()

resp = am.invoke(q, thought)

print(resp)