# This Python file uses the following encoding: utf-8
# Licensed under a 3-clause BSD style license - see LICENSE.rst


import pytest
import numpy as np
from io import StringIO

from astropy import table
from astropy.io import ascii
from astropy.table import Table, QTable
from astropy.table.table_helpers import simple_table
from astropy import units as u
from astropy.utils import console

BIG_WIDE_ARR = np.arange(2000, dtype=np.float64).reshape(100, 20)
SMALL_ARR = np.arange(18, dtype=np.int64).reshape(6, 3)


@pytest.mark.usefixtures('table_type')
class TestMultiD():

    def test_multidim(self, table_type):
        """Test printing with multidimensional column"""
        arr = [np.array([[1, 2],
                         [10, 20]], dtype=np.int64),
               np.array([[3, 4],
                         [30, 40]], dtype=np.int64),
               np.array([[5, 6],
                         [50, 60]], dtype=np.int64)]
        t = table_type(arr)
        lines = t.pformat()
        assert lines == ['col0 [2] col1 [2] col2 [2]',
                         '-------- -------- --------',
                         '  1 .. 2   3 .. 4   5 .. 6',
                         '10 .. 20 30 .. 40 50 .. 60']

        lines = t.pformat(html=True)
        assert lines == [
            f'<table id="table{id(t)}">',
            '<thead><tr><th>col0 [2]</th><th>col1 [2]</th><th>col2 [2]</th></tr></thead>',
            '<tr><td>1 .. 2</td><td>3 .. 4</td><td>5 .. 6</td></tr>',
            '<tr><td>10 .. 20</td><td>30 .. 40</td><td>50 .. 60</td></tr>',
            '</table>']
        nbclass = table.conf.default_notebook_table_class
        masked = 'masked=True ' if t.masked else ''
        assert t._repr_html_().splitlines() == [
            f'<div><i>{table_type.__name__} {masked}length=2</i>',
            f'<table id="table{id(t)}" class="{nbclass}">',
            '<thead><tr><th>col0 [2]</th><th>col1 [2]</th><th>col2 [2]</th></tr></thead>',
            '<thead><tr><th>int64</th><th>int64</th><th>int64</th></tr></thead>',
            '<tr><td>1 .. 2</td><td>3 .. 4</td><td>5 .. 6</td></tr>',
            '<tr><td>10 .. 20</td><td>30 .. 40</td><td>50 .. 60</td></tr>',
            '</table></div>']

        t = table_type([arr])
        lines = t.pformat()
        assert lines == ['col0 [2,2]',
                         '----------',
                         '   1 .. 20',
                         '   3 .. 40',
                         '   5 .. 60']

    def test_fake_multidim(self, table_type):
        """Test printing with 'fake' multidimensional column"""
        arr = [np.array([[(1,)],
                         [(10,)]], dtype=np.int64),
               np.array([[(3,)],
                         [(30,)]], dtype=np.int64),
               np.array([[(5,)],
                         [(50,)]], dtype=np.int64)]
        t = table_type(arr)
        lines = t.pformat()
        assert lines == ['col0 [1,1] col1 [1,1] col2 [1,1]',
                         '---------- ---------- ----------',
                         '         1          3          5',
                         '        10         30         50']

        lines = t.pformat(html=True)
        assert lines == [
            f'<table id="table{id(t)}">',
            '<thead><tr><th>col0 [1,1]</th><th>col1 [1,1]</th><th>col2 [1,1]</th></tr></thead>',
            '<tr><td>1</td><td>3</td><td>5</td></tr>',
            '<tr><td>10</td><td>30</td><td>50</td></tr>',
            '</table>']
        nbclass = table.conf.default_notebook_table_class
        masked = 'masked=True ' if t.masked else ''
        assert t._repr_html_().splitlines() == [
            f'<div><i>{table_type.__name__} {masked}length=2</i>',
            f'<table id="table{id(t)}" class="{nbclass}">',
            '<thead><tr><th>col0 [1,1]</th><th>col1 [1,1]</th><th>col2 [1,1]</th></tr></thead>',
            '<thead><tr><th>int64</th><th>int64</th><th>int64</th></tr></thead>',
            '<tr><td>1</td><td>3</td><td>5</td></tr>', '<tr><td>10</td><td>30</td><td>50</td></tr>',
            '</table></div>']

        t = table_type([arr])
        lines = t.pformat()
        assert lines == ['col0 [2,1,1]',
                         '------------',
                         '     1 .. 10',
                         '     3 .. 30',
                         '     5 .. 50']


