import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import pyautogui
import pytesseract

# Set the path to the Tesseract executable
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

def detect_elements(screenshot_path):
    # This function will only be used if UIED is not available
    # Read the image
    img = cv2.imread(screenshot_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # Apply multiple edge detection methods with different parameters
    edges1 = cv2.Canny(gray, 30, 100)
    edges2 = cv2.Canny(gray, 60, 200)
    edges3 = cv2.Canny(gray, 90, 250)
    edges = cv2.bitwise_or(cv2.bitwise_or(edges1, edges2), edges3)

    # Dilate edges to connect nearby elements
    kernel = np.ones((3,3), np.uint8)
    dilated = cv2.dilate(edges, kernel, iterations=2)

    # Find contours with hierarchy
    contours, hierarchy = cv2.findContours(dilated, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    # Filter and process contours
    elements = []
    min_area = 50  # Reduced minimum area
    max_area = img.shape[0] * img.shape[1] * 0.5  # Maximum area (50% of image)

    for i, contour in enumerate(contours):
        area = cv2.contourArea(contour)
        if min_area < area < max_area:
            x, y, w, h = cv2.boundingRect(contour)
            if w >= 5 and h >= 5:  # Allow smaller elements
                elements.append({
                    'id': len(elements) + 1,
                    'bbox': (x, y, w, h),
                    'area': area,
                    'hierarchy': hierarchy[0][i]
                })

    # Sort elements by area (largest first) and limit to top 100
    elements.sort(key=lambda e: e['area'], reverse=True)
    elements = elements[:100]

    # Convert to PIL Image for drawing
    pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(pil_img)
    font = ImageFont.truetype("arialbd.ttf", 28)  # Use a single, larger font size for all labels

    for element in elements:
        x, y, w, h = element['bbox']
        draw.rectangle([x, y, x+w, y+h], outline="red", width=2)
        
        # Determine label position
        if w <= 100 or h <= 40:  # For small elements, place label outside
            label_x = x + w + 5
            label_y = y + h // 2
            anchor = 'lm'  # Left-middle alignment
        else:
            label_x = x + w - 5
            label_y = y + 5
            anchor = 'ra'  # Right-top alignment

        # Draw a semi-transparent background for the label
        label_text = str(element['id'])
        text_bbox = draw.textbbox((label_x, label_y), label_text, font=font, anchor=anchor)
        text_width = text_bbox[2] - text_bbox[0]
        text_height = text_bbox[3] - text_bbox[1]
        
        if anchor == 'lm':
            bg_bbox = (label_x, label_y - text_height // 2, label_x + text_width, label_y + text_height // 2)
        else:
            bg_bbox = (label_x - text_width, label_y, label_x, label_y + text_height)
        
        draw.rectangle(bg_bbox, fill=(255, 255, 255, 180))  # Semi-transparent white background

        # Draw the label text
        draw.text((label_x, label_y), label_text, fill="red", font=font, anchor=anchor)

    # Save the annotated image
    annotated_path = 'annotated_screenshot.png'
    pil_img.save(annotated_path)

    return annotated_path, elements

def find_text_on_screen(screenshot_path, text):
    try:
        # Read the image
        img = cv2.imread(screenshot_path)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        # Perform text detection with different configurations
        configs = [
            '--psm 11',  # Sparse text. Find as much text as possible in no particular order.
            '--psm 6',   # Assume a single uniform block of text.
            '--psm 3'    # Fully automatic page segmentation, but no OSD. (Default)
        ]

        for config in configs:
            data = pytesseract.image_to_data(gray, config=config, output_type=pytesseract.Output.DICT)

            for i, word in enumerate(data['text']):
                if text.lower() in word.lower():
                    x = data['left'][i]
                    y = data['top'][i]
                    w = data['width'][i]
                    h = data['height'][i]
                    return (x + w//2, y + h//2)  # Return center of the word

        # If text not found, try to find partial matches
        for config in configs:
            data = pytesseract.image_to_data(gray, config=config, output_type=pytesseract.Output.DICT)

            for i, word in enumerate(data['text']):
                if any(part.lower() in word.lower() for part in text.lower().split()):
                    x = data['left'][i]
                    y = data['top'][i]
                    w = data['width'][i]
                    h = data['height'][i]
                    return (x + w//2, y + h//2)  # Return center of the word

        return None
    except Exception as e:
        print(f"Error in text detection: {e}")
        return None