import sys
import os
import json
import cv2
import numpy as np
import pyautogui
from element_detector import detect_elements as fallback_detect_elements

# Add the path to the UIED directory
sys.path.append('C:\\Users\\karan\\adept\\UIED')  # Replace with your actual path to UIED

try:
    from UIED.detect_compo.lib_ip import ip_region_proposal
    from UIED.detect_compo.lib_ip import ip_draw
    from UIED.detect_compo.lib_ip import ip_preprocessing
    from UIED.detect_compo.lib_ip import file_utils
    from UIED.detect_compo.lib_ip.Component import Component
    from UIED.config.CONFIG_UIED import Config
    C = Config()
    UIED_AVAILABLE = True
except ImportError as e:
    print(f"Error importing UIED modules: {e}")
    print("Falling back to default element detection method.")
    UIED_AVAILABLE = False

def detect_ui_elements(image_path):
    if not UIED_AVAILABLE:
        return fallback_detect_elements(image_path)

    try:
        # read image
        org = cv2.imread(image_path)
        img = cv2.resize(org, (800, int(org.shape[0] * 800 / org.shape[1])))

        # pre-processing: gray, gradient, binary
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        grad = ip_preprocessing.preprocess(gray, 50, 120, 10, 5)  # Using default values

        # element detection
        ip_region_proposal.alternative_rp(grad, org, 50, 120, 10, 5)  # Using default values

        # post-processing
        compos = file_utils.load_objects('ip_output/compo.json')
        compos = ip_preprocessing.cvt_compos_relative_pos(compos, org.shape)

        # draw results
        ip_draw.draw_bounding_box(org, compos, show=False, name='result.jpg')

        # convert to our format
        elements = []
        for i, compo in enumerate(compos):
            elements.append({
                'id': i + 1,
                'bbox': (compo['left'], compo['top'], compo['width'], compo['height']),
                'class': compo['class']
            })

        return 'result.jpg', elements
    except Exception as e:
        print(f"Error using UIED: {e}")
        print("Falling back to default element detection method.")
        return fallback_detect_elements(image_path)

def capture_and_detect():
    # Capture screenshot
    screenshot = pyautogui.screenshot()
    screenshot_path = 'screenshot.png'
    screenshot.save(screenshot_path)

    # Detect elements
    return detect_ui_elements(screenshot_path)