import time
import unittest
from pathlib import Path
from unittest.mock import patch

import pytest
from pytest import CaptureFixture

from huggingface_hub.utils import (
    SoftTemporaryDirectory,
    are_progress_bars_disabled,
    disable_progress_bars,
    enable_progress_bars,
    tqdm,
    tqdm_stream_file,
)


class TestTqdmUtils(unittest.TestCase):
    @pytest.fixture(autouse=True)
    def capsys(self, capsys: CaptureFixture) -> None:
        """Workaround to make capsys work in unittest framework.

        Capsys is a convenient pytest fixture to capture stdout.
        See https://waylonwalker.com/pytest-capsys/.

        Taken from https://github.com/pytest-dev/pytest/issues/2504#issuecomment-309475790.
        """
        self.capsys = capsys

    def setUp(self) -> None:
        """Get verbosity to set it back after the tests."""
        self._previous_are_progress_bars_disabled = are_progress_bars_disabled()
        return super().setUp()

    def tearDown(self) -> None:
        """Set back progress bars verbosity as before testing."""
        if self._previous_are_progress_bars_disabled:
            disable_progress_bars()
        else:
            enable_progress_bars()

    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", None)
    def test_tqdm_helpers(self) -> None:
        """Test helpers to enable/disable progress bars."""
        disable_progress_bars()
        self.assertTrue(are_progress_bars_disabled())

        enable_progress_bars()
        self.assertFalse(are_progress_bars_disabled())

    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", True)
    def test_cannot_enable_tqdm_when_env_variable_is_set(self) -> None:
        """
        Test helpers cannot enable/disable progress bars when
        `HF_HUB_DISABLE_PROGRESS_BARS` is set.
        """
        disable_progress_bars()
        self.assertTrue(are_progress_bars_disabled())

        with self.assertWarns(UserWarning):
            enable_progress_bars()
        self.assertTrue(are_progress_bars_disabled())  # Still disabled !

    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", False)
    def test_cannot_disable_tqdm_when_env_variable_is_set(self) -> None:
        """
        Test helpers cannot enable/disable progress bars when
        `HF_HUB_DISABLE_PROGRESS_BARS` is set.
        """
        enable_progress_bars()
        self.assertFalse(are_progress_bars_disabled())

        with self.assertWarns(UserWarning):
            disable_progress_bars()
        self.assertFalse(are_progress_bars_disabled())  # Still enabled !

    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", None)
    def test_tqdm_disabled(self) -> None:
        """Test TQDM not outputting anything when globally disabled."""
        disable_progress_bars()
        for _ in tqdm(range(10)):
            pass

        captured = self.capsys.readouterr()
        self.assertEqual(captured.out, "")
        self.assertEqual(captured.err, "")

    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", None)
    def test_tqdm_disabled_cannot_be_forced(self) -> None:
        """Test TQDM cannot be forced when globally disabled."""
        disable_progress_bars()
        for _ in tqdm(range(10), disable=False):
            pass

        captured = self.capsys.readouterr()
        self.assertEqual(captured.out, "")
        self.assertEqual(captured.err, "")

    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", None)
    def test_tqdm_can_be_disabled_when_globally_enabled(self) -> None:
        """Test TQDM can still be locally disabled even when globally enabled."""
        enable_progress_bars()
        for _ in tqdm(range(10), disable=True):
            pass

        captured = self.capsys.readouterr()
        self.assertEqual(captured.out, "")
        self.assertEqual(captured.err, "")

    @patch("huggingface_hub.utils._tqdm.HF_HUB_DISABLE_PROGRESS_BARS", None)
    def test_tqdm_enabled(self) -> None:
        """Test TQDM work normally when globally enabled."""
        enable_progress_bars()
        for _ in tqdm(range(10)):
            pass

        captured = self.capsys.readouterr()
        self.assertEqual(captured.out, "")
        self.assertIn("10/10", captured.err)  # tqdm log

    def test_tqdm_stream_file(self) -> None:
        with SoftTemporaryDirectory() as tmpdir:
            filepath = Path(tmpdir) / "config.json"
            with filepath.open("w") as f:
                f.write("#" * 1000)

            with tqdm_stream_file(filepath) as f:
                while True:
                    data = f.read(100)
                    if not data:
                        break
                    time.sleep(0.001)  # Simulate a delay between each chunk

            captured = self.capsys.readouterr()
            self.assertEqual(captured.out, "")
            self.assertIn("config.json: 100%", captured.err)  # log file name
            self.assertIn("|█████████", captured.err)  # tqdm bar
            self.assertIn("1.00k/1.00k", captured.err)  # size in B
