import os
from openai import AsyncOpenAI  # importing openai for API usage
import chainlit as cl  # importing chainlit for our app
from chainlit.prompt import Prompt, PromptMessage  # importing prompt tools
from chainlit.playground.providers import ChatOpenAI  # importing ChatOpenAI tools
from dotenv import load_dotenv
import asyncio
import textract  # importing textract for document processing

# Load OpenAI API key from environment variables
api_key = os.getenv("OPENAI_API_KEY")

# Templates for ChatOpenAI interaction
system_template = """You are a helpful assistant who always speaks in a pleasant tone and answers based on the provided document!
"""

user_template = """{input}
Use the document content to respond to the user query step by step.
"""

# Function to extract text from uploaded documents
async def extract_text_from_file(file):
    # Extract text from the uploaded file using textract
    return textract.process(file).decode('utf-8')

# Function to start the chat session
@cl.on_chat_start  
async def start_chat():
    settings = {
        "model": "gpt-3.5-turbo",
        "temperature": 0,
        "max_tokens": 500,
        "top_p": 1,
        "frequency_penalty": 0,
        "presence_penalty": 0,
    }

    cl.user_session.set("settings", settings)

    # Welcome message with file upload option
    await cl.Message(content="Welcome! Please upload a document to begin.").send()

# Function to handle uploaded files
@cl.on_file_upload
async def on_file_upload(files: list):
    # Handle the uploaded files, assuming it's the first file
    if files:
        file = files[0]
        file_content = await extract_text_from_file(file.path)

        # Save document content in session for later use
        cl.user_session.set("document_content", file_content)
        
        # Inform the user that the document was successfully uploaded
        await cl.Message(content=f"Document '{file.name}' uploaded successfully! You can now ask questions based on the document content.").send()

# Function to handle user messages
@cl.on_message  
async def main(message: cl.Message):
    settings = cl.user_session.get("settings")
    document_content = cl.user_session.get("document_content", "")

    if not document_content:
        # If no document is uploaded, prompt the user to upload one
        await cl.Message(content="Please upload a document first.").send()
        return

    client = AsyncOpenAI()

    prompt = Prompt(
        provider=ChatOpenAI.id,
        messages=[
            PromptMessage(
                role="system",
                template=system_template,
                formatted=system_template,
            ),
            PromptMessage(
                role="user",
                template=user_template,
                formatted=user_template.format(input=message.content),
            ),
        ],
        inputs={"input": message.content, "document": document_content},
        settings=settings,
    )

    msg = cl.Message(content="")

    # Call OpenAI
    async for stream_resp in await client.chat.completions.create(
        messages=[m.to_openai() for m in prompt.messages], stream=True, **settings
    ):
        token = stream_resp.choices[0].delta.content
        if not token:
            token = ""
        await msg.stream_token(token)

    # Update the prompt object with the completion
    prompt.completion = msg.content
    msg.prompt = prompt

    # Send and close the message stream
    await msg.send()