from dataclasses import dataclass
from typing import TYPE_CHECKING, Union
from pathlib import Path
import random


if TYPE_CHECKING:
    from divergent_gravity.functionality.flattened_rotation_curves.galaxy import (
        GalaxyDataset,
    )


@dataclass
class PlotSaver:
    directory: str
    """
    directory must be set to an absolute path.
    """
    file_name: Union[str, None] = None
    """
    An optional name. If none is set, a unique name will be generated for each plot type,
    for each galaxy. This means that the same plot will overwrite the generated plot of the same type,
    for the same galaxy if the `directory` is the same.
    """
    randomize: bool = False
    """
    If randomize=True, a random string will be appended to each file name to avoid overwriting similar plots.
    """

    def random_string(self) -> str:
        l = "abcdefghijklnopqrstuvwxyz1234567890"
        sample = random.sample(l.split(""), 8)
        return "".join(sample)

    def get_file_name(
        self, plot_key: str, gal: Union["GalaxyDataset", None] = None
    ) -> str:
        s = f"{plot_key}"
        if gal is not None:
            s = f"{s}-{gal.catalog_name}"
        if self.randomize:
            s = f"{s}-{self.random_string()}"
        return f"{s}.png"

    def to_path(self, plot_key: str, gal: Union["GalaxyDataset", None] = None) -> Path:
        return Path(self.directory) / self.get_file_name(plot_key, gal)
