"""Dot colour sample labeler — Tk-based, safe to import headlessly."""
from __future__ import annotations

import json
from pathlib import Path

VALID_LABELS: frozenset[str] = frozenset({
    "turquoise", "yellow", "dark_green", "dark_red",
    "light_green", "light_red", "gray", "skip",
})


class LabelStore:
    """Persistent dict-backed label store serialised as JSON."""

    def __init__(self, path: Path) -> None:
        self.path = Path(path)
        self._labels: dict[str, str] = {}
        if self.path.exists():
            self._labels = json.loads(self.path.read_text(encoding="utf-8"))

    def __getitem__(self, key: str) -> str:
        return self._labels[key]

    def __setitem__(self, key: str, label: str) -> None:
        self._labels[key] = label

    def get(self, key: str, default: str | None = None) -> str | None:
        return self._labels.get(key, default)

    def save(self) -> None:
        self.path.write_text(json.dumps(self._labels, indent=2), encoding="utf-8")

    def as_dict(self) -> dict[str, str]:
        return dict(self._labels)


def run_labeler(samples_dir: Path, out_path: Path) -> None:
    """Launch the Tk labeling wizard. Imports tkinter only inside this function."""
    import tkinter as tk
    from PIL import Image, ImageTk  # type: ignore[import-untyped]

    samples_dir = Path(samples_dir)
    out_path = Path(out_path)
    store = LabelStore(out_path)

    images = sorted(samples_dir.glob("*.png")) + sorted(samples_dir.glob("*.jpg"))
    unlabeled = [p for p in images if store.get(p.name) is None]

    if not unlabeled:
        print("All samples already labeled.")
        return

    root = tk.Tk()
    root.title("ATM Labeler")

    img_label = tk.Label(root)
    img_label.pack(padx=8, pady=8)
    status_var = tk.StringVar()
    tk.Label(root, textvariable=status_var).pack()

    idx = [0]
    _photo_ref: list[object] = []  # keep PhotoImage alive

    def load_current() -> None:
        i = idx[0]
        if i >= len(unlabeled):
            store.save()
            root.destroy()
            return
        img = Image.open(unlabeled[i]).resize((200, 200))
        photo = ImageTk.PhotoImage(img)
        _photo_ref[:] = [photo]
        img_label.config(image=photo)
        status_var.set(f"{i + 1}/{len(unlabeled)}: {unlabeled[i].name}")

    def label_and_next(lbl: str) -> None:
        store[unlabeled[idx[0]].name] = lbl
        idx[0] += 1
        load_current()

    btn_frame = tk.Frame(root)
    btn_frame.pack(pady=4)
    for lbl in sorted(VALID_LABELS - {"skip"}):
        tk.Button(btn_frame, text=lbl, command=lambda l=lbl: label_and_next(l)).pack(
            side=tk.LEFT, padx=2
        )
    tk.Button(btn_frame, text="skip", command=lambda: label_and_next("skip")).pack(
        side=tk.LEFT, padx=2
    )

    load_current()
    root.mainloop()


def accuracy(
    labels: dict[str, str],
    predicted: dict[str, str],
) -> dict[str, float]:
    """Per-label precision/recall/F1 and overall accuracy.

    Returns a flat dict with keys:
      ``accuracy``, ``<label>_precision``, ``<label>_recall``, ``<label>_f1``
    for every label that appears in either *labels* or *predicted*.
    """
    common = set(labels) & set(predicted)
    all_label_values = sorted(set(labels.values()) | set(predicted.values()))

    result: dict[str, float] = {}

    if common:
        correct = sum(1 for k in common if labels[k] == predicted[k])
        result["accuracy"] = correct / len(common)
    else:
        result["accuracy"] = 0.0

    for lbl in all_label_values:
        tp = sum(1 for k in common if labels[k] == lbl and predicted[k] == lbl)
        fp = sum(1 for k in common if labels[k] != lbl and predicted[k] == lbl)
        fn = sum(1 for k in common if labels[k] == lbl and predicted[k] != lbl)

        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
        f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) > 0 else 0.0

        result[f"{lbl}_precision"] = precision
        result[f"{lbl}_recall"] = recall
        result[f"{lbl}_f1"] = f1

    return result
