import os
import gradio as gr
from PIL import Image
import pandas as pd
from ultralytics import YOLO
import easyocr
import numpy as np
from tqdm import tqdm
from datetime import datetime
import zipfile
import tempfile
from pathlib import Path

# Initialize variables
annotations = []
current_index = 0
output_excel = "results.xlsx"

# Directories for saving images
original_dir = "original_images"
yoloed_dir = "yoloed_images"
os.makedirs(original_dir, exist_ok=True)
os.makedirs(yoloed_dir, exist_ok=True)

# Load YOLO model
model_path = "best.pt"
if not os.path.exists(model_path):
    raise FileNotFoundError(f"YOLO model not found at {model_path}")
model = YOLO(model_path)

# Initialize EasyOCR reader
reader = easyocr.Reader(['en'])

# Global counter for cropped images
global_crop_counter = 0

# Image validation
def is_valid_image(image_path):
    try:
        Image.open(image_path)
        return True
    except IOError:
        return False

# YOLO detection
def detect_objects(image_path):
    if not is_valid_image(image_path):
        return []
    try:
        results = model.predict(source=image_path, save=False, conf=0.25)
        detections = []
        for result in results:
            for box in result.boxes.xyxy.cpu().numpy():
                x1, y1, x2, y2 = map(int, box)
                detections.append((x1, y1, x2, y2))
        return detections
    except Exception as e:
        print(f"Error: {e}")
        return []

# Image cropping
def crop_image(image_path, detections):
    global global_crop_counter
    image = Image.open(image_path)
    cropped_images = []
    for bbox in detections:
        x1, y1, x2, y2 = bbox
        cropped_image = image.crop((x1, y1, x2, y2))
        global_crop_counter += 1
        cropped_path = os.path.join(yoloed_dir, f"crop_{global_crop_counter}.jpg")
        cropped_image.save(cropped_path)
        cropped_images.append(cropped_image)
    return cropped_images

# OCR processing
def perform_ocr(image):
    image_np = np.array(image)
    result = reader.readtext(image_np, detail=0)
    numbers = ''.join(filter(str.isdigit, ''.join(result)))
    return numbers.strip()

# GUI update
def update_gui():
    global current_index
    if not annotations:
        return "No images processed.", "", "", "0/0"

    annotation = annotations[current_index]
    original_image = annotation["original_image"]
    thumbnail_size = (600, 600)
    original_thumbnail = original_image.copy().resize(thumbnail_size)

    progress_text = f"Image {current_index + 1}/{len(annotations)}"
    return (
        original_thumbnail,
        annotation.get("room_number", ""),
        annotation.get("meter_value", ""),
        progress_text
    )

# Navigation functions
def save_current_annotation(room_number, meter_value):
    if 0 <= current_index < len(annotations):
        annotations[current_index]["room_number"] = room_number
        annotations[current_index]["meter_value"] = meter_value

def key_handler(key, room_number, meter_value):
    global current_index
    save_current_annotation(room_number, meter_value)

    if key == "-" or key == "[":
        if current_index > 0:
            current_index -= 1
    elif key == "+" or key == "=" or key == "]":
        if current_index < len(annotations) - 1:
            current_index += 1
    else:
        return [gr.update()] * 4  # No change if unsupported key pressed

    return update_gui()

def prev_image(room_number, meter_value):
    global current_index
    save_current_annotation(room_number, meter_value)
    if current_index > 0:
        current_index -= 1
    return update_gui()

def next_image(room_number, meter_value):
    global current_index
    save_current_annotation(room_number, meter_value)
    if current_index < len(annotations) - 1:
        current_index += 1
    return update_gui()

