import sys
import subprocess
import tkinter as tk
from tkinter import messagebox, filedialog, scrolledtext, ttk
import json
import csv
import threading
import time
import os

try:
    import cv2
    import numpy as np
    from PIL import Image, ImageDraw, ImageFont
except ImportError:
    root = tk.Tk()
    root.withdraw()
    if messagebox.askyesno("Dependency Error", "Required libraries are not installed. Would you like to try and install them now?"):
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install", "opencv-python", "numpy", "Pillow"])
            messagebox.showinfo("Success", "Dependencies installed successfully. Please restart the application.")
        except Exception as e:
            messagebox.showerror("Installation Failed", f"Could not install dependencies. Please run 'pip install -r requirements.txt' manually.\n\nError: {e}")
    root.destroy()
    sys.exit()

class ToolTip:
    def __init__(self, widget, text):
        self.widget = widget
        self.text = text
        self.tooltip = None
        self.widget.bind("<Enter>", self.enter)
        self.widget.bind("<Leave>", self.leave)

    def enter(self, event=None):
        x, y, _, _ = self.widget.bbox("insert")
        x += self.widget.winfo_rootx() + 25
        y += self.widget.winfo_rooty() + 25

        self.tooltip = tk.Toplevel(self.widget)
        self.tooltip.wm_overrideredirect(True)
        self.tooltip.wm_geometry(f"+{x}+{y}")
        
        label = tk.Label(self.tooltip, text=self.text, background="#ffffe0", relief="solid", borderwidth=1, wraplength=200)
        label.pack()

    def leave(self, event=None):
        if self.tooltip:
            self.tooltip.destroy()
            self.tooltip = None

