from fastapi import FastAPI, Request, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse, Response
import requests
import uvicorn
import os
import logging
import torch 
from transformers import pipeline
import subprocess
import re
from piper.voice import PiperVoice
import wave
from pydub import AudioSegment
from gradio_client import Client
import time
import gc
import threading

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

app = FastAPI()
HF_READ_TKN = os.getenv("HF_READ_TKN") 
hope_client = None
Lock = threading.Lock()
TIMESTAMP_LAST_CALL_TO_HOPE = None

#Ongoing llama chat; Dict will tell if we are in chat with this person and what llama input earlier was
ONGOING_CONVO = {}

def is_space_sleeping():
    # Stage can be SLEEPING, PAUSED, RUNNING, or APP_STARTING
    headers = {"Authorization": f"Bearer {HF_READ_TKN}"}
    response = requests.get("https://huggingface.co/api/spaces/HopeLLM/HopeV1.0", headers=headers)
    if response.status_code == 200:
        space_info = response.json()
        status = space_info.get("runtime", {}).get("stage", "") 
        if status == "SLEEPING":
            return True
        elif status == "APP_STARTING": 
            return True 
        elif status == "RESTARTING": 
            return True 
        elif status == "BUILDING": 
            return True 
        elif status == "RUNNING": 
            return False 
        elif status == "PAUSED": 
            return "unknown" 
        else: 
            return False
    return "unknown"  # If API call fails, assume the space is sleeping

def call_hope_api(input_txt, sender_id):
    global hope_client
    global TIMESTAMP_LAST_CALL_TO_HOPE

    TIMESTAMP_LAST_CALL_TO_HOPE = time.time()
    with Lock:
        if hope_client == None:
            if(is_space_sleeping()==True):
                logger.info("Hope is sleeping, trying to wake it up")
                send_whatsapp_message(sender_id, None, "Hope is sleeping, waking it up ;)\nKindly wait for some time", text=True)
                try:
                    Client("HopeLLM/HopeV1.0", hf_token=HF_READ_TKN)
                except:
                    while(is_space_sleeping()):
                        #Sleep for 30sec
                        time.sleep(30)
            elif (is_space_sleeping() == "unknown"):
                logger.error("Something wrong with Hope Check")
            hope_client = Client("HopeLLM/HopeV1.0", hf_token=HF_READ_TKN)
        return hope_client.predict(inp = input_txt, api_name = "/answer_question") 

def reset_hope_client(exit_fn=None):
    #The HOPE GPU Space may sleep; Hence we need to remove this hope_client variable before that
    #If we dont do that, we get runtime error in our WhatsAppServer Space 
    global hope_client
    global TIMESTAMP_LAST_CALL_TO_HOPE
    #If More than an hour to last access to HOPE
    try:
        while True:
            time.sleep(3600)
            with Lock:
                if hope_client and TIMESTAMP_LAST_CALL_TO_HOPE and time.time() - TIMESTAMP_LAST_CALL_TO_HOPE > 3600: 
                    logger.info("Closing Hope Client Connection. An hour since the last iteraction.")
                    hope_client = None
                    gc.collect()

            ##Added cleanup as well here after 40 min sleep time
            logger.debug("Cleaning Ongoing convo")
            for sender_id in ONGOING_CONVO.keys():
                logger.debug(f"Trying to clean up sender {sender_id}")
                if( ONGOING_CONVO[sender_id] and (time.time() - ONGOING_CONVO[sender_id].get("lst_msg_time", 0) > 2400)):
                    send_whatsapp_message(sender_id, None, "Session Time-Out", text=True)
                    clean_up_audio_files(sender_id)
                    ONGOING_CONVO[sender_id] = None
    finally:
        if exit_fn:
            exit_fn()

def notify_parent_of_exit():
    logger.error("The cleaner thread has exited. Check!!")

def clean_up_audio_files(sender_id):
    command = f"rm -f /tmp/*{sender_id}*"
    subprocess.run(command, text = True, shell = True)

# Your WhatsApp verification token
VERIFY_TOKEN = os.getenv("WHATSAPP_VERIFY_TOKEN", "your_verification_token")

# WhatsApp API credentials
WHATSAPP_BASE_URL = "https://graph.facebook.com/v22.0"
WHATSAPP_API_URL = f"{WHATSAPP_BASE_URL}/599521673246967/messages"
WHATSAPP_MEDIA_URL = f"{WHATSAPP_BASE_URL}/599521673246967/media"
WHATSAPP_ACCESS_TOKEN = os.getenv("WHATSAPP_ACCESS_TOKEN", "your_access_token")


