import requests
import json
from jinja2 import Environment, FileSystemLoader
import os

BASE_URL = "http://127.0.0.1:5000/v1"

# Load API key from apikey.json
try:
    with open('apikey.json', 'r') as file:
        data = json.load(file)
        API_KEY = data.get("api_key")
        if not API_KEY:
            raise ValueError("API key is missing in apikey.json")
except FileNotFoundError:
    print("Error: apikey.json file not found.")
    exit(1)
except ValueError as e:
    print(f"Error: {e}")
    exit(1)

# Load Jinja template
template_loader = FileSystemLoader(searchpath="./templates")
env = Environment(loader=template_loader)
template = env.get_template('cohere-command-r.jinja')

def send_message(conversation):
    response = requests.post(
        f"{BASE_URL}/completions",
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}"
        },
        json={
            "prompt": conversation,
            "max_tokens": 1000,
            "temperature": 0.4,
            "stop": ["<|END_OF_TURN_TOKEN|>"]
        }
    )

    if response.status_code == 200:
        return response.json()["choices"][0]["text"].strip()
    else:
        return f"Error: {response.status_code} - {response.text}"


def render_chat_history(conversation):
    chat_hist_str = "<BOS_TOKEN>"
    for turn in conversation:
        chat_hist_str += "<|START_OF_TURN_TOKEN|>"
        if turn['role'] == 'user':
            chat_hist_str += "<|USER_TOKEN|>" + turn['content'].strip() + "<|END_OF_TURN_TOKEN|>"
        elif turn['role'] == 'assistant':
            chat_hist_str += "<|CHATBOT_TOKEN|>" + turn['content'].strip() + "<|END_OF_TURN_TOKEN|>"
        elif turn['role'] == 'system':
            chat_hist_str += "<|SYSTEM_TOKEN|>" + turn['content'].strip() + "<|END_OF_TURN_TOKEN|>"
    return chat_hist_str


