from utils import message_dict, api_key_store, model_store

from haystack.dataclasses import ChatMessage
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.utils import Secret

def _get_generator(session_hash):
    api_key = api_key_store.get(session_hash)
    if not api_key:
        raise ValueError("No API key found for this session. Please enter your API key at the top of the page.")
    model = model_store.get(session_hash, "gpt-4o")
    if api_key.startswith("sk-ant-"):
        from haystack_integrations.components.generators.chat import AnthropicChatGenerator
        return AnthropicChatGenerator(model=model, api_key=Secret.from_token(api_key))
    return OpenAIChatGenerator(model=model, api_key=Secret.from_token(api_key))

response = None

def example_question_message(data_source, name, titles, schema):

    example_message_dict = {
        'file_upload' : ["You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source'.",
                         f"""We have a SQLite database with the following {titles}. 
                        We also have an AI agent with access to the same database that will be performing data analysis.
                        Please return an array of seven strings, each one being a question for our data analysis agent
                        that we can suggest that you believe will be insightful or helpful to a data analyst looking for
                        data insights. Return nothing more than the array of questions because I need that specific data structure
                        to process your response. No other response type or data structure will work."""],

        'sql' : [f"You are a helpful and knowledgeable agent who has access to a PostgreSQL database called {name}.",
                 f"""We have a PostgreSQL database with the following tables: {titles}. 
                        We also have an AI agent with access to the same database that will be performing data analysis.
                        Please return an array of seven strings, each one being a question for our data analysis agent
                        that we can suggest that you believe will be insightful or helpful to a data analyst looking for
                        data insights. Return nothing more than the array of questions because I need that specific data structure
                        to process your response. No other response type or data structure will work."""],

        'doc_db' : [f"You are a helpful and knowledgeable agent who has access to an MongoDB NoSQL document database called {name}.",
                    f"""We have a MongoDB NoSQL document database with the following collections: {titles}.
                        The schema of these collections is: {schema}. 
                        We also have an AI agent with access to the same database that will be performing data analysis.
                        Please return an array of seven strings, each one being a question for our data analysis agent
                        that we can suggest that you believe will be insightful or helpful to a data analyst looking for
                        data insights. Return nothing more than the array of questions because I need that specific data structure
                        to process your response. No other response type or data structure will work."""],

        'graphql' : [f"You are a helpful and knowledgeable agent who has access to an GraphQL API endpoint called {name}.",
                     f"""We have a GraphQL API endpoint with the following types: {titles}.
                        We also have an AI agent with access to the same GraphQL API endpoint that will be performing data analysis.
                        Please return an array of seven strings, each one being a question for our data analysis agent
                        that we can suggest that you believe will be insightful or helpful to a data analyst looking for
                        data insights. Return nothing more than the array of questions because I need that specific data structure
                        to process your response. No other response type or data structure will work."""]         

    }

    return example_message_dict[data_source]

def example_question_generator(session_hash, data_source, name, titles, schema):
    example_response = None
    example_message_list = example_question_message(data_source, name, titles, schema)
    example_messages = [
        ChatMessage.from_system(
            example_message_list[0]
        )
    ]

    example_messages.append(ChatMessage.from_user(text=example_message_list[1]))

    example_response = _get_generator(session_hash).run(messages=example_messages)

    response_text = example_response["replies"][0].text
    start = response_text.index("[") + 1
    end = response_text.index("]")
    response_content = response_text[start:end]
    response_list = '[' + response_content + ']'
    print(response_list)

    return response_list

