import os
import json
import asyncio
import logging
from azure.cognitiveservices.vision.computervision import ComputerVisionClient
from azure.cognitiveservices.vision.computervision.models import OperationStatusCodes
from msrest.authentication import CognitiveServicesCredentials
from io import BytesIO
from .util import debug_print, timing_decorator

DEFAULT_OCR_ENGINES = 'azure,tesseract,paddleocr,ios,android'

@timing_decorator
async def perform_ocr(directory: str, filename: str, ocr_engine: str, use_cached: bool = True):
    # Generate cache filename based on image_path and ocr_engine
    CACHE_DIR = '.ocr_cache'
    cache_filename = f"{CACHE_DIR}/{os.path.splitext(filename)[0]}_{ocr_engine}_ocr_cache.json"

    # Try to load from disk before doing perform_x_ocr
    if os.path.exists(cache_filename):
        with open(cache_filename, "r") as cache_file:
            texts_local_ocr = json.load(cache_file)
    else:
        with open(os.path.join(directory, filename), "rb") as image_file:
            file_content = image_file.read()
            if ocr_engine == 'azure':
                texts_local_ocr = await perform_azure_ocr(file_content)
            elif ocr_engine == 'tesseract':
                texts_local_ocr = await perform_tesseract_ocr(file_content)
            elif ocr_engine == 'paddleocr':
                texts_local_ocr = await perform_paddleocr_ocr(file_content)
            elif ocr_engine in ['ios', 'android']:
                print(f'for ocr_engine {ocr_engine}, only cached results are supported')
                return []

        # Cache the result to disk
        if texts_local_ocr != []:
            with open(cache_filename, "w") as cache_file:
                json.dump(texts_local_ocr, cache_file)

    return texts_local_ocr

async def perform_azure_ocr(image_content: bytes):
    # Replace with your Azure Computer Vision subscription key and endpoint
    subscription_key = os.getenv("MS_COMPUTER_VISION_KEY")
    endpoint = os.getenv("MS_COMPUTER_VISION_ENDPOINT")

    # Create an authenticated ComputerVisionClient
    credentials = CognitiveServicesCredentials(subscription_key)
    client = ComputerVisionClient(endpoint, credentials)

    # Read the image content
    image = BytesIO(image_content)

    # Call the API and get the operation location
    read_response = await asyncio.to_thread(client.read_in_stream, image, raw=True)
    operation_location = read_response.headers["Operation-Location"]

    # Get the operation ID
    operation_id = operation_location.split("/")[-1]

    # Wait for the operation to complete
    while True:
        read_result = await asyncio.to_thread(client.get_read_result, operation_id)
        if read_result.status not in [OperationStatusCodes.running, OperationStatusCodes.not_started]:
            break
        await asyncio.sleep(1)

    # Extract and return the text
    if read_result.status == OperationStatusCodes.succeeded:
        texts_with_bounding_boxes = [
            {
                'text': line.text,
                'bounding_box': line.bounding_box
            }
            for page in read_result.analyze_result.read_results
            for line in page.lines
        ]
        return texts_with_bounding_boxes
    else:
        return []

async def perform_tesseract_ocr(image_content: bytes):
    try:
        logging.info("Performing Tesseract OCR")
        import pytesseract
        from PIL import Image
        import numpy as np
        import cv2

        # Had to add this when installing on Azure ("Ubuntu 24.04.1 LTS", but AWS has the same Ubuntu?!)
        # pytesseract.pytesseract.tesseract_cmd =  '/usr/bin/tesseract'

        # Convert bytes to a numpy array
        nparr = np.frombuffer(image_content, np.uint8)
        img_np = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

        # Convert image to grayscale
        gray = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)

        # Use Tesseract to perform OCR with a specified language
        custom_config = r'--oem 3 --psm 6'
        language = 'slv'  # Replace 'eng' with the desired language code, e.g., 'fra' for French
        data = pytesseract.image_to_data(gray, config=custom_config, lang=language, output_type=pytesseract.Output.DICT)

        texts_with_bounding_boxes = []
        for i in range(len(data['text'])):
            if int(data['conf'][i]) > 0:  # Filter out weak confidence results
                text = data['text'][i]
                x, y, w, h = data['left'][i], data['top'][i], data['width'][i], data['height'][i]
                bounding_box = [x, y, x + w, y + h]
                texts_with_bounding_boxes.append({'text': text, 'bounding_box': bounding_box})

        return texts_with_bounding_boxes
    except Exception as e:
        print(f"Error performing Tesseract OCR: {str(e)}")
        logging.info(f"Error performing Tesseract OCR: {str(e)}")
        return []

async def perform_paddleocr_ocr(image_content: bytes):
    try:
        from paddleocr import PaddleOCR
        import numpy as np
        import cv2

        # Convert bytes to a numpy array
        nparr = np.frombuffer(image_content, np.uint8)
        img_np = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

        # Initialize PaddleOCR
        ocr = PaddleOCR(use_angle_cls=True, lang='sl')  # You can change the language as needed

        # Perform OCR on the image
        result = ocr.ocr(img_np, cls=True)

        texts_with_bounding_boxes = []
        for line in result:
            for word_info in line:
                text = word_info[1][0]
                bounding_box = word_info[0]
                texts_with_bounding_boxes.append({'text': text, 'bounding_box': bounding_box})

        return texts_with_bounding_boxes
    except Exception as e:
        print(f"Error performing PaddleOCR: {str(e)}")
        return []

def join_text_snippets(texts: list[dict], strategy: str = 'spaces'):
    try:
        if strategy == 'spaces':
            combined_text = " ".join(t['text'] for t in texts)
        elif strategy == 'smart_newlines':
            # TODO: this shouldn't be hardcoded, it depends on the font size and works only for horizontal text orientation
            PIXEL_OFFSET_RIGHT = 10

            combined_text = ""
            prev = None
            for t in texts:
                # TODO: yuck! bounding box needs to be non-empty and needs to be flattened,
                #   b/c some LLM engines return a list of 4 tuples, some return a list 8 points
                #   we use sum(list, []) to flatten the list (works only for 2-level nested lists) :/
                if prev and len(prev["bounding_box"]) > 0:
                    prev_box = sum(prev["bounding_box"], []) if isinstance(prev["bounding_box"][0], list) else prev["bounding_box"]
                if len(t["bounding_box"]) > 0:
                    t_box = sum(t["bounding_box"], []) if isinstance(t["bounding_box"][0], list) else t["bounding_box"]

                if prev and prev_box and t_box and t_box[0] > prev_box[0] and \
                                        t_box[1] < prev_box[1] + PIXEL_OFFSET_RIGHT and \
                                        t_box[1] > prev_box[1] - PIXEL_OFFSET_RIGHT:
                    # new snippet is to the right of the previous one and vertically approximately in line
                    combined_text += " " + t['text']
                else:
                    combined_text += "\n" + t['text']
                prev = t
        else:
            raise Exception(f"Unsupported strategy: {strategy}")
    except Exception as e:
        print(f"Error joining text snippets: {str(e)}")
        breakpoint()
        return ""

    return combined_text