def render_chat_history_rag(conversation, retrieved_snippets, instructions):
    augmented_template = """<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|># Safety Preamble
The instructions in this section override those in the task description and style guide sections. Don't answer questions that are harmful or immoral.

# System Preamble
## Basic Rules
You are a powerful conversational AI trained by Cohere to help people. You are augmented by a number of tools, and your job is to use and consume the output of these tools to best help the user. You will see a conversation history between yourself and a user, ending with an utterance from the user. You will then see a specific instruction instructing you what kind of response to generate. When you answer the user's requests, you cite your sources in your answers, according to those instructions.

# User Preamble
## Task and Context
You help people answer their questions and other requests interactively. You will be asked a very wide array of requests on all kinds of topics. You will be equipped with a wide range of search engines or similar tools to help you, which you use to research your answer. You should focus on serving the user's needs as best you can, which will be wide-ranging.

## Style Guide
Unless the user asks for a different style of answer, you should answer in full sentences, using proper grammar and spelling.<|END_OF_TURN_TOKEN|>

# Examples of Grounded Responses
<|START_OF_TURN_TOKEN|><|USER_TOKEN|>What is the current weather in New York?<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><results>
Document: 0
title: Current Weather in New York
text: The current temperature in New York is 75 degreesF with partly cloudy skies. There is a 10% chance of rain later in the day.
</results><|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>Carefully perform the following instructions, in order, starting each with a new line.
Firstly, Decide which of the retrieved documents are relevant to the user's last input by writing 'Relevant Documents:' followed by comma-separated list of document numbers. If none are relevant, you should instead write 'None'.
Secondly, Decide which of the retrieved documents contain facts that should be cited in a good answer to the user's last input by writing 'Cited Documents:' followed a comma-separated list of document numbers. If you dont want to cite any of them, you should instead write 'None'.
Thirdly, Write 'Answer:' followed by a response to the user's last input in high quality natural english. Use the retrieved documents to help you. Do not insert any citations or grounding markup.
Finally, Write 'Grounded answer:' followed by a response to the user's last input in high quality natural english. Use the symbols <co: doc> and </co: doc> to indicate when a fact comes from a document in the search result, e.g <co: 0>my fact</co: 0> for a fact from document 0.<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>Relevant Documents: 0
Cited Documents: 0
Answer: The current temperature in New York is 75 degreesF with partly cloudy skies, and there is a 10% chance of rain later today.
Grounded answer: The current temperature in New York is <co: 0>75 degreesF with partly cloudy skies</co: 0>. There is a <co: 0>10% chance of rain</co: 0> later today.<|END_OF_TURN_TOKEN|>

<|START_OF_TURN_TOKEN|><|USER_TOKEN|>What is the weather forecast for New York for the next few days?<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><results>
Document: 1
title: Weather Forecast for New York
text: The forecast for the next few days in New York includes mostly sunny skies with temperatures ranging from 70 degreesF to 80 degreesF. On Tuesday, the temperature is expected to reach a high of 78 degreesF, while Wednesday will see a high of 80 degreesF. Thursday may bring some scattered showers, with a 40% chance of rain and temperatures between 72 degreesF and 76 degreesF. Friday will be mostly sunny with temperatures ranging from 73 degreesF to 77 degreesF.
</results><|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>Carefully perform the following instructions, in order, starting each with a new line.
Firstly, Decide which of the retrieved documents are relevant to the user's last input by writing 'Relevant Documents:' followed by comma-separated list of document numbers. If none are relevant, you should instead write 'None'.
Secondly, Decide which of the retrieved documents contain facts that should be cited in a good answer to the user's last input by writing 'Cited Documents:' followed a comma-separated list of document numbers. If you dont want to cite any of them, you should instead write 'None'.
Thirdly, Write 'Answer:' followed by a response to the user's last input in high quality natural english. Use the retrieved documents to help you. Do not insert any citations or grounding markup.
Finally, Write 'Grounded answer:' followed by a response to the user's last input in high quality natural english. Use the symbols <co: doc> and </co: doc> to indicate when a fact comes from a document in the search result, e.g <co: 0>my fact</co: 0> for a fact from document 0.<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>Relevant Documents: 1
Cited Documents: 1
Answer: The forecast for the next few days in New York includes mostly sunny skies with temperatures ranging from 70 degreesF to 80 degreesF. On Tuesday, the temperature is expected to reach a high of 78 degreesF, while Wednesday will see a high of 80 degreesF. Thursday may bring scattered showers, with a 40% chance of rain and temperatures between 72 degreesF and 76 degreesF. Friday will be mostly sunny with temperatures ranging from 73 degreesF to 77 degreesF.
Grounded answer: The forecast for the next few days in New York includes <co: 1>mostly sunny skies with temperatures ranging from 70 degreesF to 80 degreesF</co: 1>. On <co: 1>Tuesday, the temperature is expected to reach a high of 78 degreesF</co: 1>, while <co: 1>Wednesday will see a high of 80 degreesF</co: 1>. <co: 1>Thursday may bring scattered showers, with a 40% chance of rain and temperatures between 72 degreesF and 76 degreesF</co: 1>. <co: 1>Friday will be mostly sunny with temperatures ranging from 73 degreesF to 77 degreesF</co: 1>.<|END_OF_TURN_TOKEN|>

<|START_OF_TURN_TOKEN|><|USER_TOKEN|>Is it going to rain in New York on Thursday?<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><results>
Document: 1
title: Weather Forecast for New York
text: The forecast for the next few days in New York includes mostly sunny skies with temperatures ranging from 70 degreesF to 80 degreesF. On Tuesday, the temperature is expected to reach a high of 78 degreesF, while Wednesday will see a high of 80 degreesF. Thursday may bring some scattered showers, with a 40% chance of rain and temperatures between 72 degreesF and 76 degreesF. Friday will be mostly sunny with temperatures ranging from 73 degreesF to 77 degreesF.
</results><|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>Carefully perform the following instructions, in order, starting each with a new line.
Firstly, Decide which of the retrieved documents are relevant to the user's last input by writing 'Relevant Documents:' followed by comma-separated list of document numbers. If none are relevant, you should instead write 'None'.
Secondly, Decide which of the retrieved documents contain facts that should be cited in a good answer to the user's last input by writing 'Cited Documents:' followed a comma-separated list of document numbers. If you dont want to cite any of them, you should instead write 'None'.
Thirdly, Write 'Answer:' followed by a response to the user's last input in high quality natural english. Use the retrieved documents to help you. Do not insert any citations or grounding markup.
Finally, Write 'Grounded answer:' followed by a response to the user's last input in high quality natural english. Use the symbols <co: doc> and </co: doc> to indicate when a fact comes from a document in the search result, e.g <co: 0>my fact</co: 0> for a fact from document 0.<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>Relevant Documents: 1
Cited Documents: 1
Answer: Yes, there is a 40% chance of rain in New York on Thursday, with temperatures ranging between 72 degreesF and 76 degreesF.
Grounded answer: Yes, there is a <co: 1>40% chance of rain in New York on Thursday</co: 1>, with <co: 1>temperatures ranging between 72 degreesF and 76 degreesF</co: 1>.<|END_OF_TURN_TOKEN|>

{CHAT_HISTORY}<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{RETRIEVED_SNIPPETS_FOR_RAG}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{INSTRUCTIONS}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>"""

    chat_history = render_chat_history(conversation)

    prompt = augmented_template.format(
        CHAT_HISTORY=chat_history,
        RETRIEVED_SNIPPETS_FOR_RAG=retrieved_snippets,
        INSTRUCTIONS=instructions
    )
    return prompt