# Export functionality
def export_to_excel(room_number, meter_value):
    global current_index, annotations
    try:
        # Save current edits
        save_current_annotation(room_number, meter_value)

        # Prepare data for Excel
        data = []
        for annotation in annotations:
            rn = annotation.get("room_number", "").strip()
            mv = annotation.get("meter_value", "").strip()
            if rn or mv:
                data.append({"Room Number": rn, "Meter Value": mv})

        df = pd.DataFrame(data)
        temp_dir = tempfile.mkdtemp()
        output_excel = Path(temp_dir) / "results.xlsx"
        df.to_excel(output_excel, index=False)

        # Prepare images for ZIP
        today_date = datetime.now().strftime("%Y-%m-%d")
        images_dir = Path(temp_dir) / f"electricity_meter_images_{today_date}"
        images_dir.mkdir(exist_ok=True)

        for annotation in annotations:
            rn = annotation.get("room_number", "").strip()
            original_image = annotation.get("original_image")
            if rn and original_image:
                new_path = images_dir / f"{rn}.jpg"
                original_image.save(new_path)

        # Create ZIP file
        zip_path = Path(temp_dir) / f"meter_images_{today_date}.zip"
        with zipfile.ZipFile(zip_path, 'w') as zipf:
            for img_file in images_dir.glob("*.jpg"):
                zipf.write(img_file, arcname=img_file.name)

        return str(output_excel), str(zip_path)

    except Exception as e:
        error_dir = tempfile.mkdtemp()
        error_file = Path(error_dir) / "error.txt"
        with open(error_file, 'w') as f:
            f.write(f"Export failed: {str(e)}")
        return str(error_file), str(error_file)

# Image processing
def process_images(uploaded_files):
    global annotations, current_index
    annotations.clear()
    current_index = 0

    if not uploaded_files:
        return None, "", "", "0/0"

    try:
        for temp_file in tqdm(uploaded_files, desc="Processing"):
            image_path = temp_file.name
            original = Image.open(image_path)
            detections = detect_objects(image_path)

            if detections:
                cropped_images = crop_image(image_path, detections)
                for cropped in cropped_images:
                    annotations.append({
                        "image_path": image_path,
                        "original_image": original.copy(),
                        "cropped_image": cropped,
                        "meter_value": perform_ocr(cropped),
                        "room_number": ""
                    })
            else:
                annotations.append({
                    "image_path": image_path,
                    "original_image": original.copy(),
                    "cropped_image": None,
                    "meter_value": "",
                    "room_number": ""
                })

        if annotations:
            current_index = 0
            return update_gui()
        return None, "", "", "0/0"

    except Exception as e:
        return str(e), "", "", "0/0"


# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## Electricity Meter Reader")

    with gr.Row():
        image_input = gr.File(label="Upload Images", file_types=["image"], file_count="multiple")

    with gr.Row():
        original_image_output = gr.Image(label="Original Image")

    with gr.Row():
        room_number_output = gr.Textbox(label="Room Number", interactive=True)
        meter_value_output = gr.Textbox(label="Meter Value", interactive=True)

    with gr.Row():
        progress_label = gr.Textbox(label="Progress", value="0/0", interactive=False)

    with gr.Row():
        prev_button = gr.Button("Previous")
        next_button = gr.Button("Next")
        export_button = gr.Button("Export to Excel & ZIP")

    key_input = gr.Textbox(visible=False, label="Key Handler")

    # Event handlers
    image_input.change(
        fn=process_images,
        inputs=[image_input],
        outputs=[original_image_output, room_number_output, meter_value_output, progress_label]
    )

    prev_button.click(
        fn=prev_image,
        inputs=[room_number_output, meter_value_output],
        outputs=[original_image_output, room_number_output, meter_value_output, progress_label]
    )

    next_button.click(
        fn=next_image,
        inputs=[room_number_output, meter_value_output],
        outputs=[original_image_output, room_number_output, meter_value_output, progress_label]
    )

    export_button.click(
        fn=export_to_excel,
        inputs=[room_number_output, meter_value_output],
        outputs=[
            gr.File(label="Download Excel File"),
            gr.File(label="Download Images ZIP")
        ]
    )

    key_input.submit(
        fn=key_handler,
        inputs=[key_input, room_number_output, meter_value_output],
        outputs=[original_image_output, room_number_output, meter_value_output, progress_label]
    )

    # JavaScript to capture keyboard events
    demo.load(
        js="""
        () => {
            document.addEventListener('keydown', function(e) {
                const allowedKeys = ['-', '=', '[', ']'];
                if (allowedKeys.includes(e.key)) {
                    e.preventDefault();
                    const hiddenInput = document.querySelector('#key_input input');
                    if (hiddenInput) {
                        hiddenInput.value = e.key;
                        hiddenInput.dispatchEvent(new Event('input'));
                        hiddenInput.dispatchEvent(new Event('change'));
                    }
                }
            });
        }
        """
    )

demo.launch()