# Initialize the Whisper model pipeline
device = "cuda" if torch.cuda.is_available() else "cpu"
transcriber = pipeline(
    task="automatic-speech-recognition",
    model="openai/whisper-small",
    device=device,
    generate_kwargs={"task": "translate", "language":"hi"}
)
transcriber.model.config.forced_decoder_ids = transcriber.tokenizer.get_decoder_prompt_ids(language="hi", task="translate")

amy_voice = PiperVoice.load("./piper/en_US-amy-medium.onnx")

FIXED_PROMPT = "You are a helpful English Teaching Bot. You need to be friendly. Since you teach kids less than 7 years old, be motivating and don't point out any mistakes. Dont have long responses, keep your responses short. Assume I am your students and refer me as your student when replying"
llama_input_template = "<s> [INST] <<SYS>> {fixed_prompt} <</SYS>> Understand the story below and ask me one basic question regarding the story.  Then wait for me to reply, when I reply, appreciate my participation and same way keep asking more questions one by one. Keep the questions one-liners and do not give much context. Keep everything brief \n {story_to_use} [/INST]"
llama_input_gen_random_story = f"<s> [INST] <<SYS>> You are a helpful English Teacher. You teach kids less than 7 years old.<</SYS>> Give me an english story for today's english learning session. Don't make it very big. Just output the story nothing else, no chatting with me [/INST]"

@app.get("/webhook")
async def verify_webhook(request: Request):
    """
    Endpoint for WhatsApp to verify the webhook during setup.
    """
    params = request.query_params
    if (
        params.get("hub.mode") == "subscribe" and
        params.get("hub.verify_token") == VERIFY_TOKEN
    ):
        logger.info("Webhook verified successfully.")
        return Response(content=params.get("hub.challenge"), media_type="text/plain")
    else:
        logger.warning("Webhook verification failed.")
        raise HTTPException(status_code=403, detail="Verification failed")

@app.post("/webhook")
async def receive_message(request: Request, background_tasks: BackgroundTasks):
    """
    Endpoint to receive incoming WhatsApp messages.
    """
    data = await request.json()
    # Extract the message and sender information
    try:
        logger.info(f"Incoming message data: {data}")
        entry = data["entry"][0]
        changes = entry["changes"][0]
        value = changes["value"]
        message = value["messages"][0]
        background_tasks.add_task(backgnd_process_message, message)
        return JSONResponse(status_code=200, content={"message": "Webhook received successfully"})
        
    except (KeyError, IndexError) as e:
        # If the expected fields are not present, return a 400 response
        logger.error(f"Error parsing incoming message: {e}")
        return JSONResponse(status_code=200, content={"message": "Webhook received successfully"})

    return {"status": "success"}


