# Constants
DOCUMENT_TITLE_PROMPT = """
INSTRUCTIONS
What is the title of the following document?

Your response MUST be the title of the document, and nothing else. DO NOT respond with anything else.

{document_title_guidance}

{truncation_message}

DOCUMENT
{document_text}
""".strip()

TRUNCATION_MESSAGE = """
Also note that the document text provided below is just the first ~{num_words} words of the document. That should be plenty for this task. Your response should still pertain to the entire document, not just the text provided below.
""".strip()

MAX_CONTENT_TOKENS = 4000
MODEL_NAME = "gpt-4o-mini"
TOKEN_ENCODER = tiktoken.encoding_for_model('gpt-3.5-turbo')

def make_llm_call(chat_messages: list[dict]) -> str:
    """
    Make an API call to the OpenAI language model.

    Args:
        chat_messages (list[dict]): A list of message dictionaries for the chat completion.

    Returns:
        str: The generated response from the language model.
    """
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    response = client.chat.completions.create(
        model=MODEL_NAME,
        messages=chat_messages,
        max_tokens=MAX_CONTENT_TOKENS,
        temperature=0.2,
    )
    return response.choices[0].message.content.strip()

def truncate_content(content: str, max_tokens: int) -> tuple[str, int]:
    """
    Truncate the content to a specified maximum number of tokens.

    Args:
        content (str): The input text to be truncated.
        max_tokens (int): The maximum number of tokens to keep.

    Returns:
        tuple[str, int]: A tuple containing the truncated content and the number of tokens.
    """
    tokens = TOKEN_ENCODER.encode(content, disallowed_special=())
    truncated_tokens = tokens[:max_tokens]
    return TOKEN_ENCODER.decode(truncated_tokens), min(len(tokens), max_tokens)

def get_document_title(document_text: str, document_title_guidance: str = "") -> str:
    """
    Extract the title of a document using a language model.

    Args:
        document_text (str): The text of the document.
        document_title_guidance (str, optional): Additional guidance for title extraction. Defaults to "".

    Returns:
        str: The extracted document title.
    """
    # Truncate the content if it's too long
    document_text, num_tokens = truncate_content(document_text, MAX_CONTENT_TOKENS)
    truncation_message = TRUNCATION_MESSAGE.format(num_words=3000) if num_tokens >= MAX_CONTENT_TOKENS else ""

    # Prepare the prompt for title extraction
    prompt = DOCUMENT_TITLE_PROMPT.format(
        document_title_guidance=document_title_guidance,
        document_text=document_text,
        truncation_message=truncation_message
    )
    chat_messages = [{"role": "user", "content": prompt}]
    
    return make_llm_call(chat_messages)

# Example usage
if __name__ == "__main__":
    # Assuming document_text is defined elsewhere
    document_title = get_document_title(document_text)
    print(f"Document Title: {document_title}")

def rerank_documents(query: str, chunks: List[str]) -> List[float]:
    """
    Use Cohere Rerank API to rerank the search results.

    Args:
        query (str): The search query.
        chunks (List[str]): List of document chunks to be reranked.

    Returns:
        List[float]: List of similarity scores for each chunk, in the original order.
    """
    MODEL = "rerank-english-v3.0"
    client = cohere.Client(api_key=os.environ["CO_API_KEY"])

    reranked_results = client.rerank(model=MODEL, query=query, documents=chunks)
    results = reranked_results.results
    reranked_indices = [result.index for result in results]
    reranked_similarity_scores = [result.relevance_score for result in results]
    
    # Convert back to order of original documents
    similarity_scores = [0] * len(chunks)
    for i, index in enumerate(reranked_indices):
        similarity_scores[index] = reranked_similarity_scores[i]

    return similarity_scores

def compare_chunk_similarities(chunk_index: int, chunks: List[str], document_title: str, query: str) -> None:
    """
    Compare similarity scores for a chunk with and without a contextual header.

    Args:
        chunk_index (int): Index of the chunk to inspect.
        chunks (List[str]): List of all document chunks.
        document_title (str): Title of the document.
        query (str): The search query to use for comparison.

    Prints:
        Chunk header, chunk text, query, and similarity scores with and without the header.
    """
    chunk_text = chunks[chunk_index]
    chunk_wo_header = chunk_text
    chunk_w_header = f"Document Title: {document_title}\n\n{chunk_text}"

    similarity_scores = rerank_documents(query, [chunk_wo_header, chunk_w_header])

    print(f"\nChunk header:\nDocument Title: {document_title}")
    print(f"\nChunk text:\n{chunk_text}")
    print(f"\nQuery: {query}")
    print(f"\nSimilarity without contextual chunk header: {similarity_scores[0]:.4f}")
    print(f"Similarity with contextual chunk header: {similarity_scores[1]:.4f}")

# Notebook cell for execution
# Assuming chunks and document_title are defined in previous cells
CHUNK_INDEX_TO_INSPECT = 86
QUERY = "Nike climate change impact"

compare_chunk_similarities(CHUNK_INDEX_TO_INSPECT, chunks, document_title, QUERY)