"""
pytest tests/test_progiter.py
"""

from __future__ import annotations

import itertools as it
import sys
import typing
from io import StringIO

from xdoctest.utils import CaptureStdout, strip_ansi

from ubelt import ProgIter


class FakeStream:
    """
    Helper to hook into and introspect when progiter writes to the display
    """

    def __init__(
        self,
        verbose: int = 0,
        callback: typing.Callable[[], None] | None = None,
    ) -> None:
        self.verbose = verbose
        self.callback = callback
        self._callcount = 0
        self.messages: list[str] = []

    def write(self, msg: str) -> None:
        self._callcount += 1
        self.messages.append(msg)
        if self.verbose:
            sys.stdout.write(msg)
        if self.callback is not None:
            self.callback()

    def flush(self, *args: object, **kw: object) -> None: ...


class FakeTimer:
    """
    Helper to hook into and introspect when progiter measures times.
    You must tic this timer yourself.
    """

    def __init__(self, times: typing.Iterable[int] = (1,)) -> None:
        self._time: float | int = 0
        self._callcount = 0
        self._iter = it.cycle(times)

    def tic(self, step: float | int | None = None) -> None:
        if step is None:
            step = next(self._iter)
        delta: float | int = step
        self._time += delta

    def __call__(self) -> float | int:
        self._callcount += 1
        return self._time


def test_rate_format_string() -> None:
    # Less of a test than a demo
    rates = [1 * 10**i for i in range(-10, 10)]

    texts = []
    for rate in rates:
        rate_format = '4.2f' if rate > 0.001 else 'g'
        # Really cool: you can embed format strings inside format strings
        msg = '{rate:{rate_format}}'.format(rate=rate, rate_format=rate_format)
        texts.append(msg)
    assert texts == [
        '1e-10',
        '1e-09',
        '1e-08',
        '1e-07',
        '1e-06',
        '1e-05',
        '0.0001',
        '0.001',
        '0.01',
        '0.10',
        '1.00',
        '10.00',
        '100.00',
        '1000.00',
        '10000.00',
        '100000.00',
        '1000000.00',
        '10000000.00',
        '100000000.00',
        '1000000000.00',
    ]


def test_rate_format() -> None:
    # Define a function that takes some time
    file = StringIO()
    prog: typing.Any = ProgIter(file=file)
    prog.begin()

    prog._iters_per_second = 0.000001
    msg = prog.format_message()
    rate_part = msg.split('rate=')[1].split(' Hz')[0]
    assert rate_part == '1e-06'

    prog._iters_per_second = 0.1
    msg = prog.format_message()
    rate_part = msg.split('rate=')[1].split(' Hz')[0]
    assert rate_part == '0.10'

    prog._iters_per_second = 10000
    msg = prog.format_message()
    rate_part = msg.split('rate=')[1].split(' Hz')[0]
    assert rate_part == '10000.00'