def backgnd_process_message(message):

    sender_id = message["from"]
    message_type = message["type"]
    message_id = message["id"]

    if (ONGOING_CONVO.get(sender_id, None) == None):
        #Send the button to start story
        welcome_prompt = "Hi there!! I am Hope, your english teaching bot :)\nGive me a story or news article and we can chat around it.\n\n_Note: To end this teaching session, send the message_ END"
        ONGOING_CONVO[sender_id] = {'Story': None, 'Llama_conv' : None, 'lst_msg_time': time.time()}
        send_whatsapp_message(sender_id, message_id, welcome_prompt, text=True)
        send_story_option(sender_id)
    else:
        ONGOING_CONVO[sender_id]['lst_msg_time'] = time.time()
        if message_type == "text":
            message_text = message["text"]["body"]
            if(message_text.lower() == "end"):
                ONGOING_CONVO[sender_id] = None
                response_text = "I hope you had a wonderful learning session!! Keep up the good work and best of luck"
                clean_up_audio_files(sender_id)
            elif (ONGOING_CONVO[sender_id]["Story"] == "PENDING"):
                ONGOING_CONVO[sender_id]["Story"] = message_text
                response_text = message_text #The input story. We will repeat the story
                ## Generate question from llama 
                llama_input  = llama_input_template.format(story_to_use = message_text, fixed_prompt=FIXED_PROMPT)
                llama_output = call_hope_api(llama_input,sender_id)
                logger.info(f"LLama Ouput: {llama_output}")
                response_text = response_text + "\n \n" + llama_output
                ONGOING_CONVO[sender_id]['Llama_conv'] = llama_input + llama_output + "</s> <s> [INST] "
            else:
                llama_input = ONGOING_CONVO[sender_id]['Llama_conv'] + message_text + "[/INST]"
                llama_output = call_hope_api(llama_input,sender_id)
                logger.info(f"LLama Ouput: {llama_output}")
                response_text = llama_output
                ONGOING_CONVO[sender_id]['Llama_conv'] = llama_input + llama_output + "</s> <s> [INST] "
        elif message_type == "audio":
            media_id = message["audio"]["id"]
            audio_url = fetch_media_url(media_id)
            audio_path = download_audio(audio_url, sender_id)
            transcription = transcribe_audio(audio_path)
            if (ONGOING_CONVO[sender_id]["Story"] == "PENDING"):
                ONGOING_CONVO[sender_id]["Story"] = transcription
                response_text = transcription
                ## Generate question from llama 
                llama_input  = llama_input_template.format(story_to_use = transcription, fixed_prompt=FIXED_PROMPT)
                llama_output = call_hope_api(llama_input,sender_id)
                logger.info(f"LLama Ouput: {llama_output}")
                response_text = response_text + "\n \n" + llama_output
                ONGOING_CONVO[sender_id]['Llama_conv'] = llama_input + llama_output + "</s> <s> [INST] "
            else:
                llama_input = ONGOING_CONVO[sender_id]['Llama_conv'] + transcription + "[/INST]"
                llama_output = call_hope_api(llama_input,sender_id)
                logger.info(f"LLama Ouput: {llama_output}")
                response_text = llama_output
                ONGOING_CONVO[sender_id]['Llama_conv'] = llama_input + llama_output + "</s> <s> [INST] "
        elif message_type == "interactive":
            if message["interactive"]["button_reply"]["id"] == "random_story":
                ONGOING_CONVO[sender_id]['Story'] = "random_story"
                #send_whatsapp_message(sender_id, message_id, "Perfect, let me find a nice story for today.")
                story_for_chat = call_hope_api(llama_input_gen_random_story,sender_id)
                ONGOING_CONVO[sender_id]['Story'] = story_for_chat
                logger.info(f"Llama gave story {story_for_chat}")

                ## Generate question from llama 
                llama_input  = llama_input_template.format(story_to_use = story_for_chat, fixed_prompt=FIXED_PROMPT)
                llama_output = call_hope_api(llama_input,sender_id)
                logger.info(f"LLama Ouput: {llama_output}")
                response_text = story_for_chat + "\n \n" + llama_output
                ONGOING_CONVO[sender_id]['Llama_conv'] = llama_input + llama_output + "</s> <s> [INST] "
            elif message["interactive"]["button_reply"]["id"] == "own_story":
                response_text = "Great!! Send me your story as the next message, it can be text or audio!!" 
                ONGOING_CONVO[sender_id]['Story'] = "PENDING"
            else:
                button_type = message["interactive"]["button_reply"]["id"]
                logger.warning(f"Unsupported button reply: {button_type}")
                response_text = "Sorry, invalid story type, send a valid type"
        else:
            logger.warning(f"Unsupported message type: {message_type}")
            response_text = "Sorry, I can only process text and audio messages at the moment."

        send_whatsapp_message(sender_id, message_id, response_text) #TO-DO remove audio files every hr or something

def process_text_message(message_text: str, sender_id) -> str:
    """
    Process the incoming message and generate a response.
    Customize this function based on your requirements.
    """
    # For demonstration, echoing back the received message
    logger.info(f"Processing message: {message_text}")
    if ONGOING_CONVO[sender_id]["Story"] == None:
        ONGOING_CONVO[sender_id]["Story"] = message_text
    return f"Received your message: {message_text}"

def fetch_media_url(media_id: str) -> str:
    """
    Fetch the media URL using the media ID.
    """
    url = f"{WHATSAPP_BASE_URL}/{media_id}"
    headers = {
        "Authorization": f"Bearer {WHATSAPP_ACCESS_TOKEN}"
    }
    response = requests.get(url, headers=headers)
    if response.status_code == 200:
        media_url = response.json().get("url")
        logger.info(f"Fetched media URL: {media_url}")
        return media_url
    else:
        logger.error(f"Failed to fetch media URL: {response.status_code} - {response.text}")
        raise HTTPException(status_code=500, detail="Failed to fetch media URL")

def download_audio(audio_url: str, sender_id) -> str:
    """
    Download the audio file from the provided URL.
    """
    headers = {
        "Authorization": f"Bearer {WHATSAPP_ACCESS_TOKEN}"
    }
    response = requests.get(audio_url, headers=headers)
    if response.status_code == 200:
        audio_path = f"/tmp/voice_message{sender_id}.ogg"
        with open(audio_path, "wb") as audio_file:
            audio_file.write(response.content)
        logger.info(f"Audio file downloaded to: {audio_path}")
        return audio_path
    else:
        logger.error(f"Failed to download audio file: {response.status_code} - {response.text}")
        raise HTTPException(status_code=500, detail="Failed to download audio file")

