from dataclasses import dataclass
from pathlib import Path

from modules.shared.ConfigDataClass import Config


@dataclass
class ModelPackageCollection(Config):
    name: str = None
    external_packages: list[str] = None
    unsorted_packages: list[str] = None
    categorized_packages: dict[str, list[str]] = None


    def __post_init__(self):
        if not self.external_packages: self.external_packages = list()
        if not self.unsorted_packages: self.unsorted_packages = list()
        if not self.categorized_packages: self.categorized_packages = dict()
        super().__post_init__()
        if not self.name: raise ValueError('ModelPackageCollection(): name must be specified')

    def add_package(self, pkg_name, category: str = None, internal=True):
        if not internal and pkg_name not in self.external_packages: self.external_packages.append(pkg_name)
        elif not category and pkg_name not in self.unsorted_packages: self.unsorted_packages.append(pkg_name)
        else:
            if category not in self.categorized_packages: self.categorized_packages[category] = list()
            if pkg_name not in self.categorized_packages[category]: self.categorized_packages[category].append(pkg_name)
        self.save()

    def get_path(self, pkg_name) -> Path:
        if pkg_name in self.external_packages: return Path('')
        elif pkg_name in self.unsorted_packages: return  Path(self.name)
        else:
            for category in self.categorized_packages:
                if pkg_name in self.categorized_packages[category]: return Path(self.name) / category

        raise FileNotFoundError(f'package {pkg_name} not in collection {self.name}')

    @property
    def paths_dict(self) -> dict[str, Path]:
        result: dict[str, Path] = dict()
        for category in self.categorized_packages:
            for pkg_name in list(self.categorized_packages[category]):
                result[pkg_name] = Path(self.name) / category

        for pkg_name in list(self.unsorted_packages): result[pkg_name] = Path(self.name)
        for pkg_name in list(self.external_packages): result[pkg_name] = Path('')

        return result