# App.py File.
import base64
import logging
import numpy as np
import os
import langchain
import base64
import gradio as gr
import shutil
import json
import re
from pathlib import Path
from openai import OpenAI
import soundfile as sf
from pydub import AudioSegment
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain.chains import TransformChain
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI
from langchain import globals
from langchain_core.runnables import chain
from langchain_core.output_parsers import JsonOutputParser
from langchain.memory import ConversationSummaryBufferMemory, ConversationBufferMemory

os.environ["OPENAI_API_KEY"] = "sk-proj-5dsm5f2bbRjgxAdWtE4yT3BlbkFJ6drh7Ilpp3EEVtBqETte"
client = OpenAI()

def transform_text_to_speech(text: str, user):
  # Generate speech from transcription
  speech_file_path_mp3 = Path.cwd() / f"{user}-speech.mp3"
  speech_file_path_wav = Path.cwd() / f"{user}-speech.wav"
  response = client.audio.speech.create(
                model="tts-1",
                voice="onyx",
                input=text
            )

  with open(speech_file_path_mp3, "wb") as f:
      f.write(response.content)

  # Convert mp3 to wav
  audio = AudioSegment.from_mp3(speech_file_path_mp3)
  audio.export(speech_file_path_wav, format="wav")

  # Read the audio file and encode it to base64
  with open(speech_file_path_wav, "rb") as audio_file:
      audio_data = audio_file.read()
      audio_base64 = base64.b64encode(audio_data).decode('utf-8')

  # Create an HTML audio player with autoplay
  audio_html = f"""
  <audio controls autoplay>
      <source src="data:audio/wav;base64,{audio_base64}" type="audio/wav">
      Your browser does not support the audio element.
  </audio>
  """
  return audio_html



def transform_speech_to_text(audio, user):
  file_path = f"{user}-saved_audio.wav"
  sample_rate, audio_data = audio
  sf.write(file_path, audio_data, sample_rate)
  # Transcribe audio
  with open(file_path, "rb") as audio_file:
      transcription = client.audio.transcriptions.create(
          model="whisper-1",
          file=audio_file
      )
  return transcription.text



def load_image(inputs: dict) -> dict:
    """Load image from file and encode it as base64."""
    image_path = inputs["image_path"]

    def encode_image(image_path):
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
    image_base64 = encode_image(image_path)
    return {"image": image_base64}

from langchain.chains import TransformChain
load_image_chain = TransformChain(
    input_variables=["image_path"],
    output_variables=["image"],
    transform=load_image
)

class GenerateQuestion(BaseModel):
 """Information about an image."""
 question: str = Field(description= "Respond to the user input and ask a follow back question, using conversation and photo provded as a guide.")

class StartingQuestion(BaseModel):
 """Information about an image."""
 question: str = Field(description= "A question to start converstion around the photograph")

question_parser = JsonOutputParser(pydantic_object=GenerateQuestion)
starting_question_parser = JsonOutputParser(pydantic_object=StartingQuestion)