class MasterImageFinderApp:
    def __init__(self, root):
        self.root = root
        self.root.title("Master Image Finder")
        self.root.geometry("800x700")

        self.layouts_path = tk.StringVar()
        self.masters_path = tk.StringVar()
        self.upscale = tk.BooleanVar()
        self.denoise = tk.BooleanVar()
        self.sharpen = tk.BooleanVar()
        self.contrast = tk.BooleanVar()

        tk.Label(root, text="Layouts Folder:").pack(pady=5)
        tk.Entry(root, textvariable=self.layouts_path, width=100).pack(pady=5)
        tk.Button(root, text="Select Layouts Folder", command=self.select_layouts_folder).pack(pady=5)

        tk.Label(root, text="Master Images Folder:").pack(pady=5)
        tk.Entry(root, textvariable=self.masters_path, width=100).pack(pady=5)
        tk.Button(root, text="Select Master Images Folder", command=self.select_masters_folder).pack(pady=5)

        enhancement_frame = tk.LabelFrame(root, text="Advanced Enhancement Options", padx=10, pady=10)
        enhancement_frame.pack(pady=10, padx=10, fill="x")

        upscale_check = tk.Checkbutton(enhancement_frame, text="Smart Upscaling", variable=self.upscale)
        upscale_check.grid(row=0, column=0, sticky="w")
        ToolTip(upscale_check, "Enlarges small images to improve feature detection. Best for images under 400x400px.")

        denoise_check = tk.Checkbutton(enhancement_frame, text="Denoising", variable=self.denoise)
        denoise_check.grid(row=0, column=1, sticky="w", padx=10)
        ToolTip(denoise_check, "Removes digital noise and compression artifacts. Can be slow on large images.")

        sharpen_check = tk.Checkbutton(enhancement_frame, text="Sharpening", variable=self.sharpen)
        sharpen_check.grid(row=1, column=0, sticky="w")
        ToolTip(sharpen_check, "Enhances edges and fine details. Very fast.")

        contrast_check = tk.Checkbutton(enhancement_frame, text="Contrast Enhancement", variable=self.contrast)
        contrast_check.grid(row=1, column=1, sticky="w", padx=10)
        ToolTip(contrast_check, "Improves local contrast, making features in dark or washed-out areas more distinct.")

        self.run_button = tk.Button(root, text="Find Matches", command=self.run_finder_thread)
        self.run_button.pack(pady=20)

        self.progress = ttk.Progressbar(root, orient="horizontal", length=780, mode="determinate")
        self.progress.pack(pady=10)

        self.log_area = scrolledtext.ScrolledText(root, wrap=tk.WORD, width=100, height=15)
        self.log_area.pack(pady=10, padx=10)

    def select_layouts_folder(self):
        self.layouts_path.set(filedialog.askdirectory())

    def select_masters_folder(self):
        self.masters_path.set(filedialog.askdirectory())

    def log(self, message):
        self.root.after(0, self._log, message)

    def _log(self, message):
        self.log_area.insert(tk.END, message + "\n")
        self.log_area.see(tk.END)

    def update_progress(self, value):
        self.root.after(0, self.progress.config, {'value': value})

    def run_finder_thread(self):
        self.run_button.config(state=tk.DISABLED)
        self.log_area.delete(1.0, tk.END)
        self.progress['value'] = 0
        thread = threading.Thread(target=self.run_finder)
        thread.start()

    def run_finder(self):
        start_time = time.time()
        layouts_dir = self.layouts_path.get()
        masters_dir = self.masters_path.get()
        
        if not layouts_dir or not masters_dir:
            self.log("Error: Please select both folders.")
            self.run_button.config(state=tk.NORMAL)
            return

        output_dir = os.path.join(layouts_dir, "reports")
        if not os.path.isdir(output_dir):
            os.makedirs(output_dir)
            self.log(f"Created output directory at {output_dir}")

        try:
            results = self.find_matches(layouts_dir, masters_dir)
            self.create_html_report(results, output_dir, layouts_dir, masters_dir)
            end_time = time.time()
            
            total_matches = sum(1 for item in results if item['found'])
            self.log("\n--- Process Complete! ---")
            self.log(f"Found matches for {total_matches} out of {len(results)} layout images.")
            self.log(f"Total time: {end_time - start_time:.2f} seconds")
            self.log(f"Reports saved in: {output_dir}")
        except Exception as e:
            self.log(f"An error occurred: {e}")
        finally:
            self.run_button.config(state=tk.NORMAL)

    def find_matches(self, layouts_path, masters_path, min_good_matches=10, inlier_threshold_ratio=0.5):
        akaze = cv2.AKAZE_create()
        bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        sharpen_kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])

        layout_images = [f for f in os.listdir(layouts_path) if f.endswith(('.png', '.jpg', '.jpeg'))]
        master_images = [f for f in os.listdir(masters_path) if f.endswith(('.png', '.jpg', '.jpeg'))]
        
        results = []

        master_descriptors = {}
        self.log("Preprocessing master images...")
        for i, master_image_name in enumerate(master_images):
            self.log(f"  - Preprocessing {i+1}/{len(master_images)}: {master_image_name}")
            master_image_path = os.path.join(masters_path, master_image_name)
            master_img = cv2.imread(master_image_path, cv2.IMREAD_GRAYSCALE)
            if master_img is None: continue
            kp, des = akaze.detectAndCompute(master_img, None)
            if des is not None: master_descriptors[master_image_name] = (kp, des)

        total_layouts = len(layout_images)
        self.log("\nProcessing layout images...")
        for i, layout_image_name in enumerate(layout_images):
            self.update_progress((i / total_layouts) * 100)
            self.log(f"  - Processing {i+1}/{total_layouts}: {layout_image_name}")
            
            layout_image_path = os.path.join(layouts_path, layout_image_name)
            layout_img_gray = cv2.imread(layout_image_path, cv2.IMREAD_GRAYSCALE)
            if layout_img_gray is None:
                self.log(f"    - Could not read layout image.")
                continue

            enhancements_applied = []
            if self.upscale.get() and (layout_img_gray.shape[0] < 400 or layout_img_gray.shape[1] < 400):
                layout_img_gray = cv2.resize(layout_img_gray, (0,0), fx=2.0, fy=2.0, interpolation=cv2.INTER_LANCZOS4)
                enhancements_applied.append("Upscaled")
            if self.denoise.get():
                layout_img_gray = cv2.fastNlMeansDenoising(layout_img_gray, None, 10, 7, 21)
                enhancements_applied.append("Denoised")
            if self.sharpen.get():
                layout_img_gray = cv2.filter2D(layout_img_gray, -1, sharpen_kernel)
                enhancements_applied.append("Sharpened")
            if self.contrast.get():
                layout_img_gray = clahe.apply(layout_img_gray)
                enhancements_applied.append("Contrast Enhanced")
            
            if enhancements_applied:
                self.log(f"    - Applied: {', '.join(enhancements_applied)}")

            kp1, des1 = akaze.detectAndCompute(layout_img_gray, None)
            if des1 is None:
                self.log(f"    - No features found.")
                results.append({"layout": layout_image_name, "masters": [], "found": False, "enhancements": enhancements_applied})
                continue

            all_possible_matches = []
            for master_image_name, (kp2, des2) in master_descriptors.items():
                matches = bf.knnMatch(des1, des2, k=2)
                good_matches = [m for m, n in matches if len(matches) > 1 and len(matches[0]) > 1 and m.distance < 0.75 * n.distance]

                if len(good_matches) > min_good_matches:
                    src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
                    dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
                    M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
                    if mask is not None:
                        inliers = np.sum(mask)
                        if inliers > min_good_matches:
                            all_possible_matches.append({"master": master_image_name, "inliers": int(inliers)})

            if not all_possible_matches:
                results.append({"layout": layout_image_name, "masters": [], "found": False, "enhancements": enhancements_applied})
                continue

            best_match = max(all_possible_matches, key=lambda x: x['inliers'])
            max_inliers = best_match['inliers']
            
            confident_matches = [best_match]
            for match in all_possible_matches:
                if match != best_match and match['inliers'] > max_inliers * inlier_threshold_ratio:
                    confident_matches.append(match)

            if confident_matches:
                self.log(f"    - Found {len(confident_matches)} confident master image(s).")
                results.append({"layout": layout_image_name, "masters": confident_matches, "found": True, "enhancements": enhancements_applied})
            else:
                results.append({"layout": layout_image_name, "masters": [], "found": False, "enhancements": enhancements_applied})
        
        self.update_progress(100)
        return results

    def create_html_report(self, data, output_path, layouts_abs_path, masters_abs_path):
        report_path = os.path.join(output_path, 'report.html')
        
        total_matches = sum(1 for item in data if item['found'])
        total_layouts = len(data)

        html = f"""
<!DOCTYPE html><html lang="en"><head><meta charset="UTF-8"><title>Image Match Report</title>
<style>
    body {{ font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; margin: 20px; background-color: #f9f9f9; }}
    .container {{ max-width: 1600px; margin: auto; }}
    h1, h2 {{ text-align: center; color: #333; border-bottom: 2px solid #eee; padding-bottom: 10px; }}
    .summary {{ text-align: center; margin-bottom: 20px; font-size: 1.2em; }}
    .card {{ background: #fff; border-radius: 8px; box-shadow: 0 4px 8px rgba(0,0,0,0.1); margin-bottom: 20px; padding: 20px; }}
    .layout-container {{ display: flex; gap: 20px; align-items: flex-start; }}
    .layout-box {{ flex: 1; text-align: center; }}
    .layout-box img {{ max-width: 400px; max-height: 400px; height: auto; border-radius: 4px; cursor: pointer; }}
    .masters-grid {{ flex: 2; display: grid; grid-template-columns: repeat(auto-fill, minmax(150px, 1fr)); gap: 15px; }}
    .master-box {{ text-align: center; border: 1px solid #ddd; border-radius: 8px; padding: 10px; }}
    .master-box img {{ width: 100%; height: 150px; object-fit: cover; border-radius: 4px; cursor: pointer; }}
    .filename {{ white-space: nowrap; overflow: hidden; text-overflow: ellipsis; font-weight: bold; }}
    .alternatives {{ font-size: 0.8em; color: #666; margin-top: 5px; }}
    .unmatched-grid {{ display: grid; grid-template-columns: repeat(auto-fill, minmax(150px, 1fr)); gap: 15px; }}
    .unmatched-card img {{ width: 100%; height: 150px; object-fit: cover; }}
    .modal {{ display: none; position: fixed; z-index: 1000; padding-top: 60px; left: 0; top: 0; width: 100%; height: 100%; overflow: auto; background-color: rgba(0,0,0,0.9); }}
    .modal-content {{ margin: auto; display: block; max-width: 90%; max-height: 90vh; }}
    .close {{ position: absolute; top: 15px; right: 35px; color: #f1f1f1; font-size: 40px; font-weight: bold; cursor: pointer; }}
</style>
</head><body><div class="container"><h1>Image Match Report</h1><div class="summary">Found matches for {total_matches} out of {total_layouts} layout images.</div>"""

        html += "<h2>Matched Layouts</h2>"
        for item in data:
            if item['found']:
                layout_img_path = os.path.join(layouts_abs_path, item['layout']).replace('\\', '/')
                enhancements_str = f"<p class='enhancements'>Enhancements: {', '.join(item['enhancements'])}</p>" if item['enhancements'] else ""
                html += f"<div class='card'><div class='layout-container'>"
                html += f"<div class='layout-box'><h3>Layout Image</h3><img src='file:///{layout_img_path}' onclick='openModal(this.src)'><p class='filename' title='{item['layout']}'>{item['layout']}</p>{enhancements_str}</div>"
                html += "<div class='masters-grid'>"
                for master_item in item['masters']:
                    master_img_path = os.path.join(masters_abs_path, master_item['master']).replace('\\', '/')
                    html += f"<div class='master-box'><img src='file:///{master_img_path}' onclick='openModal(this.src)'><p class='filename' title='{master_item['master']}'>{master_item['master']}</p><p>({master_item['inliers']} inliers)</p></div>"
                html += "</div></div></div>"

        html += "<h2>Unmatched Layouts</h2><p style='text-align:center;'>Please review these manually.</p><div class='unmatched-grid'>"
        for item in data:
            if not item['found']:
                layout_img_path = os.path.join(layouts_abs_path, item['layout']).replace('\\', '/')
                html += f"<div class='card unmatched-card'><img src='file:///{layout_img_path}' onclick='openModal(this.src)'><p class='filename' title='{item['layout']}'>{item['layout']}</p></div>"
        
        html += """</div></div>
<div id="myModal" class="modal"><span class="close" onclick="closeModal()">&times;</span><img class="modal-content" id="img01"></div>
<script>
    var modal = document.getElementById("myModal");
    var modalImg = document.getElementById("img01");
    function openModal(src) { modal.style.display = "block"; modalImg.src = src; }
    function closeModal() { modal.style.display = "none"; }
</script>
</body></html>"""

        with open(report_path, 'w', encoding='utf-8') as f:
            f.write(html)
        self.log(f"HTML report saved to {report_path}")

if __name__ == "__main__":
    root = tk.Tk()
    app = MasterImageFinderApp(root)
    root.mainloop()