from openai import OpenAI
from pathlib import Path
import streamlit as st
import PyPDF2
import os

st.set_page_config(page_title="ReadPDF")

if "output_file_name" not in st.session_state:
    st.session_state.output_file_name = ""

if "file_name" not in st.session_state:
    st.session_state.file_name = ""

if "audio_processed" not in st.session_state:
    st.session_state.audio_processed = False

def extract_text(pdf_path, from_page=10, to_page=10):
    try:
        with open(pdf_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader (file)
            total_pages = len(pdf_reader.pages)

            # Ensure from_page and to_page are within valid range
            from_page = max(1, min(from_page, total_pages))
            to_page = max(from_page, min(to_page, total_pages))

            # Limit the number of pages to extract to a maximum of 10
            to_page = min(from_page + 9, to_page)

            st.session_state.output_file_name = f"{st.session_state.file_name[:-4]} P{from_page}-P{to_page}.mp3"

            text = ""

            for page_num in range(from_page - 1, to_page):
                page = pdf_reader.pages[page_num]
                text += page.extract_text()

            return text
    except Exception as e:
        print(f"Error: {e}")
        return 0

def process(voice):

    with st.spinner("Uploading..."):

        client = OpenAI(api_key=st.session_state.api_key)

        extracted_text = extract_text(Path("Data", st.session_state.file_name), st.session_state.from_page, st.session_state.to_page)
        if extracted_text == 0:
            st.error('Page numbers out of bounds!', icon="🚨")

            return "Error!"
        
        print("THIS WILL NOT GET PRINTED!")
        text_prompt = f"{extracted_text}\n\n Make corrections to this text, such as puncuation, typing errors, missing letters, and formatting."
        text_response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "You are a text correction assistant."},
                {"role": "user", "content": text_prompt},
            ]
        )

        output = text_response.choices[0].message.content

        response = client.audio.speech.create(
            model="tts-1",
            voice=voice.lower(),
            input=output,
        )
        response.stream_to_file(st.session_state.output_file_name)
        st.session_state.audio_processed = True

    st.success('Processed!')

def upload_docs () :

    path_check = os.path.isdir("Data")
    if not path_check:
        os.makedirs("Data")

    with st.spinner("Uploading..."):
        
        doc = st.session_state.uploaded_file
        with open(Path("Data", doc.name),"wb") as f:
            f.write(doc.getbuffer())

        st.session_state.file_name = doc.name

    st.success('Done!')
    
def create_upload () :  
    uploaded_doc = st.file_uploader("Upload Files", type=["pdf"], accept_multiple_files=False, key="uploaded_file",
                        help="Upload documents to embed!", on_change=upload_docs, label_visibility="visible")

    return upload_docs

st.title("ReadPDF")

file_uploaded = not len(st.session_state.file_name) > 1

upload_button = create_upload()

st.text(f"Current Document: {st.session_state.file_name}")

st.text_input('API Key', key="api_key")

from_page = st.number_input("From Page: ", value=10, min_value=1, max_value=300, step=1, key="from_page")
to_page = st.number_input("To Page: ", value=10, min_value=1, max_value=300, step=1, key="to_page")

voice = st.selectbox(
   "Choose TTS Voice",
   ("Alloy", "Echo", "Fable", "Onyx", "Nova", "Shimmer"),
    key="voice"
)

st.markdown("[You can listen to the different voice options here!](https://platform.openai.com/docs/guides/text-to-speech/voice-options)")
st.button("Process", type="secondary", on_click=process, disabled=file_uploaded, kwargs={"voice": st.session_state.voice})

if (st.session_state.audio_processed):
    
    with open(st.session_state.output_file_name, "rb") as file:
        btn = st.download_button(
                label="Download Audio",
                data=file,
                file_name=st.session_state.output_file_name,
                mime="audio/mpeg"
            )