import os
from pathlib import Path

from torchaudio.datasets import utils as dataset_utils

from ..common_utils import (
    TempDirMixin,
    TorchaudioTestCase,
)


class TestWalkFiles(TempDirMixin, TorchaudioTestCase):
    root = None
    expected = None

    def _add_file(self, *parts):
        path = self.get_temp_path(*parts)
        self.expected.append(path)
        Path(path).touch()

    def setUp(self):
        self.root = self.get_temp_path()
        self.expected = []

        # level 1
        for filename in ['a.txt', 'b.txt', 'c.txt']:
            self._add_file(filename)

        # level 2
        for dir1 in ['d1', 'd2', 'd3']:
            for filename in ['d.txt', 'e.txt', 'f.txt']:
                self._add_file(dir1, filename)
            # level 3
            for dir2 in ['d1', 'd2', 'd3']:
                for filename in ['g.txt', 'h.txt', 'i.txt']:
                    self._add_file(dir1, dir2, filename)

        print('\n'.join(self.expected))

    def test_walk_files(self):
        """walk_files should traverse files in alphabetical order"""
        n_ites = 0
        for i, path in enumerate(dataset_utils.walk_files(self.root, '.txt', prefix=True)):
            found = os.path.join(self.root, path)
            assert found == self.expected[i]
            n_ites += 1
        assert n_ites == len(self.expected)