from agent import generate_response,chat_history
import chainlit as cl
import fitz  # PyMuPDF for PDF extraction
import pytesseract
from PIL import Image
import io


@cl.author_rename
def rename(orig_author: str):
    rename_dict = {"Chatbot": "Makewell Assitant"}
    return rename_dict.get(orig_author, orig_author)

@cl.on_chat_start
async def start():
    await cl.Message(content=f"""# Welcome to Makewell Assistant!
    Hi There,👋 We're excited to have you on board This Conversational Solution is designed to assist you in creating a detailed tentative health care plan based on your medical prescription records and conversations.
.
    """).send()

# @cl.on_message
# async def main(message: cl.Message):
#   res = generate_response(message.content)
#   if message.content.lower() == "confirm":
#     print("confirm is called")
#     user_data=extract_data(chat_history=chat_history)
#   await cl.Message(res, author="AppointmentBot").send()
@cl.on_message
async def main(message: cl.Message):
    if message.elements:
        texts = []

        images = [file for file in message.elements if "image" in file.mime]

        for image in images:
            image_data = await image.get_bytes()
            image_text = await extract_text_from_image(image_data)
            texts.append(image_text)

        if texts:
            res = "\n".join(texts)
        else:
            res = "No supported files were found in the attachments."
    else:
        res = generate_response(message.content)

    if message.content.lower() == "confirm":
        user_data = extract_data(chat_history=chat_history)

    await cl.Message(content=res, author="Makewell Assistant").send()

async def extract_text_from_pdf(attachment):
    pdf_data = await attachment.get_bytes()
    pdf_text = ""
    with fitz.open(io.BytesIO(pdf_data)) as doc:
        for page in doc:
            pdf_text += page.get_text()
    return pdf_text

async def extract_text_from_image(image_data):
    image = Image.open(io.BytesIO(image_data))
    image_text = pytesseract.image_to_string(image)
    return image_text