import yaml
from pocketflow import Node

from pf_agent.utils import callLLM, callWebSearch, callWikipediaSearch

class Decide(Node):
    def setLogger(self,logger):
        self.logger = logger

    def prep(self,shared):
        context = shared.get("context","No previous search")
        question = shared["question"]
        return question,context
    
    def exec(self,inputs):
        question,context = inputs
        prompt = f"""
### CONTEXT
You are a research assistant and I need your help to answer the following question:
Question: {question}
Previous research: {context}

### ACTION SPACE
[1] search
  Description: Look up more information on the web
  Parameters:
  - query (str): keywords to run the search on the web 

[2] wikipedia
  Description: Look up for more information in the wikipedia in English
  Parameters:
  - query (str): keywords to run the search in the English wikipedia

[3] answer
  Description: Answer the question with current knowledge
  Parameters:
  - answer (str): final answer to the question


### NEXT ACTION
Decide the next action based on the context and available actions.
Return your response in the following format:
```yaml
thinking: |
  <your step-by-step reasoning process>
action: search OR answer OR wikipedia
reason: <why you choose this action>
parameters: <yaml containing the parameter or parameters required to execute the action selected in yaml format indicating parameter name and value>
```
IMPORTANT: Make sure to:
1. Use proper indentation (4 spaces) for all multi-line fields
2. Use the | character for multi-line text fields
3. Keep single-line fields without the | character
"""
        self.logger.debug(f"=== Calling LLM to DECIDE")
        self.logger.debug(f"Context: {context}")
        response = callLLM(prompt)
        yaml_str = response.replace("|","") #.split("```yaml")[1].split("```")[0].strip()  #.replace("|","") #
        decision = yaml.safe_load(yaml_str)
        self.logger.debug("LLM responded")
        self.logger.debug(f"Thinking: {decision['thinking']}")
        self.logger.debug(f"Action: {decision['action']} {decision['parameters']}")
        return decision
    
    def post(self,shared,prep_res,exec_res):
        if exec_res["action"] != "answer":
            shared["parameters"] = exec_res["parameters"]
        return exec_res["action"]
    

class Search(Node):
    def setLogger(self,logger):
        self.logger = logger

    def prep(self,shared):
        if type(shared["parameters"]) == type([]):
            for p in shared["parameters"]:
                if 'query' in p:
                    return p['query']
        else:
            return shared["parameters"]["query"]
        return "invalid query"
    
    def exec(self,search_query):
        self.logger.debug("=== Searching the web")
        self.logger.debug(f"Search query: {search_query}")
        results = callWebSearch(search_query)
        self.logger.debug(f"Results: {results}")
        return results
    
    def post(self,shared,prep_res,exec_res):
        previous = shared.get("context","")
        shared["context"] = f"{previous}\n\nSEARCH: {shared['parameters']}\n\nRESULTS: {exec_res}"
        return "decide"
    
class Wikipedia(Node):
    def setLogger(self,logger):
        self.logger = logger

    def prep(self,shared):
        if type(shared["parameters"]) == type([]):
            for p in shared["parameters"]:
                if 'query' in p:
                    return shared["question"],p['query']
        else:
            return shared["question"],shared["parameters"]["query"]
        return shared["question"],"invalid query"
    
    def exec(self,prepared):
        question,search_query = prepared
        self.logger.debug("=== Searching wikipedia")
        self.logger.debug(f"Search query: {search_query}")
        results = callWikipediaSearch(search_query)
        self.logger.debug(f"Results: {results}")
        prompt = f'''You are my research assistant and I need your help to answer the following question: {question}

From my previous research I have found a Wikipedia page that could contain the answer or, at least, useful information for further research. 
Take a careful look at it. Ignore any HTML markup you may find. You will find sections, paragraphs, etc. formated as plain text and tables
as comma-separated values. The answer or any useful information could be anywhere.
If the text contains the answer, give it to me. 
If it is not, extract relevant information that could me help find it elsewere. 
Please don't include information not directly related to the question.
It is very important that in any case, you are concise, brief and stick to the question being asked. Find the text below:
{results}'''
        response = callLLM(prompt)
        return response
    
    def post(self,shared,prep_res,exec_res):
        previous = shared.get("context","")
        shared["context"] = f"{previous}\n\nWIKIPEDIA RESULTS: {shared['parameters']}\n\nRESULTS: {exec_res}"
        return "decide"    

class Answer(Node):
    def setLogger(self,logger):
        self.logger = logger

    def prep(self,shared):
        return shared["question"],shared.get("context","")
    
    def exec(self,inputs):
        question,context = inputs
        prompt = f'''
### CONTEXT
Based on the following information, answer the question.
Question: {question}
Research: {context}

### YOUR ANSWER
Provide a comprehensive answer using research results
'''
        response = callLLM(prompt)
        self.logger.debug(f"### CALLING LLM\n{prompt}\n### LLM RESPONSE\n{response}\n\n##########\n\n")
        return response
    
    def post(self,shared,prep_res,exec_res):
        shared["answer"] = exec_res