import os
import streamlit as st
from datetime import datetime
import json
import requests
import uuid
from datetime import date, datetime
import requests
from pydantic import BaseModel, Field
from typing import Optional
from 

placeHolderPersona1 = """
##Mission
Please create a highly targeted query for a semantic search engine. The query must represent the conversation to date. 
** You will be given the converstaion to date in the user prompt.
** If no converstaion provided then this is the first converstaion

##Rules
Ensure the query is concise
Do not respond with anything other than the query for the Semantic Search Engine.
Respond with just a plain string """

class ChatRequestClient(BaseModel):

    user_id: str
    user_input: str
    numberOfQuestions: int
    welcomeMessage: str
    llm1: str
    tokens1: int
    temperature1: float
    persona1SystemMessage: str
    persona2SystemMessage: str
    userMessage2: str
    llm2: str
    tokens2: int
    temperature2: float

def call_chat_api(data: ChatRequestClient):
    url = "https://agent-builder-api.greensea-b20be511.northeurope.azurecontainerapps.io/chat/"
    # Validate and convert the data to a dictionary
    validated_data = data.dict()
    
    # Make the POST request to the FastAPI server
    response = requests.post(url, json=validated_data)
    
    if response.status_code == 200:
        return response.json()  # Return the JSON response if successful
    else:
        return "An error occured"  # Return the raw response text if not successful

def genuuid ():
    return uuid.uuid4()

def format_elapsed_time(time):
    # Format the elapsed time to two decimal places
    return "{:.2f}".format(time)

def search_knowledgebase(query)
    return results
    

# Title of the application
# st.image('agentBuilderLogo.png')
st.title('RAG Query Designer')

# Sidebar for inputting personas
st.sidebar.image('cognizant_logo.jpg')
st.sidebar.header("Query Designer")
# st.sidebar.subheader("Welcome Message")
# welcomeMessage = st.sidebar.text_area("Define Intake Persona", value=welcomeMessage, height=300)
st.sidebar.subheader("Query Designer Config")
# numberOfQuestions = st.sidebar.slider("Number of Questions", min_value=0, max_value=10, step=1, value=5, key='persona1_questions')
persona1SystemMessage = st.sidebar.text_area("Query Designer System Message", value=placeHolderPersona1, height=300)

llm1 = st.sidebar.selectbox("Model Selection", ['GPT-4', 'GPT3.5'], key='persona1_size')
temp1 = st.sidebar.slider("Temperature", min_value=0.0, max_value=1.0, step=0.1, value=0.6, key='persona1_temp')
tokens1 = st.sidebar.slider("Tokens", min_value=0, max_value=4000, step=100, value=500, key='persona1_tokens')

# # Persona 2
# st.sidebar.subheader("Recommendation and Next Best Action AI")
# persona2SystemMessage = st.sidebar.text_area("Define Recommendation Persona", value=placeHolderPersona2, height=300)
# with st.sidebar.expander("See explanation"):
#     st.write("This AI persona uses the output of the symptom intake AI as its input. This AI’s job is to augment a health professional by assisting with a diagnosis and possible next best action. The teams will need to determine if this should be a tool used directly by the patient, as an assistant to the health professional or a hybrid of the two. ")
#     st.image("agentPersona2.png")
# llm2 = st.sidebar.selectbox("Model Selection", ['GPT-4', 'GPT3.5'], key='persona2_size')
# temp2 = st.sidebar.slider("Temperature", min_value=0.0, max_value=1.0, step=0.1, value=0.5, key='persona2_temp')
# tokens2 = st.sidebar.slider("Tokens", min_value=0, max_value=4000, step=100, value=500, key='persona2_tokens')
# userMessage2 = st.sidebar.text_area("Define User Message", value="This is the conversation todate, ", height=150)
st.sidebar.caption(f"Session ID: {genuuid()}")


# Main chat interface
st.markdown("""#### Query Translation in RAG Architecture

Query translation in a Retrieval-Augmented Generation (RAG) architecture is the process where an LLM acts as a translator between the user's natural language input and the retrieval system. 

##### Key Functions of Query Translation:
1. **Adds Context**  
   The LLM enriches the user's input with relevant context (e.g., expanding vague questions or specifying details) to make it more precise.
   
2. **Converts to Concise Query**  
   The LLM reformulates the input into a succinct and effective query optimized for the retrieval system's semantic search capabilities.

##### Purpose
This ensures that the retrieval system receives a clear and focused query, increasing the relevance of the information it retrieves. The query translator acts as a bridge between human conversational language and the technical requirements of a semantic retrieval system.""")
# User ID Input
user_id = st.text_input("Experiment ID:", key="user_id")

# Ensure user_id is defined or fallback to a default value
if not user_id:
    st.warning("Please provide an experiment ID to start the chat.")
else:
    # Initialize chat history in session state
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # Display chat messages from history on app rerun
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Collect user input
    if user_input := st.chat_input("Start chat:"):
        # Add user message to the chat history
        st.session_state.messages.append({"role": "user", "content": user_input})
        st.chat_message("user").markdown(user_input)

        # Prepare data for API call
        data = ChatRequestClient(
            user_id=user_id,  # Ensure user_id is passed correctly
            user_input=user_input,
            numberOfQuestions=1000,
            welcomeMessage="",
            llm1=llm1,
            tokens1=tokens1,
            temperature1=temp1,
            persona1SystemMessage=persona1SystemMessage,
            persona2SystemMessage="",
            userMessage2="",
            llm2="GPT3.5",
            tokens2=1000,
            temperature2=0.2
        )

        # Call the API
        response = call_chat_api(data)

        # Process the API response
        agent_message = response.get("content", "No response received from the agent.")
        elapsed_time = response.get("elapsed_time", 0)
        count = response.get("count", 0)

        # Add agent response to the chat history
        st.session_state.messages.append({"role": "assistant", "content": agent_message})
        with st.chat_message("assistant"):
            st.markdown(agent_message)

        # Display additional metadata
        st.caption(f"##### Time taken: {format_elapsed_time(elapsed_time)} seconds")
        # st.caption(f"##### Question Count: {count} of {numberOfQuestions}")