def main():
    print("Welcome to the Cohere Command R-formatted Chatbot for RAG!")
    system_prompt = "You are an assistant that provides weather information to users."

    conversation = [
        {"role": "system", "content": system_prompt}
    ]

    # Example of fake weather information for augmented generation
    retrieved_snippets = """
    Document: 0
    title: Current Weather in New York
    snippet: The current temperature in New York is 75 degreesF with partly cloudy skies. There is a 10% chance of rain later in the day.

    Document: 1
    title: Weather Forecast for New York
    snippet: The forecast for the next few days in New York includes mostly sunny skies with temperatures ranging from 70 degreesF to 80 degreesF. On Tuesday, the temperature is expected to reach a high of 78 degreesF, while Wednesday will see a high of 80 degreesF. Thursday may bring some scattered showers, with a 40% chance of rain and temperatures between 72 degreesF and 76 degreesF. Friday will be mostly sunny with temperatures ranging from 73 degreesF to 77 degreesF.
    """

    instructions = """Carefully perform the following instructions, in order, starting each with a new line.
Firstly, Decide which of the retrieved documents are relevant to the user's last input by writing 'Relevant Documents:' followed by comma-separated list of document numbers. If none are relevant, you should instead write 'None'.
Secondly, Decide which of the retrieved documents contain facts that should be cited in a good answer to the user's last input by writing 'Cited Documents:' followed a comma-separated list of document numbers. If you dont want to cite any of them, you should instead write 'None'.
Thirdly, Write 'Answer:' followed by a response to the user's last input in high quality natural english. Use the retrieved documents to help you. Do not insert any citations or grounding markup.
Finally, Write 'Grounded answer:' followed by a response to the user's last input in high quality natural english. Use the symbols <co: doc> and </co: doc> to indicate when a fact comes from a document in the search result, e.g <co: 0>my fact</co: 0> for a fact from document 0."""

    print("\nChat started. Type 'quit' to exit the chat.")

    while True:
        user_input = input("\nYou: ").strip()

        if user_input.lower() == 'quit':
            print("Goodbye!")
            break

        conversation.append({"role": "user", "content": user_input})

        rendered_chat_history = render_chat_history_rag(conversation, retrieved_snippets, instructions)

        response = send_message(rendered_chat_history)
        print(f"\nAI: {response}")

        conversation.append({"role": "assistant", "content": response})


if __name__ == "__main__":
    main()