import pytorch_test

import uvicorn
from fastapi import FastAPI, Request

from media_download import YoutubeDownloader
from transcription import StableWhisper
from summarizer import Extract_Summary, AudioBookNarration
from audiobook import AudioBook

app = FastAPI()
output_folder = 'Output'

# Create a context variable to store the contexts for each user
users_context = dict()


@app.get("/get_media_info")
async def get_media_info(request: Request, url: str):

    # Getting User's IP
    user_ip = request.client.host

    # Getting User's Youtube Downloader
    youtube_downloader = YoutubeDownloader(url, output_folder)

    # Getting Youtube Media Info
    media_info = youtube_downloader.get_media_info()
    
    # Storing Info in the context for this user's session
    users_context[user_ip] = dict()
    users_context[user_ip]['downloader'] = youtube_downloader
    # users_context[user_ip]['media_info'] = media_info
    users_context[user_ip]['url'] = url
    
    return media_info


@app.get("/download_media")
async def download_media(request: Request, media_type: str, media_format: str, media_quality: str):

    # Getting User's IP
    user_ip = request.client.host
    
    # Downloading Media for User
    media_path = users_context[user_ip]['downloader'].download(media_type, media_format, media_quality)
    
    # Getting Status
    status = 1 if media_path else 0

    if status:
        # Storing Media Info in the context for this user's session
        users_context[user_ip]['media_path'] = media_path
        users_context[user_ip]['media_type'] = media_type
    
    return {"status": status}


@app.get("/get_transcript")
async def get_transcript(request: Request, subtitle_format: str = 'srt', word_level: bool = True):
    
    # Getting User's IP
    user_ip = request.client.host
    
    # Retrieving the media_path from the context for this user's session
    media_path = users_context[user_ip]['media_path']

    # Checking if the media_type is Video, then extract it's audio
    media_type = users_context[user_ip]['media_type']
    if media_type == 'video':
        media_path = users_context[user_ip]['downloader'].extract_audio(media_path)

    # Whisper based transcription
    stable_whisper_transcript = StableWhisper(media_path, output_folder, subtitle_format=subtitle_format, word_level=word_level)
    transcript = stable_whisper_transcript.generate_transcript()
    transcript_path = stable_whisper_transcript.save_transcript()
    
    # Getting Status
    status = 1 if transcript else 0

    if status:
        # Storing Transcript Info in the context for this user's session
        users_context[user_ip]['transcript'] = transcript
        users_context[user_ip]['transcript_path'] = transcript_path
    
    return {"status": status, "transcript": transcript}


@app.get("/get_summary")
async def get_summary(request: Request, Summary_type: str, Summary_strategy: str, Target_Person_type: str,
                      Response_length:  str, Writing_style: str, text_input: str = None):
    
    # Getting User's IP
    user_ip = request.client.host

    # Getting Transcript if not provided
    if not text_input:
        text_input = users_context[user_ip]['transcript']
    
    # Extracting Summary
    summary_extractor = Extract_Summary(text_input=text_input)
    output = summary_extractor.define_chain(Summary_type=Summary_type,
                                            Summary_strategy=Summary_strategy,
                                            Target_Person_type=Target_Person_type,
                                            Response_length=Response_length,
                                            Writing_style=Writing_style,
                                            key_information=False)
    
    # Getting Status
    status = 1 if output else 0

    if status:
        # Storing Summary Info in the context for this user's session
        users_context[user_ip]['summary'] = output
    
    return {"status": status, "summary": output}


@app.get("/get_key_info")
async def get_key_info(request: Request, Summary_type: str, Summary_strategy: str, Target_Person_type: str, 
                       Response_length: str, Writing_style: str, text_input: str = None):
    
    # Getting User's IP
    user_ip = request.client.host

    # Getting Transcript if not provided
    if not text_input:
        text_input = users_context[user_ip]['transcript']
    
    # Extracting Summary
    summary_extractor = Extract_Summary(text_input=text_input)
    output = summary_extractor.define_chain(Summary_type=Summary_type,
                                            Summary_strategy=Summary_strategy,
                                            Target_Person_type=Target_Person_type,
                                            Response_length=Response_length,
                                            Writing_style=Writing_style,
                                            key_information=True)
    
    # Getting Status
    status = 1 if output else 0

    if status:
        # Storing Key Info in the context for this user's session
        users_context[user_ip]['key_info'] = output
    
    return {"status": status, "key_info": output}


@app.get("/get_narration")
async def get_narration(request: Request, Narration_style: str, text_input: str = None):
    
    # Getting User's IP
    user_ip = request.client.host

    # Getting Transcript if not provided
    if not text_input:
        text_input = users_context[user_ip]['transcript']
    
    # Extracting Narration
    narrator = AudioBookNarration(text_input=text_input)
    output = narrator.define_chain(Narration_style=Narration_style)

    # Getting Status
    status = 1 if output else 0

    if status:
        # Storing Narration Info in the context for this user's session
        users_context[user_ip]['narration'] = output
    
    return {"status": status, "narration": output}


@app.get("/get_audiobook")
async def get_audiobook(request: Request, speaker: str = "male", text_input: str = None):

    # Getting User's IP
    user_ip = request.client.host

    # Getting Transcript if not provided
    if not text_input:
        text_input = users_context[user_ip]['narration']
    
    # Generating Audiobook
    audiobook = AudioBook(output_folder=output_folder)
    audio_path = audiobook.generate_audio_from_text(text_input, speaker=speaker, filename="output_audio")
    
    # Getting Status
    status = 1 if audio_path else 0

    if status:
        # Storing Audiobook path in the context for this user's session
        users_context[user_ip]['audiobook_path'] = audio_path
    
    return {"status": status, "audiobook_path": audio_path}


if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8000)