def test_html_escaping():
    t = table.Table([('<script>alert("gotcha");</script>', 2, 3)])
    nbclass = table.conf.default_notebook_table_class
    assert t._repr_html_().splitlines() == [
        '<div><i>Table length=3</i>',
        f'<table id="table{id(t)}" class="{nbclass}">',
        '<thead><tr><th>col0</th></tr></thead>',
        '<thead><tr><th>str33</th></tr></thead>',
        '<tr><td>&lt;script&gt;alert(&quot;gotcha&quot;);&lt;/script&gt;</td></tr>',
        '<tr><td>2</td></tr>',
        '<tr><td>3</td></tr>',
        '</table></div>']


@pytest.mark.usefixtures('table_type')
class TestPprint():

    def _setup(self, table_type):
        self.tb = table_type(BIG_WIDE_ARR)
        self.tb['col0'].format = 'e'
        self.tb['col1'].format = '.6f'

        self.tb['col0'].unit = 'km**2'
        self.tb['col19'].unit = 'kg s m**-2'
        self.ts = table_type(SMALL_ARR)

    def test_empty_table(self, table_type):
        t = table_type()
        lines = t.pformat()
        assert lines == ['<No columns>']
        c = repr(t)
        masked = 'masked=True ' if t.masked else ''
        assert c.splitlines() == [f'<{table_type.__name__} {masked}length=0>',
                                  '<No columns>']

    def test_format0(self, table_type):
        """Try getting screen size but fail to defaults because testing doesn't
        have access to screen (fcntl.ioctl fails).
        """
        self._setup(table_type)
        arr = np.arange(4000, dtype=np.float64).reshape(100, 40)
        lines = table_type(arr).pformat()
        nlines, width = console.terminal_size()
        assert len(lines) == nlines
        for line in lines[:-1]:  # skip last "Length = .. rows" line
            assert width - 10 < len(line) <= width

    def test_format1(self, table_type):
        """Basic test of formatting, unit header row included"""
        self._setup(table_type)
        lines = self.tb.pformat(max_lines=8, max_width=40)
        assert lines == ['    col0         col1    ...   col19  ',
                         '    km2                  ... kg s / m2',
                         '------------ ----------- ... ---------',
                         '0.000000e+00    1.000000 ...      19.0',
                         '         ...         ... 
...       ...',
                         '1.960000e+03 1961.000000 ...    1979.0',
                         '1.980000e+03 1981.000000 ...
    1999.0',
                         'Length = 100 rows']

    def test_format2(self, table_type):
        """Basic test of formatting, unit header row excluded"""
        self._setup(table_type)
        lines = self.tb.pformat(max_lines=8, max_width=40, show_unit=False)
        assert lines == ['    col0         col1    ... col19 ',
                         '------------ ----------- ... ------',
                         '0.000000e+00    1.000000 ...   19.0',
                         '2.000000e+01   21.000000 ...   39.0',
                         '         ...         ... ...    ...',
                         '1.960000e+03 1961.000000 ... 1979.0',
                         '1.980000e+03 1981.000000 ... 1999.0',
                         'Length = 100 rows']

    def test_format3(self, table_type):
        """Include the unit header row"""
        self._setup(table_type)
        lines = self.tb.pformat(max_lines=8, max_width=40, show_unit=True)

        assert lines == ['    col0         col1    ...   col19  ',
                         '    km2                  ... kg s / m2',
                         '------------ ----------- ... ---------',
                         '0.000000e+00    1.000000 ...      19.0',
                         '         ...         ... ...       ...',
                         '1.960000e+03 1961.000000 ...    1979.0',
                         '1.980000e+03 1981.000000 ...    1999.0',
   