import os

import numpy as np
import tinygrad


class Manager:
    def __init__(self) -> None:
        self._classes = []
        self._class_to_idx = {}
        self._samples = []

        print("Manager initialized")

    def handle_index_imagefolder(self, root: str = "/dataset") -> None:
        classes = sorted(
            d for d in os.listdir(root)
            if os.path.isdir(os.path.join(root, d))
        )

        class_to_idx = {c: i for i, c in enumerate(classes)}

        samples = []
        for c in classes:
            cdir = os.path.join(root, c)
            for name in os.listdir(cdir):
                if name.lower().endswith((".jpg", ".jpeg", ".png")):
                    samples.append((os.path.join(cdir, name), class_to_idx[c]))

        print(f"Indexed ImageFolder: {len(samples)} samples found in {len(classes)} classes.")

        self._classes = classes
        self._class_to_idx = class_to_idx
        self._samples = samples

    def handle_get_random_sample(self) -> dict:
        if not self._samples:
            raise RuntimeError("No samples indexed. Please upload an image folder first.")

        idx = np.random.randint(0, len(self._samples))
        path, class_idx = self._samples[idx]
        class_name = self._classes[class_idx]
        return {"imagePath": path, "label": class_name}