# ---------------------------------- web code ----------------------------------

import json
from http.server import BaseHTTPRequestHandler

class handler(BaseHTTPRequestHandler):

    def do_POST(self):

        self.send_response(200)
        self.send_header('Content-type', 'application/json')
        self.end_headers()

        content_length = int(self.headers['Content-Length'])
        post_data = self.rfile.read(content_length)
        data = json.loads(post_data)

        self.wfile.write(chat(data['query'], data['history']).encode('utf-8'))

# ------------------------------- env, constants -------------------------------

from api.get_blocks import get_top_k_blocks, Block

from typing import List, Dict
import openai
import os
import requests
import tiktoken

# OpenAI models
EMBEDDING_MODEL = "text-embedding-ada-002"
COMPLETIONS_MODEL = "gpt-3.5-turbo"
MODERATION_ENDPOINT = "https://api.openai.com/v1/moderations"

# OpenAI parameters
LEN_EMBEDDINGS = 1536
MAX_TOKEN_LEN_PROMPT = 4095 # This may be 8191, unsure.
TRUNCATE_CONTEXT = 2000

# OpenAI API key
OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
openai.api_key = OPENAI_API_KEY

# --------------------------------- prompt code --------------------------------




def limit_tokens(text: str, max_tokens: int, encoding_name: str = "cl100k_base") -> str:
    encoding = tiktoken.get_encoding(encoding_name)
    tokens = encoding.encode(text)[:max_tokens]
    return encoding.decode(tokens)






def construct_prompt(query: str, history: List[Dict[str, str]], context: List[Block]) -> List[Dict[str, str]]:
    """
    
    Args:
        query (str): The user query.
        history (List[Dict[str, str]]): The previous dialogue. Defaults to [].
        blocks (List[Block]): The top-k most relevant blocks from the Alignment Research Dataset. Defaults to [].

    History takes the format: history=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Who won the world series in 2020?"},
        {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."},
        {"role": "user", "content": "Where was it played?"}
        {"role": "assistant", "content": "Los Angeles, California."}
    ]

    Returns: List[Dict[str, str]]: The prompt in messages format.
    """

    # Initialize prompt with system description
    prompt = [{"role": "system", "content": "You are a helpful assistant knowledgeable about AI Alignment and Safety."}]
    # Add previous dialogue
    prompt.extend(history) 

    instruction_prompt = "Please give a clear and coherent answer to my question (written after \"Q:\") using the following sources:"
    prompt.append({"role": "user", "content": instruction_prompt})
    
    # Add context from top-k blocks
    context_prompt = ""
    for i, block in enumerate(context):
        context_prompt += f"[{chr(ord('a') + i)}] {block.title} - {block.author} - {block.date}\n\n{block.text}\n\n"

    context_prompt = context_prompt[:-2] # trim last two newlines

    context_prompt = limit_tokens(context_prompt, TRUNCATE_CONTEXT) # truncate to about 2k tokens

    prompt.append({"role": "user", "content": f"{context_prompt}"})
    
    # Add user query
    question_prompt = "In your answer, please cite any claims you make back to each source using the format: [a], [b], etc."
    question_prompt += "\n\nQ: " + query
    prompt.append({"role": "user", "content": question_prompt})
    
    return prompt




# ------------------------------------------------------------------------------

        
def normal_completion(prompt: List[Dict[str, str]]) -> str:
    """
    This function uses the OpenAI ChatCompletions API to answer a user query.

    Args:
        messages (Dict[str, str]): A dictionary containing the system prompt and user prompt, in addition to any previous dialogue.

    Returns:
        str: The answer generated by the API.
        
    Raises:
        Exception: If the API call fails.
    """
    try:
        return openai.ChatCompletion.create(
            model=COMPLETIONS_MODEL,
            messages=prompt
        )["choices"][0]["message"]["content"]
    except Exception as e:
        print(e)
        return "I'm sorry, I failed to process your query. Please try again. If the problem persists, please contact the administrator."

def chat(query: str, history: List[Dict[str, str]] = [], k: str = 10, HyDE: bool = False) -> str:
    """
    This function uses the OpenAI ChatCompletions API to answer a user query.
    It first checks if the query is offensive, and if so, raises an exception.
    Then, it finds the top-k most relevant blocks from the Alignment Research Dataset and uses them as context for the ChatCompletions API.
    It uses the blocks to generate a prompt for the ChatCompletions API.
    Finally, it uses the ChatCompletions API to generate an answer to the user query.

    Args:
        query (str): The user query.
        previous_dialogue (List[Dict[str, str]]): The previous dialogue. Defaults to [].
        k (str): The number of blocks to use as context.
        mode (str): The mode to use for the ChatCompletions API. Defaults to "standard".
        HyDE (bool): Whether to use the HyDE technique for semantic search. This makes search slower, but better. Defaults to False.
        stream (bool): Whether to stream the results from the ChatCompletions API. Defaults to True.
        stream_delay (float): The delay between each word in the streamed response when streaming a hard-coded response. Defaults to 0.1.

    Returns:
        str: The answer to the user query.
    
    Raises:
        Exception: If the query is offensive.
    """

    

    # 2. Find the top-k most relevant blocks from the Alignment Research Dataset
    top_k_blocks: List[Block] = get_top_k_blocks(query, k, HyDE)
    
    # 3. Generate a prompt for the ChatCompletions API
    prompt: List[Dict[str, str]] = construct_prompt(query, history, top_k_blocks)
    
    # 4. Use the top-k most relevant blocks as context for the ChatCompletions API, and generate an answer to the user query
    completion: str = normal_completion(prompt)
    return completion