# This file will be used to store the interim conversational state of the doctor and patient and then store the various diagnostic results and the final diagnosis and next best action.
# To store the data we will use Azure blob store and each file will be stored as a json file with the conversation id as the file name.

# Steps
# 1. Create a class to store the conversation state
# 2. Create a function to check if this is a new conversation or an existing conversation
# 3. Create a function to count interactions in the conversation if existing exists
# 4. Create a function to store the conversation in Azure blob storage
# 5. Create a function to retrieve the conversation from Azure blob storage
# 6. Create a function to update the conversation in Azure blob storage
# 7. Create a function to delete the conversation from Azure blob storage

# Import the required libraries
import os
import json
from azure.storage.blob import BlobServiceClient, BlobClient, ContainerClient

# Azure connection string
connect_str = "DefaultEndpointsProtocol=https;AccountName=chatlogs;AccountKey=7tiHxmMbbdEp3/yrydWOEoAc7PPDkVHXV5QXZWR1jTH0hX8/O2gIrSHwMNqXwEnEwWWAD8pqIwuV+AStWqHVNA==;EndpointSuffix=core.windows.net"
container_name = "vme25"

# Create a class to store the conversation state
class Conversation:
    def __init__(self, conversation_id, patient, conversation):
        self.conversation_id = conversation_id
        self.patient = patient
        self.conversation = conversation

# Create a function to check if this is a new conversation or an existing conversation
def check_conversation(conversation_id):
    # Create a blob service client
    blob_service_client = BlobServiceClient.from_connection_string(connect_str)
    container_client = blob_service_client.get_container_client(container_name)
    blob_client = container_client.get_blob_client(conversation_id)
    if blob_client.exists():
        return True
    else:
        return False

# Create a function to count interactions in the conversation if existing exists
def count_interactions(conversation_id):
    # Create a blob service client
    blob_service_client = BlobServiceClient.from_connection_string(connect_str)
    container_client = blob_service_client.get_container_client(container_name)
    blob_client = container_client.get_blob_client(conversation_id)
    
    # Retrieve the existing conversation
    conversation_json = blob_client.download_blob().readall()
    conversation_obj = json.loads(conversation_json)
    
    # Count the number of user interactions
    user_interactions = sum(1 for message in conversation_obj['conversation'] if message['role'] == 'user')
    
    return user_interactions

# Create a function to store the conversation in Azure blob storage
def initail_conversation(conversation_id, patient, from_user, response_from_ai):
    # Create a blob service client
    blob_service_client = BlobServiceClient.from_connection_string(connect_str)
    container_client = blob_service_client.get_container_client(container_name)
    blob_client = container_client.get_blob_client(conversation_id)

    # Initialize the conversation with the first user and assistant messages
    conversation = [
        {'role': 'user', 'content': from_user},
        {'role': 'assistant', 'content': response_from_ai}
    ]

    conversation_obj = Conversation(conversation_id, patient, conversation)
    conversation_json = json.dumps(conversation_obj.__dict__)
    blob_client.upload_blob(conversation_json)

# Create a function to retrieve the conversation from Azure blob storage
def retrieve_conversation(conversation_id):
    print("Retrieving conversation")
    # Create a blob service client
    blob_service_client = BlobServiceClient.from_connection_string(connect_str)
    container_client = blob_service_client.get_container_client(container_name)
    blob_client = container_client.get_blob_client(conversation_id)
    conversation_json = blob_client.download_blob().readall()
    conversation_obj = json.loads(conversation_json)
    return conversation_obj

# Create a function to update the conversation in Azure blob storage
def update_conversation(conversation_id, from_user, response_from_ai):
    paitent = "patient"
    print("Updating conversation")
    # Create a blob service client
    blob_service_client = BlobServiceClient.from_connection_string(connect_str)
    container_client = blob_service_client.get_container_client(container_name)
    blob_client = container_client.get_blob_client(conversation_id)
    
    # Retrieve the existing conversation
    conversation_obj = retrieve_conversation(conversation_id)
    
    # Update the conversation with new messages
    conversation_obj['conversation'].append({'role': 'user', 'content': from_user})
    conversation_obj['conversation'].append({'role': 'assistant', 'content': response_from_ai})
    
    # Convert the updated conversation to JSON
    conversation_json = json.dumps(conversation_obj)
    
    # Upload the updated conversation, overwriting the existing blob
    blob_client.upload_blob(conversation_json, overwrite=True)
    return conversation_json

def write_diagnosis(conversation_id, diagnosis):
    # Create a blob service client
    blob_service_client = BlobServiceClient.from_connection_string(connect_str)
    container_client = blob_service_client.get_container_client(container_name)
    blob_client = container_client.get_blob_client(conversation_id)
    
    # Retrieve the existing conversation
    conversation_obj = retrieve_conversation(conversation_id)
    
    # Update the conversation with new messages
    conversation_obj['conversation'].append({'role': 'diagnosis', 'content': diagnosis})
    
    # Convert the updated conversation to JSON
    conversation_json = json.dumps(conversation_obj)
    
    # Upload the updated conversation, overwriting the existing blob
    blob_client.upload_blob(conversation_json, overwrite=True)

# Create a function to delete the conversation from Azure blob storage
def delete_conversation(conversation_id):
    # Create a blob service client
    blob_service_client = BlobServiceClient.from_connection_string(connect_str)
    container_client = blob_service_client.get_container_client(container_name)
    blob_client = container_client.get_blob_client(conversation_id)
    blob_client.delete_blob()

def store_conversation(conversation_id, from_user, response_from_ai):
    patient = "patient"
    # Create a blob service client
    try:
        print("Storing conversation")
        res_check = check_conversation(conversation_id)
        print(res_check)
        if res_check == False:
            print("Conversation doesn't exists, starting new conversation")
            initail_conversation(conversation_id, patient, from_user, response_from_ai)
            print("Conversation stored")
        else:            
            update_conversation(conversation_id, patient, from_user, response_from_ai)
            print("Conversation updated")
    except Exception as e:
        return "FAIL", str(e)

def get_conversation(conversation_id):
    # Create a blob service client
    try:
        res_check = check_conversation(conversation_id)
        if res_check == True:
            print("Conversation exists")
            count = count_interactions(conversation_id)
            conversation_obj = retrieve_conversation(conversation_id)
            return "PASS", count, conversation_obj
        else:
            print("Conversation doesn't exists")
            return "FAIL", 0, "Conversation does not exist"
    except Exception as e:
        print("an error occured")
        return "FAIL", str(e)



# # Test the functions
# conversation_id = "12345623"
# patient = "John Doe"
# from_user = "I do feel a bit sick"
# response_from_ai = "thank you for this information"
# # print(get_conversation(conversation_id))

# store_conversation(conversation_id, patient, from_user, response_from_ai)
#print(check_conversation(conversation_id))
# print(count_interactions(conversation_id))
# conversation_obj = retrieve_conversation(conversation_id)
# print(conversation_obj)

# update_conversation(conversation_id, patient, "What is your name?", "I am an AI assistant.")
# conversation_obj = retrieve_conversation(conversation_id)
# print(conversation_obj)
# delete_conversation(conversation_id)