import gradio as gr
from io import BytesIO
import fitz
import tempfile
import openai


class TranslationAgent:
    def __init__(self, openai_key):
        self.memory = []
        system_msg = "You are a translator from english to chinese (简体中文).\n" \
                     " The only thing you do is to translate.\n" \
                     " You don't write anything other then the translation of the text you get.\n" \
                     " The user will only provide the text without asking anything, but what he wants is the translation.\n" \
                     " Never return the translation of a previously translated part!\n " \
                     " The text you will need to translate will often include nonesense stuff because it is coming from a text extraction of a pdf file including images and table.\n" \
                     " Do your best to translate also this messy parts."

        self.memory.append({"role": "system", "content": system_msg})

        openai.api_key = openai_key

    def fade_memory(self):
        if len(self.memory) >= 5:
            del self.memory[1:3]

    def translate_chunk(self, chunk):
        self.memory.append({"role": "user", "content": chunk})
        response = openai.ChatCompletion.create(
            model="gpt-4-turbo-preview",
            messages=self.memory
        )
        reply = response["choices"][0]["message"]["content"]
        self.memory.append({"role": "assistant", "content": reply})
        self.fade_memory()
        return reply


def extract_text_from_pdf(pdf, start, stop):
    text = ""
    with fitz.open(stream=BytesIO(pdf), filetype='pdf') as doc:  # remove .read()
        for i, page in enumerate(doc):
            if start <= i:
                if i <= stop:
                    text += page.get_text()
                else:
                    break
    return text


def split_text(text, chunk_size=100):
    words = text.split()
    chunks = []
    current_chunk_words = []

    for word in words:
        current_chunk_words.append(word)
        if word.endswith('.') and len(current_chunk_words) >= chunk_size:
            chunks.append(' '.join(current_chunk_words))
            current_chunk_words = []

    # add the last chunk if any words remain
    if current_chunk_words:
        chunks.append(' '.join(current_chunk_words))

    return chunks


def translate_pdf(openai_key, pdf, start, stop):
    translator = TranslationAgent(openai_key)
    translated_text = ""
    error_message = "Translation Successful"

    try:
        # extract text
        if pdf is not None:
            text = extract_text_from_pdf(pdf, start=start, stop=stop)
            chunks = split_text(text)

            translated_chunks = []
            for chunk in chunks:
                translated_chunk = translator.translate_chunk(chunk)
                translated_chunks.append(translated_chunk + " ")

            translated_text = ' '.join(translated_chunks)
    except Exception as e:
        error_message = f"Translation Failed: {e}"

    # Create a temporary file with a specific prefix
    temp = tempfile.NamedTemporaryFile(delete=False, prefix="translatedPDF_", suffix=".txt")

    # Write to the temporary file
    with open(temp.name, 'w', encoding='utf-8') as f:
        f.write(translated_text)

    return translated_text, error_message, temp.name


iface = gr.Interface(
    fn=translate_pdf,
    inputs=[
        gr.Textbox(lines=1, label="OpenAI API key",
                          placeholder="Enter your OpenAI API key here"),
        gr.File(type="binary", label="PDF file", ),
        gr.Number(label="Starting Page", ),
        gr.Number(label="Final Page")
    ],
    outputs=["text", "text", gr.File(label="Translated Text File")],
    title="Pdf Translator: English ==> Chinese",
)

iface.launch()