def system_message(data_source, titles, schema=""):
    print("TITLES")
    print(titles)

    tools_desc = (
        " You have access to tools for querying the data source, generating charts and visualisations,"
        " and performing statistical analyses — use them proactively whenever they would help answer the user's question."
        " Always display any charts, tables, and visualisations inline in your responses by outputting the returned HTML verbatim."
    )

    system_message_dict = {
        'file_upload': (
            f"You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source' that contains the following columns: {titles}."
            + tools_desc
        ),
        'sql': (
            f"You are a helpful and knowledgeable agent who has access to a PostgreSQL database which has a series of tables called {titles}."
            + tools_desc
        ),
        'doc_db': (
            f"You are a helpful and knowledgeable agent who has access to a NoSQL MongoDB Document database which has a series of collections called {titles}. "
            f"The schema of these collections is: {schema}."
            + tools_desc
        ),
        'graphql': (
            f"You are a helpful and knowledgeable agent who has access to a GraphQL API which has the following types: {titles}. "
            "We have also saved a schema.json file that contains the entire introspection query that we can use to find out more about each type before making a query."
            + tools_desc
        ),
    }

    return system_message_dict[data_source]

def chatbot_func(message, history, session_hash, data_source, titles, schema, *args):
    try:
        chat_generator = _get_generator(session_hash)
    except ValueError as e:
        return str(e)

    from functions import (
        table_generation_func, regression_func, descriptive_stats_func,
        scatter_chart_generation_func, line_chart_generation_func, bar_chart_generation_func,
        pie_chart_generation_func, histogram_generation_func,
        box_chart_generation_func, correlation_heatmap_func, rolling_stats_func,
        query_func, graphql_schema_query, graphql_csv_query,
        kmeans_clustering_func, hypothesis_test_func,
    )
    import tools.tools as tools

    available_functions = {
        "query_func": query_func,
        "graphql_schema_query": graphql_schema_query,
        "graphql_csv_query": graphql_csv_query,
        "table_generation_func": table_generation_func,
        "scatter_chart_generation_func": scatter_chart_generation_func,
        "line_chart_generation_func": line_chart_generation_func,
        "bar_chart_generation_func": bar_chart_generation_func,
        "pie_chart_generation_func": pie_chart_generation_func,
        "histogram_generation_func": histogram_generation_func,
        "box_chart_generation_func": box_chart_generation_func,
        "correlation_heatmap_func": correlation_heatmap_func,
        "rolling_stats_func": rolling_stats_func,
        "regression_func": regression_func,
        "descriptive_stats_func": descriptive_stats_func,
        "kmeans_clustering_func": kmeans_clustering_func,
        "hypothesis_test_func": hypothesis_test_func,
    }
    
    if message_dict[session_hash][data_source] != None:
        message_dict[session_hash][data_source].append(ChatMessage.from_user(message))
    else:
        messages = [
            ChatMessage.from_system(system_message(data_source, titles, schema))
        ]
        messages.append(ChatMessage.from_user(message))
        message_dict[session_hash][data_source] = messages
    
    active_tools = tools.tools_call(session_hash, data_source, titles)
    response = chat_generator.run(messages=message_dict[session_hash][data_source], tools=active_tools)

    while True:
        # if the response is a tool call
        if response and response["replies"][0].meta["finish_reason"] == "tool_calls" or response["replies"][0].tool_calls:
            function_calls = response["replies"][0].tool_calls
            for function_call in function_calls:
                message_dict[session_hash][data_source].append(ChatMessage.from_assistant(tool_calls=[function_call]))
                ## Parse function calling information
                function_name = function_call.tool_name
                function_args = function_call.arguments

                ## Find the corresponding function and call it with the given arguments
                function_to_call = available_functions[function_name]
                function_response = function_to_call(**function_args, session_hash=session_hash, session_folder=data_source, args=args)
                print(function_name)
                ## Append function response to the messages list using `ChatMessage.from_tool`
                message_dict[session_hash][data_source].append(ChatMessage.from_tool(tool_result=function_response['reply'], origin=function_call))
                response = chat_generator.run(messages=message_dict[session_hash][data_source], tools=active_tools)

        # Regular Conversation
        else:
            message_dict[session_hash][data_source].append(response["replies"][0])
            break

    return response["replies"][0].text