def test_progiter() -> None:
    # Define a function that takes some time
    def is_prime(n: int) -> bool:
        return n >= 2 and not any(n % i == 0 for i in range(2, n))

    N = 500

    if False:
        file = StringIO()
        prog = ProgIter(
            range(N), clearline=False, file=file, freq=N // 10, adjust=False
        )
        file.seek(0)
        print(file.read())

        prog = ProgIter(range(N), clearline=False)
        for n in prog:
            was_prime = is_prime(n)
            prog.set_extra(
                'n=%r, was_prime=%r'
                % (
                    n,
                    was_prime,
                )
            )
            if (n + 1) % 128 == 0 and was_prime:
                prog.set_extra(
                    'n=%r, was_prime=%r EXTRA'
                    % (
                        n,
                        was_prime,
                    )
                )
        file.seek(0)
        print(file.read())

    total = 200
    N = 5000
    N0 = N - total
    print('N = %r' % (N,))
    print('N0 = %r' % (N0,))

    print('\n-----')
    print('Demo #0: progress can be disabled and incur essentially 0 overhead')
    print(
        'However, the overhead of enabled progress is minimal and typically '
        'insignificant'
    )
    print('this is verbosity mode verbose=0')
    sequence = (is_prime(n) for n in range(N0, N))
    if True:
        psequence = ProgIter(sequence, total=total, desc='demo0', enabled=False)
        list(psequence)

    print('\n-----')
    print('Demo #1: progress is shown by default in the same line')
    print('this is verbosity mode verbose=1')
    sequence = (is_prime(n) for n in range(N0, N))
    if True:
        psequence = ProgIter(sequence, total=total, desc='demo1')
        list(psequence)

    # Default behavior adjusts frequency of progress reporting so
    # the performance of the loop is minimally impacted
    print('\n-----')
    print('Demo #2: clearline=False prints multiple lines.')
    print('Progress is only printed as needed')
    print('Notice the adjustment behavior of the print frequency')
    print('this is verbosity mode verbose=2')
    if True:
        sequence = (is_prime(n) for n in range(N0, N))
        psequence = ProgIter(
            sequence, total=total, clearline=False, desc='demo2'
        )
        list(psequence)
        # import utool as ut
        # print(ut.repr4(psequence.__dict__))

    print('\n-----')
    print('Demo #3: Adjustments can be turned off to give constant feedback')
    print('this is verbosity mode verbose=3')
    sequence = (is_prime(n) for n in range(N0, N))
    if True:
        psequence = ProgIter(
            sequence,
            total=total,
            adjust=False,
            clearline=False,
            freq=100,
            desc='demo3',
        )
        list(psequence)


def test_progiter_offset_10() -> None:
    """
    pytest -s  tests/test_progiter.py::test_progiter_offset_10
    """
    # Define a function that takes some time
    file = StringIO()
    list(
        ProgIter(
            range(10),
            total=20,
            verbose=3,
            start=10,
            file=file,
            freq=5,
            show_rate=False,
            show_eta=False,
            show_total=False,
            time_thresh=0,
        )
    )
    file.seek(0)
    want = ['50.00% 10/20...', '75.00% 15/20...', '100.00% 20/20...']
    got = [line.strip() for line in file.readlines()]
    if sys.platform.startswith('win32'):  # nocover
        # on windows \r seems to be mixed up with ansi sequences
        from xdoctest.utils import strip_ansi

        got = [strip_ansi(line).strip() for line in got]
    assert got == want


def test_progiter_offset_0() -> None:
    """
    pytest -s  tests/test_progiter.py::test_progiter_offset_0
    """
    # Define a function that takes some time
    file = StringIO()
    for _ in ProgIter(
        range(10),
        total=20,
        verbose=3,
        start=0,
        file=file,
        freq=5,
        show_rate=False,
        show_eta=False,
        show_total=False,
        time_thresh=0,
    ):
        pass
    file.seek(0)
    want = ['0.00%  0/20...', '25.00%  5/20...', '50.00% 10/20...']
    got = [line.strip() for line in file.readlines()]
    if sys.platform.startswith('win32'):  # nocover
        # on windows \r seems to be mixed up with ansi sequences
        from xdoctest.utils import strip_ansi

        got = [strip_ansi(line).strip() for line in got]
    assert got == want


def test_unknown_total() -> None:
    """
    Make sure a question mark is printed if the total is unknown
    """
    iterable = (_ for _ in range(0, 10))
    file = StringIO()
    prog = ProgIter(
        iterable, desc='unknown seq', file=file, show_times=False, verbose=1
    )
    for n in prog:
        pass
    file.seek(0)
    got = [line.strip() for line in file.readlines()]
    # prints an eroteme if total is unknown
    assert len(got) > 0, 'should have gotten something'
    assert all('?' in line for line in got), 'all lines should have an eroteme'


def test_initial() -> None:
    """
    Make sure a question mark is printed if the total is unknown
    """
    file = StringIO()
    prog: typing.Any = ProgIter(
        initial=9001, file=file, show_times=False, clearline=False
    )
    message = prog.format_message_parts()[1]
    assert strip_ansi(message) == ' 9001/?... '