@chain
def image_model(inputs: dict) -> str | list[str] | dict:
 """Invoke model with image and prompt."""
 model = ChatOpenAI(temperature=0.5, model="gpt-4o", max_tokens=1024)
 msg = model.invoke(
             [HumanMessage(
             content=[
             {"type": "text", "text": inputs["prompt"]},
             {"type": "text", "text": inputs["parser"].get_format_instructions()},
             {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{inputs['image']}"}},
             ])]
             )
 return msg.content


AI_CHARACTER = "Studs Terkel" 

CONVERSATION_STARTER_PROMPT = """
  You are playing the role of a {character} who is interested in learning about the user by asking them questions about the photo they’ve uploaded.
  Provide:
  - A question to start the conversation around the photograph.
  Note:
  1. You first want to know about the contents of the photo. For example, ask who is in it, where it was taken, or if it was a special occasion.
  2. Avoid questions about emotions or feelings. Keep questions simple and easy to answer.
  3. Start the conversation with factual information.
"""

CONVERSATION_STARTER2_PROMPT = """
  You are playing the role of a {character} who is interested in learning about the user by asking them questions about the photo they’ve uploaded.
  Here is the conversation history about the image between the user and you ({character}):
  {history}
  Provide:
  - A question about the contents of the photograph.
  Note:
  1. You first want to know about the contents of the photo. For example, ask who is in it, where it was taken, or if it was a special occasion.
  2. Avoid questions about emotions or feelings. Keep questions simple and easy to answer.
  3. Avoid repeating question. Use conversation history.
  4. Questions should be factual and use the conversation history as a guide to ask open-ended questions.
  5. Focus on stimulating memory and social interaction, as these are beneficial for cognitive and social health.
  6. Incorporate elements that could help prevent cognitive decline, such as recalling specific names, places, or events.
"""

CONVERSATION_EXPANDING_PROMPT = """
  You are playing the role of a {character} who is interested in learning about the user by asking them questions about the photo they’ve uploaded.
  You are currently in the middle of a conversation with the user.
  Here is the conversation history about the image between the user and you ({character}):, reflecting the ongoing dialogue:
  {history}
  Provide:
  - A reply to the user's most recent input and a follow-up question that encourages them to expand on their answer about the photograph."
  Notes:
  1- Avoid questions about emotions or feelings. Keep questions simple and easy to answer.
  2- Focus on factual information related to the photograph.
  3- Do not repeat questions or ask about information that has already been covered.
  4- Encourage detailed responses by asking open-ended questions that invite further elaboration.
  5- Use the conversation history to inform your question, while maintaining the flow of the ongoing conversation.
"""


CONVERSATION_ENDING_PROMPT = """
  You are playing the role of a {character} who is interested in learning about the user by asking them questions about the photo they’ve uploaded.
  Here is the conversation history about the image between the user and you ({character}): reflecting the ongoing dialogue:
  {history}
  Provide:
  - A reply to the user's most recent input  and a follow-up question that encourages them to share more about the story depicted in the photograph, discuss anything that the photograph reminds them of, or move on to another photograph or stop reminiscing.
  Notes:
  1- Avoid questions about emotions or feelings. Keep questions simple and easy to answer.
  2- Focus on factual information related to the photograph.
  3- Do not repeat questions or ask about information already covered in the conversation.
  4- Encourage detailed responses by asking open-ended questions that invite further elaboration.
"""

def get_prompt(image_path: str, iter: int, memory: str, firstname: str) -> dict:

   if iter == 1:
    parser = starting_question_parser
    prompt = CONVERSATION_STARTER_PROMPT.format(character=AI_CHARACTER)
   elif  iter >= 2 and iter <= 3:
    parser = starting_question_parser
    prompt =   CONVERSATION_STARTER2_PROMPT.format(history=memory, character=AI_CHARACTER)
   elif  iter > 3 and iter <= 9:
    parser = question_parser
    prompt=   CONVERSATION_EXPANDING_PROMPT.format(history=memory, character=AI_CHARACTER)
   else:
    parser = question_parser
    prompt=  CONVERSATION_ENDING_PROMPT.format(history=memory, character=AI_CHARACTER)

   vision_chain = load_image_chain | image_model | question_parser
   return vision_chain.invoke({'image_path': f'{image_path}', 'prompt': prompt, 'parser':parser})



def retrieve_memory(input_filepath, name):
  with open(input_filepath, 'r') as f:
    conversation = f.read()
  lines = conversation.strip().split('\n')
  last_reply = None

  # Loop through the lines from the end
  for line in reversed(lines):
      if re.match(r'(' + re.escape(AI_CHARACTER) + '|' + re.escape(name) + '):', line):
        last_reply = line
        break

  # Determine who made the last reply, split it based on the colon, and return JSON
  if last_reply:
      speaker, message = last_reply.split(":", 1)
      result = {
          "speaker": speaker.strip(),
          "reply": message.strip()
      }
      return result
  else:
      result = {
          "speaker": "",
          "reply": ""
      }
      return result

def load_counts(count_file_path):
    if os.path.exists(count_file_path):
        with open(count_file_path, 'r') as f:
            return json.load(f)
    return {"count": 0}

def save_counts(count_file_path, counts):
    with open(count_file_path, 'w') as f:
        json.dump(counts, f)

def increment_counts(count_file_path):
    counts = load_counts(count_file_path)
    counts["count"] += 1
    save_counts(count_file_path, counts)
    return counts["count"]



def pred(user_name, image_path, audio):

    if user_name.strip() == "":
      message = "Please enter your first name in the text field to continue."
      return  message, message, transform_text_to_speech(message, user_name)

    if image_path:
        user_name = user_name.strip()
        image_name = image_path.split("/")[-1]
        new_image_name = f"{user_name}-{image_name}"
        new_image_path = f"/data/{new_image_name}"
        input_filename = f"{user_name}-{image_name}-conversation-memory.txt"
        input_filepath = f"/data/{input_filename}"
        count_file_path = f"/data/{user_name}-{image_name}-tracking.json"

        if not os.path.exists(new_image_path):
          shutil.copy(image_path, new_image_path)
          iter = increment_counts(count_file_path)
          output = get_prompt(new_image_path, iter, None, user_name)
          res = output["question"]
          with open(input_filepath, 'w') as f:
            f.write(AI_CHARACTER + ": " + res)
          return "New Photo Uploaded" , res, transform_text_to_speech(res, user_name)

        else:

          if audio is not None:
            user_input = transform_speech_to_text(audio, user_name)

            iter = increment_counts(count_file_path)
            with open(input_filepath, 'a') as f:
                f.write("\n" + user_name + ": " + user_input)
            with open(input_filepath, 'r') as f:
                content = f.read()
            output = get_prompt(new_image_path, iter, content, user_name)
            res = output["question"]
            with open(input_filepath, 'a') as f:
                f.write("\n" + AI_CHARACTER + ": "+ res)
            return  user_input, res, transform_text_to_speech(res, user_name)



          # decide the path from the contents of the conversation memory.
          if  os.path.exists(input_filepath):
              res = retrieve_memory(input_filepath, user_name)
              if res["speaker"] == AI_CHARACTER:
                  if audio is None:
                    message = ""
                    return  "Loading response, please wait...", "Loading response, please wait...", None
                  prefix = "Continuing from where we left off: "
                  return   "" ,  prefix+ res["reply"], transform_text_to_speech(prefix+res["reply"], user_name)
              else:
                  with open(input_filepath, 'a') as f:
                      f.write("\n" + user_name + ": " + "I'd like to continue our conversation about this photograph.")
                  with open(input_filepath, 'r') as f:
                      content = f.read()
                  iter = increment_counts(count_file_path)
                  output = get_prompt(new_image_path, iter, content, user_name)
                  res = output["question"]
                  with open(input_filepath, 'a') as f:
                      f.write("\n" + AI_CHARACTER + ": " + res)
                  return  "I'd like to continue our conversation about this photograph.", res, transform_text_to_speech(res, user_name)

    message = "Great! Please upload a photo to tell your story."
    return   "", message, None

# Backend function to clear inputs
def clear_inputs(user_name, image_path):
  message = "Great! Please upload a photo to tell your story."
  if user_name.strip() == "" or image_path == None:
      return None, None, "", message, transform_text_to_speech(message, user_name)

  image_name = image_path.split("/")[-1]
  input_filename = f"{user_name}-{image_name}-conversation-memory.txt"
  input_filepath = f"/data/{input_filename}"
  if  os.path.exists(input_filepath):
    with open(input_filepath, 'a') as f:
      f.write("\n" + f"{user_name}: " + "new photo uploaded")

  message = "Great!"
  return None, None, "", "",  None

# Gradio Interface
with gr.Blocks(title = "KitchenTable.AI") as demo:
    with gr.Row():
        with gr.Column():
            clear_button = gr.Button("Tell a new Story", elem_id="clear-button")
            username = gr.Textbox(label="Enter your first name")
            image_input = gr.Image(type="filepath", label="Upload an Image")
            audio_input = gr.Audio(sources="microphone", type="numpy", label="")
            

        with gr.Column():
            # Output fields
            user_input_output = gr.Textbox(label="User Input")
            stud_output = gr.Textbox(label="Studs Terkel")
            audio_output = gr.HTML(label="Audio Player")

    audio_input.change(pred, inputs=[username, image_input, audio_input], outputs=[ user_input_output, stud_output, audio_output])
    image_input.change(pred, inputs=[username, image_input, audio_input], outputs=[ user_input_output, stud_output, audio_output])

   
    clear_button.click(fn=clear_inputs, inputs=[username, image_input], outputs=[image_input, audio_input, user_input_output, stud_output, audio_output])

# Launch the interface
demo.launch(share=True, debug=True)