import os
import base64
import glob

import random
import string
import json

import requests
import openai
import time

from pdf2image import convert_from_path
from PIL import Image
from pathlib import Path

PACKAGE_ROOT = str(Path(__package__).absolute())
from aaw.globals import APIs, STRINGS

# openai>=1 removed the openai.error module. Keep compatibility with older code paths.
try:
    # Legacy (<1.0)
    from openai.error import InvalidRequestError  # type: ignore
except Exception:
    # Newer openai versions expose BadRequestError instead.
    InvalidRequestError = getattr(openai, "BadRequestError", Exception)


def create_dataset():
    # Create folder for raw images
    raw_dir = "raw_data"
    p_sub = Path(raw_dir)
    p_sub.mkdir(parents=True, exist_ok=True)

    # Create folder for temporary data
    tmp_dir = "tmp"
    p_tmp = Path(tmp_dir)
    p_tmp.mkdir(parents=True, exist_ok=True)


def get_random_string(length) -> str:
    # choose from all lowercase letter
    letters = string.ascii_lowercase
    result_str = ''.join(random.choice(letters) for i in range(length))
    return result_str


def image_to_text(image):
    # Step 1: Save uploaded file
    filename = image.name
    print(f'uploaded: {filename}')

    image_type = filename.split(".")[-1]
    print("Image type: " + image_type)
    bytes_data = image.getvalue()

    path = "raw_data/" + filename

    # Save bytes
    with open(path, 'wb') as f:
        f.write(bytes_data)

    # Step 2: Convert file to jpg-format if needed
    print("Convert input to JPG format...")
    random_str = get_random_string(5)
    temp_files = []

    if image_type == "pdf":
        # Convert pdf to image
        images = convert_from_path(path)
        for i, img in enumerate(images):
            temp_path = f"tmp/{random_str}_{i}.jpg"
            img.save(temp_path, 'JPEG')
            temp_files.append(temp_path)
            img.close()
    else:
        image_pillow = Image.open(path)
        image_pillow = image_pillow.convert('RGB')
        temp_path = f"tmp/{random_str}_0.jpg"
        image_pillow.save(temp_path, 'JPEG')
        temp_files.append(temp_path)
        image_pillow.close()

    # Step 3: Call GPT-4o API for OCR
    print("Calling GPT-4o API...")

    openai.api_key = APIs["openai"]

    ocr_text = ""

    for temp_file in temp_files:
        with open(temp_file, 'rb') as f:
            file_data = f.read()
            file_data_base64 = base64.b64encode(file_data).decode('utf-8')

        data = {
            "model": "gpt-4o",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            #"text": 'You are a transcription and correction assistant. Transcribe the text from the given images as accurately as possible, preserving the original line breaks. Use HTML <span style="color:blue;">unreadable character</span> for any characters that are not clearly readable. Ensure contextual accuracy in the transcription. Additionally, use the XML tags <grammarerror> for grammatical errors, <spellingerror> for spelling errors, and <punctuationerror> for punctuation errors.'
                            "text": 'You are a transcription and correction assistant. Transcribe the text from the given images as accurately as possible, preserving the original line breaks. Ensure contextual accuracy in the transcription.'
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{file_data_base64}"
                            }
                        }
                    ]
                }
            ],
            "max_tokens": 2048
        }

        response = openai.ChatCompletion.create(
            model=data["model"],
            messages=data["messages"],
            max_tokens=data["max_tokens"]
        )

        response_text = response['choices'][0]['message']['content']
        ocr_text += response_text

    # Cleanup temporary files
    for temp_file in temp_files:
        os.remove(temp_file)

    os.remove(path)

    if not ocr_text:
        print("No text found!!")
        return "", "Kein Text erkannt"

    # Step 4: Use GPT to improve the text (if needed)
    # No additional improvement needed if the OCR already provides cleaned text

    return ocr_text.strip(), ""

def ocr(image_name: str, num_requests=1) -> str:
    output_text = []

    # Very important that Mathpix does not save the data
    options_json = json.dumps({"metadata": {"improve_mathpix": False}})

    for i in range(num_requests):
        r = requests.post("https://api.mathpix.com/v3/text",
                          files={"file": open("tmp/" + image_name + str(i) + '.jpg', "rb")},
                          data={"options_json": options_json},
                          headers={
                              "app_id": APIs["ocr_app_id"],
                              "app_key": APIs["ocr_app_key"],
                          },
                          )
        try:
            output_text.append(r.json()["text"])
        except KeyError:
            print("------ There was an error processing the input image... ----------")

    total_output = " ".join(output_text)
    return total_output

def store_prompt(prompt, filename="stored_prompts.txt"):
    with open(filename, "a") as f:
        f.write(prompt + "\n\n")

def run_gpt(messages: list, task: str, essay: str, engine="text-davinci-003", max_tokens=1000, error_tmp=None):
    print("using engine " + engine)

    openai.api_key = APIs["openai"]

    exception = None
    # create a completion
    for i in range(10):
        try:
            if engine.startswith("text-davinci"):
                return run_gpt3(messages=messages, task=task, essay=essay, engine=engine, max_tokens=max_tokens), None
            elif engine.startswith("gpt-"):
                return run_chatgpt(messages=messages, task=task, essay=essay, engine=engine, max_tokens=max_tokens), None
            else:
                print("Unknown engine " + engine + "!")
                return "", None
        except InvalidRequestError as ire:
            # If this happens directly stop trying and return
            print("####  Invalid Request!  ####")
            print(ire)
            return "", ire
        except Exception as e:
            # For all other exceptions, try multiple times
            print("#####" + str(e) + "#############")
            time.sleep(5)
            exception = e

    return "", exception

def run_gpt3(messages, task, essay, engine="text-davinci-003", max_tokens=1000):
    # Verarbeite die Textbereiche
    processed_task = task.strip()
    processed_essay = essay.strip()
    
    # Füge die Textbereiche den Nachrichten hinzu
    if processed_task:
        messages.append({"role": "user", "content": processed_task})
    if processed_essay:
        messages.append({"role": "user", "content": processed_essay})
    
    # Erstelle den vollständigen Prompt für die Speicherung
    prompt = ""
    for message in messages:
        prompt += message["role"] + ": " + message["content"] + "\n"
    
    # Speichere den Prompt lokal
    store_prompt(prompt)
    
    # Sende die Anfrage an OpenAI
    completion = openai.Completion.create(engine=engine, prompt=prompt, max_tokens=max_tokens)
    
    # Gib das Ergebnis zurück
    return completion.choices[0].text

def run_chatgpt(messages: list, task: str, essay: str, engine="chatgpt-4o", max_tokens=1000):
    # Verarbeite die Textbereiche
    processed_task = task.strip()
    processed_essay = essay.strip()

    # Füge die Textbereiche den Nachrichten hinzu
    if processed_task:
        messages.append({"role": "user", "content": processed_task})
    if processed_essay:
        messages.append({"role": "user", "content": processed_essay})

    # Speichere die Nachrichten lokal
    store_prompt(json.dumps(messages, indent=4))

    # Sende die Anfrage an OpenAI
    completion = openai.ChatCompletion.create(
        model=engine,
        messages=messages,
        max_tokens=max_tokens
    )
    return completion['choices'][0]['message']['content']