import json
from actions.duck_search import duckduckgo_search
from processing.text import read_txt_files
from agent.llm_utils import llm_response, llm_stream_response
from config import Config
from agent import prompts
import os
import string

CFG = Config()


class ResearchAgent:
    def __init__(self, question, agent):
        """ Initializes the research assistant with the given question.
            Args: question (str): The question to research
            Returns: None
        """

        self.question = question
        self.agent = agent
        self.visited_urls = set()
        self.search_summary = ""
        self.directory_name = ''.join(c for c in question if c.isascii() and c not in string.punctuation)[:100]
        self.dir_path = os.path.dirname(f"./outputs/{self.directory_name}/")

    def call_agent(self, action):
        messages = [{
            "role": "system",
            "content": prompts.generate_agent_role_prompt(self.agent),
        }, {
            "role": "user",
            "content": action,
        }]
        return llm_response(
                    model=CFG.fast_llm_model,
                    messages=messages,
                )

    def call_agent_stream(self, action):
        messages = [{
            "role": "system",
            "content": prompts.generate_agent_role_prompt(self.agent),
        }, {
            "role": "user",
            "content": action,
        }]
        yield from llm_stream_response(
                model=CFG.fast_llm_model,
                messages=messages
            )

    def create_search_queries(self):
        """ Creates the search queries for the given question.
        Args: None
        Returns: list[str]: The search queries for the given question
        """
        result = self.call_agent(prompts.generate_search_queries_prompt(self.question))
        return json.loads(result)

    def search_single_query(self, query):
        """ Runs the async search for the given query.
        Args: query (str): The query to run the async search for
        Returns: list[str]: The async search for the given query
        """
        return duckduckgo_search(query, max_search_result=3)

    def run_search_summary(self, query):
        """ Runs the search summary for the given query.
        Args: query (str): The query to run the search summary for
        Returns: str: The search summary for the given query
        """
        responses = self.search_single_query(query)
        
        print(f"Searching for {query}")
        query = hash(query)
        file_path = f"./outputs/{self.directory_name}/research-{query}.txt"
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        with open(file_path, "w") as f:
            json.dump(responses, f)
            print(f"Saved {query} to {file_path}")
        return responses

    def search_online(self):
        """ Conducts the search for the given question.
        Args: None
        Returns: str: The search results for the given question
        """

        self.search_summary = read_txt_files(self.dir_path) if os.path.isdir(self.dir_path) else ""
        
        if not self.search_summary:
            search_queries = self.create_search_queries()
            for _, query in search_queries.items():
                search_result = self.run_search_summary(query)
                self.search_summary += f"=Query=:\n{query}\n=Search Result=:\n{search_result}\n================\n"

        return self.search_summary

    def write_report(self, report_type):
        """ Writes the report for the given question.
        Args: None
        Returns: str: The report for the given question
        """
        # yield "Searching online..."

        report_type_func = prompts.get_report_by_type(report_type)
        
        yield from self.call_agent_stream(report_type_func(self.question, self.search_online()))