def test_clearline() -> None:
    """
    Make sure a question mark is printed if the total is unknown

    pytest tests/test_progiter.py::test_clearline
    """
    file = StringIO()
    # Clearline=False version should simply have a newline at the end.
    prog: typing.Any = ProgIter(file=file, show_times=False, clearline=False)
    before, message, after = prog.format_message_parts()
    assert before == ''
    assert strip_ansi(message).strip(' ') == '0/?...'
    # Clearline=True version should carriage return at the beginning and have no
    # newline at the end.
    prog = ProgIter(file=file, show_times=False, clearline=True)
    before, message, after = prog.format_message_parts()
    assert before == '\r'
    assert strip_ansi(message).strip(' ') == '0/?...'


def test_disabled() -> None:
    prog = ProgIter(range(20), enabled=True)
    prog.begin()
    assert prog.started

    prog = ProgIter(range(20), enabled=False)
    prog.begin()
    prog.step()
    assert not prog.started


def test_eta_window_None() -> None:
    # nothing to check (that I can think of) run test for coverage
    prog = ProgIter(range(20), enabled=True, eta_window=None)  # type: ignore
    for _ in prog:
        pass


def test_adjust_freq() -> None:
    # nothing to check (that I can think of) run test for coverage
    prog: typing.Any = ProgIter(
        range(20),
        enabled=True,
        eta_window=typing.cast(typing.Any, None),
        rel_adjust_limit=4.0,  # type: ignore
    )

    # Adjust frequency up to have each update happen every 1sec or so
    prog.freq = 1
    prog.time_thresh = 1.0
    prog._max_between_count = -1.0  # type: ignore
    prog._max_between_time = -1.0  # type: ignore
    prog._measure_timedelta = 1
    prog._measure_countdelta = 1000
    prog._adjust_frequency()
    assert prog.freq == 4

    # Adjust frequency down to have each update happen every 1sec or so
    prog.freq = 1000
    prog.time_thresh = 1.0
    prog._max_between_count = -1.0
    prog._max_between_time = -1.0
    prog._measure_timedelta = 1
    prog._measure_countdelta = 1
    prog._adjust_frequency()
    assert prog.freq == 250

    # No need to adjust frequency to have each update happen every 1sec or so
    prog.freq = 1
    prog.time_thresh = 1.0
    prog._max_between_count = -1.0
    prog._max_between_time = -1.0
    prog._measure_timedelta = 1
    prog._measure_countdelta = 1
    prog._adjust_frequency()
    assert prog.freq == 1


def test_tqdm_compatibility() -> None:
    prog = ProgIter(range(20), total=20, miniters=17, show_times=False)
    assert prog.pos == 0
    assert prog.freq == 17
    for _ in prog:
        pass

    with CaptureStdout() as cap:
        ProgIter.write('foo')
    assert cap.text is not None
    assert cap.text.strip() == 'foo'

    with CaptureStdout() as cap:
        prog = ProgIter(show_times=False)
        prog.set_description('new desc', refresh=False)
        prog.begin()
        prog.refresh()
        prog.close()
    assert prog.label == 'new desc'
    assert cap.text is not None
    assert 'new desc' in cap.text.strip()

    with CaptureStdout() as cap:
        prog = ProgIter(show_times=False)
        prog.set_description('new desc', refresh=True)
        prog.close()
    assert prog.label == 'new desc'
    assert cap.text is not None
    assert 'new desc' in cap.text.strip()

    with CaptureStdout() as cap:
        prog = ProgIter(show_times=False)
        prog.set_description_str('new desc')
        prog.begin()
        prog.refresh()
        prog.close()
    assert prog.label == 'new desc'
    assert cap.text is not None
    assert 'new desc' in cap.text.strip()

    with CaptureStdout() as cap:
        prog = ProgIter(show_times=False)
        prog.set_postfix({'foo': 'bar'}, baz='biz', x=object(), y=2)
        prog.begin()
    assert prog.length is None
    assert cap.text is not None
    assert 'foo=bar' in cap.text.strip()
    assert 'baz=biz' in cap.text.strip()
    assert 'y=2' in cap.text.strip()
    assert 'x=<object' in cap.text.strip()

    with CaptureStdout() as cap:
        prog = ProgIter(show_times=False)
        prog.set_postfix_str('bar baz', refresh=False)
    assert cap.text is not None
    assert 'bar baz' not in cap.text.strip()

    with CaptureStdout() as cap:
        prog = ProgIter(show_times=False)
        prog.set_postfix('bar baz', refresh=False)
    assert cap.text is not None
    assert 'bar baz' not in cap.text.strip()


