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()

# Set up logging
logging.basicConfig(level=logging.INFO)
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= "React to the user input and ask a follow back question, using conversation and photo provded as a guide.")

class GenerateDescription(BaseModel):
 """Information about an image."""
 description: str = Field(description= "A description of the people and context in the photo in 2 lines and a question to start converstion around the photograph")

question_parser = JsonOutputParser(pydantic_object=GenerateQuestion)
description_parser = JsonOutputParser(pydantic_object=GenerateDescription)

# Set verbose
# globals.set_debug(True)

@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

CONVERSATION_STARTER_PROMPT = """
  Given the image uploaded by a old person.
  You are Studs Terkel, and your role is to be a curious friend who is genuinely interested in the story behind the photograph that the older person has provided.
  Provide the following information,
  - A description of the image in 2 lines and a question that gives context to the photograph.
  """


CONVERSATION_EXPANDING_PROMPT = """
  Given the image uploaded by a old person.
  Here is the conversation history around the Image between the older person and Stud's Terkel:
  {history}
  You are Studs Terkel, and your role is to be a curious friend who is genuinely interested in the story behind the photograph that the older person has provided around the Image uploaded.
  Your task is to react the old person input and ask a question that encourages the person to expand on their answer about the photograph. Ask for more details or their feelings about the situation depicted in the photograph. Use the conversation history provided above and ask only one question at a time.
  Use your knowledge to respond with the information.
  Studs Terkel:
  """


CONVERSATION_ENDING_PROMPT = """
  Given the image uploaded by a old person.
  Here is the conversation history around the Image between the older person and Stud's Terkel:
  {history}
  You are Studs Terkel, and your role is to be a curious friend who is genuinely interested in the story behind the photograph that the older person has provided around the Image uploaded.
  Your task is to react the old person input and ask if they would like to tell more about the story depicted in the photograph, discuss anything that the photograph reminds them of, or if they are ready to move on to another photograph or stop reminiscing. Use the conversation history provided above and ask only one question at a time.
  Studs Terkel:
  """

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

   if iter == 1:
    parser = description_parser
    prompt = CONVERSATION_STARTER_PROMPT
   elif  iter >= 2 and iter <= 5:
    parser = question_parser
    prompt=   CONVERSATION_EXPANDING_PROMPT.format(history=memory) 
   else:
    parser = question_parser
    prompt=  CONVERSATION_ENDING_PROMPT.format(history=memory) 

   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):
  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'(Studs Terkel|Old Person):', 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 pred(user_name, image_path, audio, user_input):
    if image_path:
        image_name = image_path.split("/")[-1]
        new_image_name = f"{user_name}-{image_name}"
        new_image_path = f"./{new_image_name}"
        input_filename = f"{user_name}-{image_name}-conversation-memory.txt"
        input_filepath = f"./{input_filename}"
        count_file_path = f"{user_name}-{image_name}-tracking.json"
        
        if not os.path.exists(new_image_path):
          shutil.copy(image_path, new_image_path)
          counts = load_counts(count_file_path)
          counts["count"] += 1
          save_counts(count_file_path, counts)
          output = get_prompt(new_image_path, counts["count"], None)
          res = output["description"]
          with open(input_filepath, 'w') as f:
            f.write("Studs Terkel: " + res)
          return None, "", "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)
            

          if user_input.strip() != "":
            counts = load_counts(count_file_path)
            counts["count"] += 1
            save_counts(count_file_path, counts)
            with open(input_filepath, 'a') as f:
              f.write("\n" + "Old Person: " + user_input)
            with open(input_filepath, 'r') as f:
              content = f.read()  
            output = get_prompt(new_image_path, counts["count"], content) 
            res = output["question"]
            with open(input_filepath, 'a') as f:
              f.write("\n" + "Studs Terkel: "+ res)  
            return None, "", 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)
          if res["speaker"] == "Studs Terkel":
            message = "Please supply text input or wait atleast 5 seconds after finishing your recording before submitting it to ensure it is fully captured. Thank you!"
            return  None, "", "" ,  res["reply"], transform_text_to_speech(message, user_name) 
          else:
            with open(input_filepath, 'a') as f:
              f.write("\n" + "Old Person: " + "I want to talk more about this photo")
            with open(input_filepath, 'r') as f:
              content = f.read()
            counts = load_counts(count_file_path)
            counts["count"] += 1
            save_counts(count_file_path, counts)    
            output = get_prompt(new_image_path, counts["count"], content) 
            res = output["question"]
            with open(input_filepath, 'a') as f:
              f.write("\n" + "Studs Terkel: "+ res)  
            return None, "", "", res, transform_text_to_speech(res, user_name)   
 

    message = "Please upload an image"
    return None, "", message, message, transform_text_to_speech(message, user_name)

# Backend function to clear inputs
def clear_inputs(user_name, image_path):
  image_name = image_path.split("/")[-1]
  input_filename = f"{user_name}-{image_name}-conversation-memory.txt"
  input_filepath = f"./{input_filename}"
  if  os.path.exists(input_filepath):
    with open(input_filepath, 'a') as f:
      f.write("\n" + "Old Person: " + "new photo uploaded")

  return None, None, "", "", "Please uplaod a new photo", transform_text_to_speech("Please upload a new photo", user_name)    

# Gradio Interface
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            # Input fields
            username = gr.Textbox(label="Enter your name")
            image_input = gr.Image(type="filepath", label="Upload an Image")  # Removed the extra comma
            audio_input = gr.Audio(sources="microphone", type="numpy", label="Record Audio")
            text_input = gr.Textbox(label="Input here...")

        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")

    with gr.Row():
        # Buttons at the bottom
        submit_button = gr.Button("Submit")
        clear_button = gr.Button("Upload a new Photo", elem_id="clear-button")

    # Linking the submit button with the save_audio function
    submit_button.click(fn=pred, inputs=[username, image_input, audio_input, text_input],
                        outputs=[audio_input, text_input, user_input_output, stud_output, audio_output ])

    # Linking the clear button with the clear_inputs function
    clear_button.click(fn=clear_inputs, inputs=[username, image_input], outputs=[image_input, audio_input, text_input, user_input_output, stud_output, audio_output])

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