# Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-


import io

import pytest

from . import test_progress_bar_func
from astropy.utils import console
from astropy import units as u


class FakeTTY(io.StringIO):
    """IOStream that fakes a TTY; provide an encoding to emulate an output
    stream with a specific encoding.
    """

    def __new__(cls, encoding=None):
        # Return a new subclass of FakeTTY with the requested encoding
        if encoding is None:
            return super().__new__(cls)

        encoding = encoding
        cls = type(encoding.title() + cls.__name__, (cls,),
                   {'encoding': encoding})

        return cls.__new__(cls)

    def __init__(self, encoding=None):
        super().__init__()

    def write(self, s):
        if isinstance(s, bytes):
            # Just allow this case to work
            s = s.decode('latin-1')
        elif self.encoding is not None:
            s.encode(self.encoding)

        return super().write(s)

    def isatty(self):
        return True


def test_fake_tty():
    # First test without a specified encoding; we should be able to write
    # arbitrary unicode strings
    f1 = FakeTTY()
    assert f1.isatty()
    f1.write('☃')
    assert f1.getvalue() == '☃'

    # Now test an ASCII-only TTY--it should raise a UnicodeEncodeError when
    # trying to write a string containing non-ASCII characters
    f2 = FakeTTY('ascii')
    assert f2.isatty()
    assert f2.__class__.__name__ == 'AsciiFakeTTY'
    assert pytest.raises(UnicodeEncodeError, f2.write, '☃')
    assert f2.getvalue() == ''


@pytest.mark.skipif("sys.platform.startswith('win')")
def test_color_text():
    assert console._color_text("foo", "green") == '\033[0;32mfoo\033[0m'


def test_color_print():
    # This stuff is hard to test, at least smoke test it
    console.color_print("foo", "green")

    console.color_print("foo", "green", "bar", "red")


def test_color_print2():
    # Test that this automatically detects that io.StringIO is
    # not a tty
    stream = io.StringIO()
    console.color_print("foo", "green", file=stream)
    assert stream.getvalue() == 'foo\n'

    stream = io.StringIO()
    console.color_print("foo", "green", "bar", "red", "baz", file=stream)
    assert stream.getvalue() == 'foobarbaz\n'


@pytest.mark.skipif("sys.platform.startswith('win')")
def test_color_print3():
    # Test that this thinks the FakeTTY is a tty and applies colors.

    stream = FakeTTY()
    console.color_print("foo", "green", file=stream)
    assert stream.getvalue() == '\x1b[0;32mfoo\x1b[0m\n'

    stream = FakeTTY()
    console.color_print("foo", "green", "bar", "red", "baz", file=stream)
    assert stream.getvalue() == '\x1b[0;32mfoo\x1b[0m\x1b[0;31mbar\x1b[0mbaz\n'


def test_color_print_unicode():
    console.color_print("überbær", "red")


def test_color_print_invalid_color():
    console.color_print("foo", "unknown")


def test_spinner_non_unicode_console():
    """Regression test for #1760

    Ensures that the spinner can fall go into fallback mode when using the
    unicode spinner on a terminal whose default encoding cannot encode the
    unicode characters.
    """

    stream = FakeTTY('ascii')
    chars = console.Spinner._default_unicode_chars

    with console.Spinner("Reticulating splines", file=stream,
                         chars=chars) as s:
        next(s)


def test_progress_bar():
    # This stuff is hard to test, at least smoke test it
    with console.ProgressBar(50) as bar:
        for i in range(50):
            bar.update()


def test_progress_bar2():
    for x in console.ProgressBar(range(50)):
        pass


def test_progress_bar3():
    def do_nothing(*args, **kwargs):
        pass

    console.ProgressBar.map(do_nothing, range(50))


def test_zero_progress_bar():
    with console.ProgressBar(0) as bar:
        pass


def test_progress_bar_as_generator():
    sum = 0
    for x in console.ProgressBar(range(50)):
        sum += x
    assert sum == 1225

    sum = 0
    for x in console.ProgressBar(50):
        sum += x
    assert sum == 1225


def test_progress_bar_map():
    items = list(range(100))
    result = console.ProgressBar.map(test_progress_bar_func.func,
                                     items, step=10, multiprocess=True)
    assert items == result

    result1 = console.ProgressBar.map(test_progress_bar_func.func,
                                      items, step=10, multiprocess=2)

    assert items == result1


@pytest.mark.parametrize(("seconds", "string"),
       [(864088, " 1w 3d"),
       (187213, " 2d 4h"),
       (3905, " 1h 5m"),
       (64, " 1m 4s"),
       (15, "   15s"),
       (2, "    2s")]
)
def test_human_time(seconds, string):
    human_time = console.human_time(seconds)
    assert human_time == string


@pytest.mark.parametrize(("size", "string"),
       [(8640882, "8.6M"),
       (187213, "187k"),
       (3905, "3.9k"),
       (64, " 64 "),
       (2, "  2 "),
       (10*u.GB, " 10G")]
)
def test_human_file_size(size, string):
    human_time = console.human_file_size(size)
    assert human_time == string


@pytest.mark.parametrize("size", (50*u.km, 100*u.g))
def test_bad_human_file_size(size):
    assert pytest.raises(u.UnitConversionError, console.human_file_size, size)