from PIL import Image
from io import BytesIO
import base64
from openai import OpenAI


def encode_image_in_base64(image: bytes):
    return base64.b64encode(image).decode("utf-8")


def replace_transparent_pixels(image_bytes: bytes):
    """
    Opens a PNG image, and replaces transparent pixels with white pixels.

    Args:
    image_path: The path to the PNG image.

    Returns:
    The path to the modified image.
    """
    try:
        img = Image.open(BytesIO(image_bytes))
        img = img.convert("RGBA")

        pixels = img.getdata()

        new_pixels = []
        for item in pixels:
            if item[3] == 0:
                new_pixels.append((255, 255, 255, 255))
            else:
                new_pixels.append(item)

        img.putdata(new_pixels)

        img_byte_arr = BytesIO()
        img.save(img_byte_arr, format="PNG")
        img_byte_arr = img_byte_arr.getvalue()

        return img_byte_arr

    except FileNotFoundError:
        print(f"Error: The file was not found.")
        return None
    except Exception as e:
        print(f"An error occurred: {e}")
        return None


def process_image_for_gpt(image_bytes: bytes) -> str:
    image_bytes = replace_transparent_pixels(image_bytes)
    base64_image = encode_image_in_base64(image_bytes)
    return base64_image


def format_final_answer(question: str, answer: str) -> str:
    """Always call to format final answer"""
    client = OpenAI()
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": """ You're tasked with correcting/reformatting an answer from an unreliable AI into the expected format as per their instructions.
            <instructions>
            You are a general AI assistant. I will ask you a question. Your answer should only be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
            </instructions>
            For context, here is the question the AI answered
            <question>"""
                + question
                + """</question>"""
                + """
            Now here is their answer. Only reply with the corrected formatting
            """,
            },
            {"role": "user", "content": str(answer)},
        ],
    )
    return response.choices[0].message.content.strip()