from __future__ import annotations

import argparse
import copy
import queue
import re
import threading
import tkinter as tk
from collections.abc import Sequence
from pathlib import Path
from tkinter import messagebox

from PIL import Image, ImageDraw, ImageFont, ImageTk

from copienator import (
    EvaluationWorkspace,
    ExitCode,
    atomic_write_json,
    execute,
    read_json,
    target_parser,
    workspace_from_target,
)
from copienator.platform import open_path
from copienator.utils import natural_key, read_all_labels

# --- Configuration & Globals ---
padding = 60
MISSING_LABEL_COLOR = "orange"
COMMON_MISSING_LABEL_COLOR = "#403a00"
COMMON_MISSING_THRESHOLD = 0.66

try:
    font = ImageFont.truetype("DejaVuSans.ttf", size=30)
except OSError:
    font = ImageFont.load_default()

# --- Helper Functions (Shared) ---

def page_number(b, nb_pages):
    column_width = 1000 // nb_pages
    center_x = (b[1] + b[3]) // 2
    return center_x // column_width


def sort_bounding_boxes(bounding_boxes, nb_pages):
    """Return label boxes in reading order: columns first, then top to bottom."""
    return sorted(
        bounding_boxes,
        key=lambda entry: (
            page_number(entry["box_2d"], nb_pages),
            entry["box_2d"][0],
            entry["box_2d"][1],
        ),
    )