class IntObject:
    def __init__(self) -> None:
        self.n = 0

    def inc(self, *args: typing.Any, **kwargs: typing.Any) -> None:
        self.n += 1


def test_adjust_fast_early_slow_late_doesnt_get_stuck() -> None:
    cnt = IntObject()
    fake_stream = FakeStream(verbose=0, callback=cnt.inc)
    fake_timer = FakeTimer()

    prog = ProgIter(
        range(1000),
        enabled=True,
        adjust=True,
        time_thresh=1.0,
        rel_adjust_limit=1000000.0,
        homogeneous=False,
        timer=fake_timer,
        stream=fake_stream,  # type: ignore
    )
    it = iter(prog)
    # Few fast updates at the beginning
    for i in range(10):
        next(it)
        fake_timer.tic(100)
    # Followed by some extremely slow updates
    for i in range(10):
        next(it)
        fake_timer.tic(0.00001)
    # Outputs should not get stuck at the few fast updates
    assert cnt.n > 8


def test_adjust_slow_early_fast_late_doesnt_spam() -> None:
    cnt = IntObject()
    fake_stream = FakeStream(verbose=0, callback=cnt.inc)
    fake_timer = FakeTimer()

    prog = ProgIter(
        range(1000),
        enabled=True,
        adjust=True,
        time_thresh=1.0,
        rel_adjust_limit=1000000.0,
        homogeneous=False,
        timer=fake_timer,
        stream=fake_stream,  # type: ignore
    )
    it = iter(prog)
    # Few slow updates at the beginning
    for i in range(10):
        next(it)
        fake_timer.tic(100)
    # Followed by a ton of extremely fast updates
    for i in range(990):
        next(it)
        fake_timer.tic(0.00001)
    # Outputs should not spam the screen with messages
    assert cnt.n < 20


def test_homogeneous_heuristic_with_iter_lengths() -> None:
    for size in range(0, 10):
        list(ProgIter(range(size), homogeneous='auto'))


def test_mixed_iteration_and_step() -> None:
    # Check to ensure nothing breaks
    for adjust in [0, 1]:
        for homogeneous in [0, 1] if adjust else [0]:
            for size in range(0, 10):
                for n_inner_steps in range(size):
                    prog = ProgIter(
                        range(size),
                        adjust=typing.cast(typing.Any, adjust),
                        homogeneous=homogeneous,  # type: ignore
                    )
                    iprog = iter(prog)
                    try:
                        while True:
                            next(iprog)
                            for k in range(n_inner_steps):
                                prog.step()
                    except StopIteration:
                        ...


def check_issue_32_non_homogeneous_time_threshold_prints() -> None:
    """
    xdoctest ~/code/progiter/tests/test_progiter.py check_issue_32_non_homogeneous_time_threshold_prints
    """
    from ubelt import ProgIter

    fake_stream = FakeStream(verbose=1)
    fake_timer = FakeTimer([10, 1, 30, 40, 3, 4, 10, 10, 10, 10, 10, 10])
    time_thresh = 50
    # fake_timer = FakeTimer([.5 * factor])
    # time_thresh = 2.9 * factor

    N = 20
    prog = ProgIter(
        range(N),
        timer=fake_timer,
        time_thresh=time_thresh,
        homogeneous='auto',
        stream=fake_stream,  # type: ignore
        clearline=False,
    )

    static_state = {
        'time_thresh': prog.time_thresh,
        'adjust': prog.adjust,
        'homogeneous': prog.homogeneous,
    }

    states = []

    def record_state() -> dict[str, typing.Any]:
        real_display_timedelta = (
            fake_timer._time - prog._display_measurement.time
        )
        state = {
            'iter_idx': prog._iter_idx,
            'next_idx': prog._next_measure_idx,
            'time': fake_timer._time,
            'freq': prog.freq,
            'curr': prog._curr_measurement,
            'disp': prog._display_measurement,
            'meas_td': prog._measure_timedelta,
            'disp_td': prog._display_timedelta,
            'real_disp_td': real_display_timedelta,
            'n_disp': fake_stream._callcount,
            'n_times': fake_timer._callcount,
        }
        states.append(state)
        return state

    _iter = iter(prog)
    prog.begin()
    record_state()

    for _ in range(N):
        next(_iter)
        record_state()
        fake_timer.tic()

    assert fake_stream._callcount == len(fake_stream.messages)

    prog.end()
    record_state()

    try:
        import pandas as pd
        import rich

        import ubelt as ub

        print(
            'fake_stream.messages = {}'.format(
                ub.urepr(fake_stream.messages, nl=1)
            )
        )
        print(f'prog._likely_homogeneous={prog._likely_homogeneous}')
        rich.print(pd.Series(static_state))
        df = pd.DataFrame(states)
        df['displayed'] = df['n_disp'].diff().astype(bool)
        df['timed'] = df['n_times'].diff().astype(bool)
        rich.print(df.to_string())
    except ImportError:
        ...

    # TODO: write actual asserts that check that displays, measurements, and
    # adjustments happen at the write times


