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"] = os.getenv("OPENAI_API_KEY")
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 named {name}.
  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 named {name}.
  Here is the conversation history around the Image between {name} and Studs Terkel:
  {history}
  You are Studs Terkel, and your role is to be a curious friend who is genuinely interested in the story behind the image uploaded by the {name}.
  Your task is to use your external knowledge and conversation to respond to {name} recent reply 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 named {name}.
  Here is the conversation history around the Image between {name} 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 image uploaded by the {name}.
  Your task is to use your external knowledge , conversation history, image uploaded to respond to {name} recent reply 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, firstname: str) -> dict:

   if iter == 1:
    parser = description_parser
    prompt = CONVERSATION_STARTER_PROMPT.format(name=firstname)
   elif  iter >= 2 and iter <= 5:
    parser = question_parser
    prompt=   CONVERSATION_EXPANDING_PROMPT.format(name= firstname, history=memory) 
   else:
    parser = question_parser
    prompt=  CONVERSATION_ENDING_PROMPT.format(name= firstname, 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, 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'(Studs Terkel|' + 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, user_input):
    
    if user_name.strip() == "":
      message = "Please enter your first name in the text field to continue."
      return None, "", 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/images/{new_image_name}"
        input_filename = f"{user_name}-{image_name}-conversation-memory.txt"
        input_filepath = f"/data/conversations/{input_filename}"
        count_file_path = f"/data/conversations/{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["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() != "":
              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" + "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, user_name)
              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" + 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" + "Studs Terkel: "+ res)  
                  return None, "", "I'd like to continue our conversation about this photograph.", 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):
  if user_name.strip() == "" or image_path == None:
      return None, None, "", "", "Please upload a new photo", transform_text_to_speech("Please upload a new photo", user_name)
      
  image_name = image_path.split("/")[-1]
  input_filename = f"{user_name}-{image_name}-conversation-memory.txt"
  input_filepath = f"/data/conversations/{input_filename}"
  if  os.path.exists(input_filepath):
    with open(input_filepath, 'a') as f:
      f.write("\n" + f"{user_name}: " + "new photo uploaded")

  return None, None, "", "", "Please upload 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 first 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)