import os
import streamlit as st

from dotenv import load_dotenv
from engines import get_engine

import api

# Load environment variables from .env file before importing any other modules
load_dotenv()


def main():
    st.set_page_config(
        page_title='Note Taker',
        page_icon='🎙️',
        layout='centered',
        initial_sidebar_state='expanded',
    )

    title = '🎙️ Meetings Note Taker 🎙️'
    st.title(title)
    st.write(
        'Upload an audio file, transcribe it using WhisperX or AssemblyAI, and generate meeting notes using your selected model.'
    )

    openai_api_key = os.environ.get('OPENAI_API_KEY') or st.text_input(
        'Enter your OpenAI API key:', type='password'
    )

    engine_type = os.environ.get('TRANSCRIPTION_ENGINE') or st.selectbox(
        'Select a transcription engine:', ['WhisperX', 'AssemblyAI']
    )
    if engine_type == 'WhisperX':
        device = os.environ.get('PYTORCH_DEVICE') or st.selectbox(
            'Select a PyTorch device:', ['cuda', 'cpu']
        )
        compute_type = os.environ.get('PYTORCH_COMPUTE_TYPE') or st.selectbox(
            'Select a compute type:', ['int8', 'float16']
        )
        batch_size = os.environ.get('PYTORCH_BATCH_SIZE') or st.selectbox(
            'Select a batch size:', [4, 8, 16, 32, 64]
        )
        whisper_model = os.environ.get('WHISPER_MODEL') or st.selectbox(
            'Select a Whisper model:', ['large-v2', 'base']
        )
    else:
        device = None
        compute_type = None
        batch_size = None
        whisper_model = None

    engine_api_key = os.environ.get(
        f'{engine_type.upper()}_API_KEY'
    ) or st.text_input(f'Enter your {engine_type} API key:', type='password')
    openai_model = os.environ.get('OPENAI_MODEL') or st.selectbox(
        'Select a model:', ['gpt-3.5-turbo-16k', 'gpt-4-0613']
    )

    uploaded_audio = st.file_uploader(
        'Upload an audio file',
        type=['aac', 'm4a', 'mp3', 'webm', 'mp4', 'mpga', 'wav', 'mpeg'],
        accept_multiple_files=False,
    )
    language = os.environ.get('AUDIO_LANGUAGE') or st.selectbox(
        'Language code of the audio:', ['en', 'es']
    )
    prompt = os.environ.get('SUMMARIZE_PROMPT') or st.text_area(
        'Enter a custom prompt for the summary:', 'Summarize the following audio transcription with a list of the key points with the speakers in the original language'
    )

    if st.button('Generate Notes'):
        if uploaded_audio:
            if openai_api_key:
                st.markdown('Transcribing the audio...')
                engine = get_engine(
                    engine_type,
                    api_key=engine_api_key,
                    device=device,
                    compute_type=compute_type,
                    batch_size=batch_size,
                    whisper_model=whisper_model,
                )
                transcription = api.transcribe(engine, language, uploaded_audio)

                st.markdown(
                    f'###  Transcription:\n\n<details><summary>Click to view</summary><p><pre><code>{transcription}</code></pre></p></details>',
                    unsafe_allow_html=True,
                )

                st.markdown('Summarizing the transcription...')

                summary = api.summarize_transcript(
                    openai_api_key=openai_api_key,
                    transcript=transcription,
                    openai_model=openai_model,
                    prompt=prompt,
                )

                st.markdown(f'### Summary:')
                st.write(summary)
            else:
                st.error('We need valid OpenAI API keys')


if __name__ == '__main__':
    main()