def test_end_message_is_displayed() -> None:
    """
    Older versions of progiter had a bug where the end step would not trigger
    if calculations were updated without a display
    """
    import io

    stream = io.StringIO()
    prog = ProgIter(range(1000), stream=stream)
    for i in prog:
        ...
    stream.seek(0)
    text = stream.read()
    assert '1000/1000' in text, 'end message should have printed'


def test_standalone_display() -> None:
    from ubelt import ProgIter

    fake_stream = FakeStream(verbose=1)
    fake_timer = FakeTimer()
    time_thresh = 50

    N = 20
    prog = ProgIter(
        range(N),
        timer=fake_timer,
        time_thresh=time_thresh,
        homogeneous=True,
        stream=fake_stream,  # type: ignore
        clearline=True,
    )

    prog.begin()

    _iter = iter(prog)

    prog.display_message()
    prog.display_message()
    prog.display_message()
    fake_timer.tic(1)
    prog.display_message()

    next(_iter)
    prog.display_message()
    prog.display_message()

    fake_timer.tic(1)
    next(_iter)
    fake_timer.tic(1)
    next(_iter)
    fake_timer.tic(1)
    next(_iter)
    prog.display_message()

    assert fake_stream.messages == [
        '\r 0.00%  0/20... rate=0 Hz, eta=?, total=0:00:00',
        '\r 0.00%  0/20... rate=0 Hz, eta=?, total=0:00:00',
        '\r 0.00%  0/20... rate=0 Hz, eta=?, total=0:00:00',
        '\r 0.00%  0/20... rate=0 Hz, eta=?, total=0:00:00',
        '\r 0.00%  0/20... rate=0 Hz, eta=?, total=0:00:00',
        '\r 5.00%  1/20... rate=1.00 Hz, eta=0:00:19, total=0:00:01',
        '\r 5.00%  1/20... rate=1.00 Hz, eta=0:00:19, total=0:00:01',
        '\r 20.00%  4/20... rate=1.00 Hz, eta=0:00:16, total=0:00:04',
    ]


def test_no_percent() -> None:
    from ubelt import ProgIter

    fake_stream = FakeStream(verbose=1)
    fake_timer = FakeTimer()
    time_thresh = 50

    N = 20
    prog = ProgIter(
        range(N),
        timer=fake_timer,
        time_thresh=time_thresh,
        show_percent=False,
        homogeneous=True,
        stream=fake_stream,  # type: ignore
        clearline=True,
    )

    prog.begin()

    _iter = iter(prog)

    prog.display_message()
    prog.display_message()
    prog.display_message()
    fake_timer.tic(1)
    prog.display_message()

    next(_iter)
    prog.display_message()
    prog.display_message()

    fake_timer.tic(1)
    next(_iter)
    fake_timer.tic(1)
    next(_iter)
    fake_timer.tic(1)
    next(_iter)
    prog.display_message()
    assert fake_stream.messages == [
        '\r  0/20... rate=0 Hz, eta=?, total=0:00:00',
        '\r  0/20... rate=0 Hz, eta=?, total=0:00:00',
        '\r  0/20... rate=0 Hz, eta=?, total=0:00:00',
        '\r  0/20... rate=0 Hz, eta=?, total=0:00:00',
        '\r  0/20... rate=0 Hz, eta=?, total=0:00:00',
        '\r  1/20... rate=1.00 Hz, eta=0:00:19, total=0:00:01',
        '\r  1/20... rate=1.00 Hz, eta=0:00:19, total=0:00:01',
        '\r  4/20... rate=1.00 Hz, eta=0:00:16, total=0:00:04',
    ]