def transcribe_audio(audio_path: str) -> str:
    """
    Transcribe the audio file to text.
    """
    try:
        logger.info(f"Transcribing audio file: {audio_path}")
        transcription = transcriber(audio_path, return_timestamps=True)["text"]
        logger.info(f"Transcription: {transcription}")
        return transcription
    except Exception as e:
        logger.error(f"Failed to transcribe audio: {e}")
        return "I'm sorry, I couldn't understand the audio message."

def convert_text_to_mp3(text, mp3_name = '/tmp/amy_said.mp3'):
    text = re.sub(r':', ". ", text) #To add pause need to add period and not space
    text = re.sub(r'!', ". ", text) #To add pause need to add period and not space
    text = re.sub(r'\n', ". ", text) #To add pause need to add period and not space
    text = re.sub(r"\*[^*]+\*", " ", text)
    text = re.sub(r'[^A-Za-z0-9.,$ \'?]+', " ", text)
    text = text.encode('cp1252', 'ignore').decode('cp1252')
    logger.info(f"Input to amy: {text}")
    if text.isspace() or text == "":
        return None

    with wave.open(mp3_name+'.wav', 'w') as wav_file:
        amy_voice.synthesize(text, wav_file, length_scale=1.3)

    audio = AudioSegment.from_wav(mp3_name+'.wav')
    audio.export(mp3_name, format='mp3')

def get_meta_id(audio_file_path, file_name):
    # Open the audio file in binary mode
    with open(audio_file_path, 'rb') as audio_file:
        # Prepare the data for the request
        files = {
            'file': (file_name, audio_file, 'audio/mpeg')
        }
        payload = {
            'messaging_product': 'whatsapp',
            'type': 'audio/mpeg'
        }
        headers = {
            'Authorization': f'Bearer {WHATSAPP_ACCESS_TOKEN}'
        }

        # Send the POST request to upload media
        response = requests.post(WHATSAPP_MEDIA_URL, headers=headers, data=payload, files=files)

    # Check for successful upload
    if response.status_code == 200:
        media_id = response.json().get('id')
        logger.info(f"Media uploaded successfully. Media ID: {media_id}")
    else:
        media_id = -1
        logger.info(f"Failed to upload media: {response.status_code} {response.text}")

    return media_id

def send_whatsapp_message(recipient_id: str,message_id:str, message_text: str, text=False):
    """
    Send a message back to the WhatsApp user.
    """
    headers = {
        "Authorization": f"Bearer {WHATSAPP_ACCESS_TOKEN}",
        "Content-Type": "application/json"
    }
    if text:
        payload = {
            "messaging_product": "whatsapp",
            "to": recipient_id,
            "type": "text",
            "text": {"body": message_text}
        }
    else:
        payload = {
        'messaging_product': 'whatsapp',
        'to': recipient_id,
        'type': 'audio',
        'audio': {
            'id': None
            }
        }   
        audio_file = convert_text_to_mp3(message_text, "/tmp/" + recipient_id + message_id +".mp3")
        payload['audio']['id'] = get_meta_id("/tmp/" + recipient_id + message_id +".mp3", recipient_id + message_id +".mp3")

    response = requests.post(WHATSAPP_API_URL, json=payload, headers=headers)
    if response.status_code == 200:
        logger.info(f"Message sent successfully to {recipient_id}.")
    else:
        logger.error(f"Failed to send message: {response.status_code} - {response.text}")

def send_story_option(recipient_id):
    headers = {
        "Authorization": f"Bearer {WHATSAPP_ACCESS_TOKEN}",
        "Content-Type": "application/json"
    }
    payload = {
        "messaging_product": "whatsapp",
        "to": recipient_id,
        "type": "interactive",
        "interactive": {
          "type": "button",
          "body": {
            "text": "Story option:"
          },
          "action": {
            "buttons": [
              {
                "type": "reply",
                "reply": {
                  "id": "random_story",
                  "title": "Pick a story"
                }
              },
              {
                "type": "reply",
                "reply": {
                  "id": "own_story",
                  "title": "I have my own story"
                }
              }
            ]
          }
        }
    }
    response = requests.post(WHATSAPP_API_URL, json=payload, headers=headers)
    if response.status_code == 200:
        logger.info(f"Message sent successfully to {recipient_id}.")
    else:
        logger.error(f"Failed to send message: {response.status_code} - {response.text}")



if __name__ == "__main__":
    TIMESTAMP_LAST_CALL_TO_HOPE = time.time()
    threading.Thread(target=reset_hope_client, kwargs={'exit_fn': notify_parent_of_exit},daemon=True).start()
    uvicorn.run(app, host="0.0.0.0", port=7860)