import glob
import os
from typing import List
import matplotlib.pyplot as plt

class SavefigRedirect:
    """Контекст, который перенаправляет plt.savefig('data/...') в нужную папку."""
    def __init__(self, base_dir: str):
        self.base_dir = base_dir
        self._orig = None

    def __enter__(self):
        self._orig = plt.savefig
        def _patched(fname, *args, **kwargs):
            if isinstance(fname, str) and (fname.startswith("data\\") or fname.startswith("data/")):
                out = os.path.join(self.base_dir, fname.replace("data\\", "data/").replace("data/", "data/"))
                os.makedirs(os.path.dirname(out), exist_ok=True)
                return self._orig(out, *args, **kwargs)
            return self._orig(fname, *args, **kwargs)
        plt.savefig = _patched
        return self

    def __exit__(self, exc_type, exc, tb):
        if self._orig:
            plt.savefig = self._orig

def collect_outputs(work_dir: str) -> List[str]:
    candidates = []
    candidates += glob.glob(os.path.join(work_dir, "data", "**", "*.*"), recursive=True)
    candidates += glob.glob(os.path.join(work_dir, "*.png"))
    out = [p for p in sorted(set(candidates)) if os.path.isfile(p)]
    return out