def convert_box2d(b, pn_ori, npn, tot_ori, tot_dest):
    l = b.copy()
    l[1] = (l[1] - (1000 // tot_ori) * (pn_ori-1)) * tot_ori // tot_dest\
        + (1000 // tot_dest) * (npn - 1)
    l[3] = (l[3] - (1000 // tot_ori) * (pn_ori-1)) * tot_ori // tot_dest\
        + (1000 // tot_dest) * (npn - 1)
    return l

def convert_list(l, group_id, json_schema):
    ll = []
    nb_pages = json_schema["columns_per_file"][group_id-1]
    nb_previous_pages = sum([json_schema["columns_per_file"][i] for i in range(group_id-1)])
    nb_tot_pages = sum([e for e in json_schema["columns_per_file"]])
    for e in l:
        ee = e.copy()
        pn = page_number(e["box_2d"], nb_pages)
        npn = pn + nb_previous_pages
        ee["box_2d"] = convert_box2d(ee["box_2d"], pn, npn, nb_pages, nb_tot_pages)
        ee["part"] = group_id
        ee["pn"] = npn
        ll.append(ee)
    return ll


def normalized_labels(entries):
    return [
        str(value["label"]).removeprefix("|").removesuffix("|")
        for value in entries
        if str(value["label"]) != "_"
    ]


def frequently_missing_labels(
    copies_dir: Path,
    all_labels: list[str],
    threshold: float = COMMON_MISSING_THRESHOLD,
) -> set[str]:
    """Return labels absent from at least ``threshold`` of detected copies."""
    labels_by_copy: dict[str, set[str]] = {}
    for json_path in copies_dir.glob("*.json"):
        match = re.fullmatch(r"(.+)_\d+", json_path.stem)
        if match is None:
            continue
        try:
            data = read_json(json_path)
            entries = data["list"]
            if not isinstance(entries, list):
                continue
            present = set(normalized_labels(entries))
        except (OSError, KeyError, TypeError, ValueError):
            continue
        labels_by_copy.setdefault(match.group(1), set()).update(present)

    copy_count = len(labels_by_copy)
    if copy_count == 0:
        return set()
    return {
        label
        for label in all_labels
        if sum(label not in present for present in labels_by_copy.values())
        >= threshold * copy_count
    }


def label_color(
    label: str | None,
    all_labels: list[str],
    last_label_index: int,
    common_missing: set[str],
) -> tuple[str, int]:
    """Choose a label color and return the updated chronological index."""
    color = "black"
    if not label or label not in all_labels:
        return color, last_label_index

    current_index = all_labels.index(label)
    if current_index < last_label_index or (
        last_label_index == -1 and current_index != 0
    ):
        color = "red"
    elif current_index > last_label_index + 1:
        only_previous_is_missing = current_index == last_label_index + 2
        previous_label = all_labels[current_index - 1]
        color = (
            COMMON_MISSING_LABEL_COLOR
            if only_previous_is_missing and previous_label in common_missing
            else MISSING_LABEL_COLOR
        )
    return color, current_index


def prepare_image(
    image_path: str,
    bounding_boxes,
    all_labels,
    nb_pages,
    last_label_index,
    common_missing: set[str] | None = None,
):
    im = Image.open(image_path)
    im.load()
    width, height = im.size
    new_im = Image.new(im.mode, (width + padding, height), "white")
    new_im.paste(im, (0, 0))
    draw = ImageDraw.Draw(new_im)
    common_missing = common_missing or set()

    for bbox in sort_bounding_boxes(bounding_boxes, nb_pages):
        raw_y_min = int(bbox["box_2d"][0] * height / 1000)
        raw_x_min = int(bbox["box_2d"][1] * width  / 1000)
        raw_y_max = int(bbox["box_2d"][2] * height / 1000)
        raw_x_max = int(bbox["box_2d"][3] * width  / 1000)
        abs_y_min = max(0, raw_y_min - 10)
        abs_x_min = max(0, raw_x_min - 10)
        abs_y_max = min(height, raw_y_max + 10)
        abs_x_max = min(width, raw_x_max + 10)

        label = bbox.get("label")
        color, last_label_index = label_color(
            label,
            all_labels,
            last_label_index,
            common_missing,
        )

        draw.rectangle(((abs_x_min, abs_y_min), (abs_x_max, abs_y_max)), outline=color, width=4)
        if label:
            if abs_y_min > 80:
                draw.text((abs_x_min + 8, abs_y_min - 30), label, fill=color, font=font)
            else:
                draw.text((abs_x_min + 8, abs_y_max + 6), label, fill=color, font=font)
    return (new_im, last_label_index)

# --- Processing Logic (Worker Thread) ---

def _worker_items(base_dir, files_to_process, all_labels, output_queue):
    """
    Iterates through files, prepares VISUALS only, and puts metadata in queue.
    Does NOT write final JSON files anymore.
    """
    previous_copie = None
    last_label_index = None
    common_missing = frequently_missing_labels(base_dir / "Copies", all_labels)
    for img_path in files_to_process:
        json_path = base_dir / "Copies" / f"{img_path.stem}.json"
        copie_part = int(img_path.stem[-2:])
        copie = img_path.stem[:-3]
        if copie != previous_copie:
            last_label_index = -1
            previous_copie = copie
        json_schema_path = base_dir / 'Cutleft' / f"{copie}_schema.json"

        try:
            json_schema = read_json(json_schema_path)
        except (OSError, TypeError, ValueError):
            print("No json_schema : ", json_schema_path)
            continue

        nb_pages = json_schema["columns_per_file"][copie_part-1]

        if json_path.exists():
            # Read strictly for visualization purposes
            bb_list = []
            json_name = ""
            try:
                json_result = read_json(json_path)
                bb_list = json_result.get("list", [])
                json_name = json_result.get("name", "")
            except Exception as e:  # noqa: BLE001 - malformed user-editable JSON
                print(f"Warning: {json_path.name} is malformed! Loading blank. {e}")
                # We do NOT skip; we continue so the user can fix it in the GUI

            try:
                print(f"Buffering {img_path.name}...")
                (pil_image, last_label_index) = \
                    prepare_image(
                        str(img_path),
                        bb_list,
                        all_labels,
                        nb_pages,
                        last_label_index,
                        common_missing,
                    )
                error_msg = None

            except Exception as e:  # noqa: BLE001 - keep the item editable in the GUI
                print(f"Error processing {img_path.name}: {e}")
                pil_image = Image.open(str(img_path))
                error_msg = str(e)

            metadata = {
                "copie": copie,
                "part": copie_part,
                "schema": json_schema,
                "name": json_name,
                "error": error_msg
            }

            output_queue.put((pil_image, json_path, metadata))

def worker_thread(base_dir, files_to_process, all_labels, output_queue):
    """Prepare queue items and always terminate the GUI stream."""
    failure = None
    try:
        _worker_items(base_dir, files_to_process, all_labels, output_queue)
    except Exception as exc:  # noqa: BLE001 - worker boundary
        failure = str(exc)
        print(f"Plotting worker failed: {exc}")
    finally:
        metadata = {"worker_error": failure} if failure else None
        output_queue.put((None, None, metadata))

# --- GUI Logic (Main Thread) ---

class ImageViewer:
    def __init__(self, root, workspace, valid_labels, input_queue):
        self.root = root
        self.root.resizable(False, False) # If you resize, coordinates will be wrong

        screen_w = root.winfo_screenwidth()
        screen_h = root.winfo_screenheight()

        x = int(screen_w * 0.1)
        y = int(screen_h * 0.05)

        root.geometry(f"+{x}+{y}")

        self.workspace = workspace
        self.base_dir = workspace.root
        self.valid_labels = valid_labels
        self.image_queue = input_queue
        self.root.title("Bounding Box Viewer")
        self.label = tk.Label(root, text="Waiting for images...")
        self.label.pack(expand=True, fill="both")

        # Display State
        self.current_image = None
        self.current_json_path = None
        self.current_meta = None # Stores schema/copie info
        self.is_viewing = False
        self.scale_factor = 1.0
        self.orig_size = (1, 1)

        # Data Aggregation State
        self.active_copie_name = None
        self.accumulated_results = None # Dict with "name" and "list"

        # To go back
        self.history = []
        self.forward_stack = []
        self.current_pil_image = None
        self.failed = False

        from config import PLOTTING_KB

        # Bindings
        self.root.bind(PLOTTING_KB["OK"], self.on_enter)
        self.root.bind(PLOTTING_KB["previous"], self.on_previous)
        self.root.bind(PLOTTING_KB["edit"], self.on_edit)
        self.root.bind(PLOTTING_KB["open pdf"], self.on_open_pdf)
        self.root.bind(PLOTTING_KB["open original pdf"], self.on_open_ori_pdf)
        self.root.bind(PLOTTING_KB["open eval"], self.on_open_interro)
        self.root.bind('<Escape>', lambda _event: self.close())
        self.root.protocol("WM_DELETE_WINDOW", self.close)
        self.label.bind('<Button-1>', self.on_click)

        self.poll_queue()

    def poll_queue(self):
        if not self.is_viewing:
            try:
                # pil_image, json_path, metadata = image_queue.get_nowait()
                if self.forward_stack:
                    pil_image, json_path, metadata = self.forward_stack.pop()
                else:
                    pil_image, json_path, metadata = self.image_queue.get_nowait()

                # Handle End of Stream
                if pil_image is None:
                    if metadata and metadata.get("worker_error"):
                        self.failed = True
                        messagebox.showerror(
                            "Processing Error", metadata["worker_error"]
                        )
                    self.save_current_batch() # Save any remaining data
                    print("All images processed.")
                    self.root.quit()
                    return

                # Check if we switched to a new "Copie" group
                if self.active_copie_name != metadata["copie"]:
                    self.save_current_batch() # Write previous group to disk
                    # Start new batch
                    self.active_copie_name = metadata["copie"]
                    self.accumulated_results = {"name": metadata["name"], "list": []}

                self.display_image(pil_image, json_path, metadata)
            except queue.Empty:
                pass
        self.root.after(100, self.poll_queue)

    def save_current_batch(self):
        """Writes the accumulated data to the main JSON file."""
        if self.active_copie_name and self.accumulated_results:
            main_json_path = self.base_dir / "Copies" / f"{self.active_copie_name}.json"
            print(f"Writing aggregated result to {main_json_path}")
            atomic_write_json(main_json_path, self.accumulated_results)
            self.accumulated_results = None

    def close(self):
        self.root.quit()


    def on_previous(self, event):
        if self.is_viewing and self.history:
            print("Going back to previous image...")
            (
                prev_pil,
                prev_json,
                prev_meta,
                previous_copie_name,
                previous_results,
            ) = self.history.pop()

            # Push current image to the forward stack so we don't lose it
            self.forward_stack.append((self.current_pil_image,
                                       self.current_json_path, self.current_meta))

            # Restore the aggregation exactly as it was before validating the
            # previous image. This also makes navigation across copies safe.
            self.active_copie_name = previous_copie_name
            self.accumulated_results = previous_results

            # Display the previous image immediately
            self.display_image(prev_pil, prev_json, prev_meta)
    def display_image(self, pil_image, json_path, metadata):
        self.current_pil_image = pil_image # ADD THIS LINE
        self.orig_size = pil_image.size
        self.scale_factor = 1.0
        screen_h = self.root.winfo_screenheight() - 100
        if pil_image.height > screen_h:
            self.scale_factor = screen_h / pil_image.height
            pil_image = pil_image.resize((int(pil_image.width * self.scale_factor),
                                          int(pil_image.height * self.scale_factor)))

        self.tk_image = ImageTk.PhotoImage(pil_image)
        self.label.config(image=self.tk_image, text=f"Processing: {json_path.name}")
        self.current_json_path = json_path
        self.current_meta = metadata
        self.is_viewing = True
        self.root.lift()

        if metadata.get("error"):
            msg = f"Error generating boxes for {json_path.name}:\n\n{metadata['error']}\n\nPlease press 'e' to fix the JSON file, then press Enter to retry."
            messagebox.showerror("Processing Error", msg)

    def on_enter(self, event):
        if self.is_viewing:
            print(f"Committing data for {self.current_json_path.name}...")
            previous_copie_name = self.active_copie_name
            previous_results = copy.deepcopy(self.accumulated_results)

            try:
                current_data = read_json(self.current_json_path)
                items = current_data["list"]

                # Perform the conversion now, post-edit
                converted_items = convert_list(
                    items,
                    self.current_meta["part"],
                    self.current_meta["schema"]
                )

                labels = normalized_labels(items)
                false_labels = [
                    label for label in labels if label not in self.valid_labels
                ]

                if false_labels:
                    msg = f"Wrong label in {self.current_json_path.name}: {false_labels}\n\n\tPlease press 'e' to fix it, then press Enter again."
                    print(msg)
                    messagebox.showerror("Label Error", msg)
                    return
                # Add to accumulator
                if self.accumulated_results:
                    self.accumulated_results["list"].extend(converted_items)
                    # Update name just in case (though usually consistent per group)
                    if "name" in current_data and current_data["name"] != "Continued":
                        self.accumulated_results["name"] = current_data["name"]

            except Exception as e:  # noqa: BLE001 - interactive validation boundary
                # Warn user and STOP (do not advance to next image)
                msg = f"Error reading {self.current_json_path.name}:\n\n{e}\n\nPlease press 'e' to fix it, then press Enter again."
                print(msg)
                messagebox.showerror("JSON Error", msg)
                return  # Abort advancement

            self.history.append(
                (
                    self.current_pil_image,
                    self.current_json_path,
                    self.current_meta,
                    previous_copie_name,
                    previous_results,
                )
            )

            # Advance UI
            self.is_viewing = False
            self.label.config(image="", text="Loading next...")

    def on_open_pdf(self, event):
        if self.is_viewing and self.current_json_path:
            a = self.current_json_path.stem.split('_')[0] + ".pdf"
            pdf_path = self.current_json_path.with_name(a)
            print(f"Opening {pdf_path}")
            open_path(pdf_path)

    def on_open_interro(self, event):
        if self.is_viewing and self.current_json_path:
        # Check local directory first
            local_accent = self.base_dir / "énoncé.pdf"
            local_plain = self.base_dir / "enonce.pdf"

            if local_accent.exists():
                pdf_path = str(local_accent)
            elif local_plain.exists():
                pdf_path = local_plain
            else:
                messagebox.showerror(
                    "PDF not found",
                    f"Neither {local_accent.name} nor {local_plain.name} exists in {self.base_dir}.",
                )
                return

            print(f"Opening {pdf_path}")
            open_path(pdf_path)

    def on_open_ori_pdf(self, event):
        if self.is_viewing and self.current_json_path:
            new_filename = self.current_json_path.stem.split('_')[0] + ".pdf"
            pdf_path = self.base_dir / "Copies Originales" / new_filename
            print(f"Opening {pdf_path}")
            open_path(pdf_path)

    def on_edit(self, event):
        if self.is_viewing and self.current_json_path:
            print(f"Opening {self.current_json_path}")
            open_path(self.current_json_path)

    def on_click(self, event):
        if not self.is_viewing: return
        x = int(event.x / self.scale_factor)
        y = int(event.y / self.scale_factor)
        w, h = self.orig_size
        box = [
            int(max(0, y - 5) / h * 1000),
            int(max(0, x - 5) / (w- padding) * 1000),
            int(min(h, y + 5) / h * 1000),
            int(min(w, x + 5) / (w - padding) * 1000),
        ]
        box_str = "{ \"box_2d\": " + str(box) + ", \"label\": \"\" },"
        print(f"Copied box at ({x},{y}): {box_str}")
        self.root.clipboard_clear()
        self.root.clipboard_append(box_str)

def _selected_images(workspace: EvaluationWorkspace, target: Path) -> list[Path]:
    workspace.require_directories("Cutleft", "Copies")
    if target.is_file():
        stem = target.stem
        exact = workspace.cutleft_dir / f"{stem}.jpg"
        if exact.is_file():
            return [exact]
        return sorted(
            workspace.cutleft_dir.glob(f"{stem}_*.jpg"),
            key=natural_key,
        )
    return sorted(workspace.cutleft_dir.glob("*.jpg"), key=natural_key)


def run(workspace: EvaluationWorkspace, target: Path) -> ExitCode:
    workspace.require_files("labels")
    all_labels = read_all_labels(workspace.root)
    files_to_process = _selected_images(workspace, target)
    if not files_to_process:
        print(f"No Cutleft images found for {target}")
        return ExitCode.PARTIAL

    print(
        "o to open pdf, O original pdf, e to edit part, p to go back, "
        "i to open the statement, click for coordinates"
    )
    input_queue = queue.Queue(maxsize=5)
    worker = threading.Thread(
        target=worker_thread,
        args=(workspace.root, files_to_process, all_labels, input_queue),
        daemon=True,
    )
    worker.start()
    root = tk.Tk()
    application = ImageViewer(root, workspace, set(all_labels), input_queue)
    root.mainloop()
    return ExitCode.PARTIAL if application.failed else ExitCode.SUCCESS


def build_parser() -> argparse.ArgumentParser:
    return target_parser("Interactively verify detected label coordinates")


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()

    def handle(args: argparse.Namespace) -> ExitCode:
        workspace, target = workspace_from_target(args)
        return run(workspace, target)

    return execute(parser, argv, handle)


if __name__ == "__main__":
    raise SystemExit(main())