def test_clearline_padding() -> None:
    """
    Ensure we overwrite the entire previous message
    """
    from ubelt import ProgIter

    fake_stream = FakeStream(verbose=1)
    prog = ProgIter(
        range(20),
        time_thresh=99999999,
        show_percent=False,
        homogeneous=True,
        stream=fake_stream,  # type: ignore
        clearline=True,
    )
    prog.start()
    prog.display_message()
    msg1_len = prog._prev_msg_len
    assert prog._prev_msg_len > 30
    assert prog._prev_msg_len < 50
    assert prog.clearline, 'test requirement'
    prog.set_extra('a very long message')
    prog.step()
    assert prog._prev_msg_len == msg1_len, (
        'We are under the time threshold. '
        'We should not have updated the display message yet'
    )
    prog.display_message()
    msg2_len = prog._prev_msg_len
    assert msg2_len > msg1_len, 'should have a longer message'

    # Now make a shorter line length
    prog.set_extra('shorter')
    prog.step()
    prog.display_message()
    msg3_len = prog._prev_msg_len
    assert msg3_len < msg2_len, 'should have a shorter message'

    msg1 = fake_stream.messages[-3]
    msg2 = fake_stream.messages[-2]
    msg3 = fake_stream.messages[-1]
    assert len(msg1) == msg1_len + 1
    assert len(msg2) == msg2_len + 1
    assert len(msg3) >= msg3_len + 1, (
        'the real third message should include padding'
    )
    assert len(msg3) == msg2_len + 1, (
        'the real third message should include padding to clear msg2'
    )


def test_extra_callback() -> None:
    from ubelt import ProgIter

    fake_stream = FakeStream(verbose=1)
    fake_timer = FakeTimer()
    time_thresh = 50

    def build_extra() -> str:
        return chr(prog._iter_idx % 26 + 97) * 3

    N = 20
    prog: typing.Any = ProgIter(
        range(N),
        timer=fake_timer,
        time_thresh=time_thresh,
        homogeneous=True,
        stream=fake_stream,  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
        clearline=True,
    )
    prog.set_extra(build_extra)

    prog.begin()

    _iter = iter(prog)

    prog.display_message()
    prog.display_message()
    prog.display_message()
    fake_timer.tic(1)
    prog.display_message()

    next(_iter)
    prog.display_message()
    prog.display_message()

    fake_timer.tic(1)
    next(_iter)
    fake_timer.tic(1)
    next(_iter)
    fake_timer.tic(1)
    next(_iter)
    prog.display_message()

    assert fake_stream.messages == [
        '\r 0.00%  0/20...aaa rate=0 Hz, eta=?, total=0:00:00',
        '\r 0.00%  0/20...aaa rate=0 Hz, eta=?, total=0:00:00',
        '\r 0.00%  0/20...aaa rate=0 Hz, eta=?, total=0:00:00',
        '\r 0.00%  0/20...aaa rate=0 Hz, eta=?, total=0:00:00',
        '\r 0.00%  0/20...aaa rate=0 Hz, eta=?, total=0:00:00',
        '\r 5.00%  1/20...bbb rate=1.00 Hz, eta=0:00:19, total=0:00:01',
        '\r 5.00%  1/20...bbb rate=1.00 Hz, eta=0:00:19, total=0:00:01',
        '\r 20.00%  4/20...eee rate=1.00 Hz, eta=0:00:16, total=0:00:04',
    ]


if __name__ == '__main__':
    import pytest

    pytest.main([__file__])
