from openai import AzureOpenAI
from agents.PromptAgent import PromptAgent
import os
import random
import json
import io
import fitz
import pytesseract
from PIL import Image


class Expert(PromptAgent):
    def __init__(self, model="chatgpt") -> None:
        # List of strings. Each element (string) is a message alternating between solver and expert
        super(Expert, self).__init__(model)

        self.prompt = self.prompts["EXPERT"]
        self.init_input = self.manual
        self.image_data = None

    def set_module(self, module):
        super().set_module(module)
        # self.init_input = self.manual
        # self.init_input = self.read_pdf("config\KeepTalkingAndNobodyExplodes-BombDefusalManual-v1.pdf")
        # self.init_input = self.extract_text_from_pdf("config\KeepTalkingAndNobodyExplodes-BombDefusalManual-v1.pdf")
        # self.init_input, self.image_data = self.extract_text_and_images_from_pdf("config\KeepTalkingAndNobodyExplodes-BombDefusalManual-v1.pdf")
        self.init_input, self.image_data = self.extract_text_and_images_from_pdf(
            "config/manual.pdf")

    def read_pdf(self, file_path):
        doc = fitz.open(file_path)
        text = ""
        for page_num in range(len(doc)):
            page = doc.load_page(page_num)
            text += page.get_text()
        return text

    def extract_text_from_pdf(self, pdf_path):
        # open the pdf file
        doc = fitz.open(pdf_path)

        full_text = ""

        for page_num in range(len(doc)):
            page = doc.load_page(page_num)
            text = page.get_text()  # extract text from the page
            full_text += text

            # extract images from the page
            image_list = page.get_images(full=True)
            for img_index, img in enumerate(image_list):
                xref = img[0]
                base_image = doc.extract_image(xref)
                image_bytes = base_image["image"]
                image = Image.open(io.BytesIO(image_bytes))

                # apply OCR to the image
                ocr_text = pytesseract.image_to_string(image)
                full_text += ocr_text

        return full_text

    def extract_text_and_images_from_pdf(self, pdf_path):
        # Open the PDF file
        doc = fitz.open(pdf_path)

        # Initialize text and images content
        text_content = ""
        images = []

        # Iterate through each page
        for page_num in range(doc.page_count):
            page = doc.load_page(page_num)

            # Extract text
            text_content += page.get_text()

            # Extract images
            image_list = page.get_images(full=True)
            for img_index, img in enumerate(image_list):
                xref = img[0]
                base_image = doc.extract_image(xref)
                image_bytes = base_image["image"]
                image = Image.open(io.BytesIO(image_bytes))
                images.append(image)

        return text_content, images

    def get_llm_input(self, image_data, actions, message=None):
        return self.get_conversation_history_string(image_data=self.image_data, actions=None, message=message)

    def append_conversation(self, llm_input):
        speaker = "SOLVER"

        for mess in self.conversation:
            if speaker == "SOLVER":
                llm_input += f"SOLVER: {mess}\n\n"
                speaker = "EXPERT"
            elif speaker == "EXPERT":
                llm_input += f"EXPERT: {mess}\n\n"
                speaker = "SOLVER"

        llm_input += f"{speaker}: "

        return llm_input
