from src.DataManager.base import BaseDataManager
from src.DataManager.utils import imread_rgb, imwrite_rgb

import numpy as np
from pathlib import Path


class ImageDataManager(BaseDataManager):
    def __init__(self, src_data: Path, output_dir: Path):
        self.output_dir: Path = output_dir
        self.output_dir.mkdir(exist_ok=True)
        self.output_dir = output_dir / "img"
        self.output_dir.mkdir(exist_ok=True)

        self.data_paths = []
        if src_data.is_file():
            self.data_paths.append(src_data)
        elif src_data.is_dir():
            self.data_paths = (
                list(src_data.glob("*.jpg"))
                + list(src_data.glob("*.jpeg"))
                + list(src_data.glob("*.png"))
            )

        assert len(self.data_paths), "Data must be supplied!"

        self.data_paths_iter = iter(self.data_paths)

        self.last_idx = -1

    def __len__(self):
        return len(self.data_paths)

    def get(self) -> np.ndarray:
        img_path = next(self.data_paths_iter)
        self.last_idx += 1
        return imread_rgb(img_path)

    def save(self, img: np.ndarray):
        filename = "swap_" + Path(self.data_paths[self.last_idx]).name

        imwrite_rgb(self.output_dir / filename, img)