"""Helpers to assemble and render tabular data."""

__all__ = ["Column", "RenderTarget", "Table"]

from abc import ABCMeta, abstractmethod
import collections
from collections.abc import Iterable
import csv
import enum
from io import StringIO
import json
from os import linesep

import colorclass
import terminaltables
import yaml


class RenderTarget(enum.Enum):

    # Human-readable tabluar output.
    plain = "plain"
    pretty = "pretty"

    # Machine-readable output.
    yaml = "yaml"
    json = "json"
    csv = "csv"

    def __str__(self):
        # Return the value. This makes it better to use as a choices option to
        # argparse at the command-line.
        return self.value


class Table(metaclass=ABCMeta):
    def __init__(self, *columns, visible_columns=None):
        super(Table, self).__init__()
        self.columns = collections.OrderedDict(
            (column.name, column) for column in columns
        )
        if visible_columns is None:
            self.visible_columns = collections.OrderedDict(self.columns.items())
        else:
            self.visible_columns = collections.OrderedDict(
                (column.name, column)
                for column in columns
                if column.name in visible_columns
            )

    def __getitem__(self, name):
        return self.columns[name]

    @abstractmethod
    def get_rows(self, target, data):
        """Get the rows for the table."""

    def _filter_rows(self, rows):
        """Filter `rows` based on the visible columns."""
        filtered_rows = []
        for row in rows:
            filtered_row = []
            for idx, name in enumerate(self.columns.keys()):
                if name in self.visible_columns:
                    filtered_row.append(row[idx])
            filtered_rows.append(filtered_row)
        return filtered_rows

    def render(self, target, data):
        """Render the table."""
        rows = self.get_rows(target, data)
        rows = self._filter_rows(rows)
        renderer = getattr(self, "_render_%s" % target.name, None)
        if renderer is None:
            raise ValueError("Cannot render %r for %s." % (self.value, target))
        else:
            return renderer(rows)

    def _flatten_columns(self, columns):
        cols = []
        for _, column in columns.items():
            if isinstance(column, NestedTableColumn):
                cols.extend(self._flatten_columns(column.get_columns()))
            else:
                cols.append(column)
        return cols

    def _compute_rows(self, target, data, duplicate=False):
        columns = self.visible_columns.values()
        computed = []
        for row_data in data:
            row = []
            rows = [row]
            for datum, column in zip(row_data, columns):
                if isinstance(column, NestedTableColumn):
                    nested_rows = column.get_rows(target, datum, duplicate=duplicate)
                    orig_row = list(row)
                    row.extend(nested_rows[0])
                    for nested_row in nested_rows[1:]:
                        new_row = list(orig_row)
                        if not duplicate:
                            for idx in range(len(new_row)):
                                new_row[idx] = ""
                        new_row.extend(nested_row)
                        rows.append(new_row)
                else:
                    row.append(column.render(target, datum))
            computed.extend(rows)
        return computed

    def _render_plain(self, data):
        rows = self._compute_rows(RenderTarget.plain, data)
        rows.insert(
            0, [column.title for column in self._flatten_columns(self.visible_columns)]
        )
        return terminaltables.AsciiTable(rows).table

    def _render_pretty(self, data):
        rows = self._compute_rows(RenderTarget.pretty, data)
        rows.insert(
            0, [column.title for column in self._flatten_columns(self.visible_columns)]
        )
        return terminaltables.SingleTable(rows).table

    def _render_yaml(self, data):
        columns = self.visible_columns.values()
        rows = [
            [
                column.render(RenderTarget.yaml, datum)
                for datum, column in zip(row, columns)
            ]
            for row in data
        ]
        return yaml.safe_dump(
            {
                "columns": [
                    {"name": column.name, "title": column.title} for column in columns
                ],
                "data": [
                    {column.name: datum for column, datum in zip(columns, row)}
                    for row in rows
                ],
            },
            default_flow_style=False,
        ).rstrip(linesep)

    def _render_json(self, data):
        columns = self.visible_columns.values()
        rows = [
            [
                column.render(RenderTarget.json, datum)
                for datum, column in zip(row, columns)
            ]
            for row in data
        ]
        return json.dumps(
            {
                "columns": [
                    {"name": column.name, "title": column.title} for column in columns
                ],
                "data": [
                    {column.name: datum for column, datum in zip(columns, row)}
                    for row in rows
                ],
            }
        )

    def _render_csv(self, data):
        output = StringIO()
        writer = csv.writer(output)
        writer.writerow(
            [column.name for column in self._flatten_columns(self.visible_columns)]
        )
        writer.writerows(self._compute_rows(RenderTarget.csv, data, duplicate=True))
        return output.getvalue().rstrip(linesep)

    def __repr__(self):
        return "<%s [%s]>" % (self.__class__.__name__, " ".join(self.visible_columns))


class DetailTable(Table):
    def _filter_rows(self, rows, visible_columns=None):
        """Filter `rows` based on the visible columns."""
        if visible_columns is None:
            visible_columns = self.visible_columns
        filtered_row = []
        for idx, name in enumerate(self.columns.keys()):
            if name in self.visible_columns:
                filtered_row.append(rows[idx])
        return filtered_row

    def render(self, target, data):
        renderer = getattr(self, "_render_%s" % target.name, None)
        if renderer is None:
            raise ValueError("Cannot render %r for %s." % (self.value, target))
        else:
            return renderer(data)

    def _split_nested_tables(self):
        table_columns = collections.OrderedDict(
            (name, column)
            for name, column in self.visible_columns.items()
            if isinstance(column, NestedTableColumn)
        )
        data_columns = collections.OrderedDict(
            (name, column)
            for name, column in self.visible_columns.items()
            if not isinstance(column, NestedTableColumn)
        )
        return data_columns, table_columns

    def _render_nested_table(self, target, data, column):
        data_idx = list(self.columns.keys()).index(column.name)
        data = data[data_idx]
        table = column.get_table()
        return table.render(target, data)

    def _render_nested_tables(self, target, data, columns):
        tables = [
            self._render_nested_table(target, data, column)
            for _, column in columns.items()
        ]
        if len(tables) > 0:
            return "\n" + "\n".join(tables)
        else:
            return ""

    def _render_table(self, target, terminaltable, data):
        columns, table_columns = self._split_nested_tables()
        all_rows = self.get_rows(target, data)
        rows = self._filter_rows(all_rows, visible_columns=columns)
        rows = [
            column.render(target, datum)
            for column, datum in zip(columns.values(), rows)
        ]
        table = terminaltable(
            [(column.title, datum) for column, datum in zip(columns.values(), rows)]
        )
        table.inner_heading_row_border = False
        return table.table + self._render_nested_tables(target, all_rows, table_columns)

    def _render_plain(self, data):
        return self._render_table(RenderTarget.plain, terminaltables.AsciiTable, data)

    def _render_pretty(self, data):
        return self._render_table(RenderTarget.pretty, terminaltables.SingleTable, data)

    def _render_yaml(self, data):
        columns = self.visible_columns.values()
        rows = self.get_rows(RenderTarget.yaml, data)
        rows = self._filter_rows(rows)
        rows = [
            column.render(RenderTarget.yaml, datum)
            for column, datum in zip(columns, rows)
        ]
        return yaml.safe_dump(
            {column.name: datum for column, datum in zip(columns, rows)},
            default_flow_style=False,
        ).rstrip(linesep)

    def _render_json(self, data):
        columns = self.visible_columns.values()
        rows = self.get_rows(RenderTarget.json, data)
        rows = self._filter_rows(rows)
        rows = [
            column.render(RenderTarget.json, datum)
            for column, datum in zip(columns, rows)
        ]
        return json.dumps({column.name: datum for column, datum in zip(columns, rows)})

    def _render_csv(self, data):
        columns, table_columns = self._split_nested_tables()
        all_rows = self.get_rows(RenderTarget.csv, data)
        rows = self._filter_rows(all_rows, visible_columns=columns)
        rows = [
            column.render(RenderTarget.csv, datum)
            for column, datum in zip(columns.values(), rows)
        ]
        output = StringIO()
        writer = csv.writer(output)
        writer.writerows(
            [(column.name, datum) for column, datum in zip(columns.values(), rows)]
        )
        return output.getvalue().rstrip(linesep) + (
            self._render_nested_tables(RenderTarget.csv, all_rows, table_columns)
        )


class Column:
    def __init__(self, name, title=None):
        super(Column, self).__init__()
        self.name = name
        self.title = name if title is None else title

    def render(self, target, datum):
        if target is RenderTarget.yaml:
            return datum
        elif target is RenderTarget.json:
            return datum
        elif target is RenderTarget.csv:
            if isinstance(datum, Iterable) and not isinstance(datum, (str, bytes)):
                return ",".join(datum)
            else:
                return datum
        elif target is RenderTarget.plain:
            if datum is None:
                return ""
            elif isinstance(datum, colorclass.Color):
                return datum.value_no_colors
            elif isinstance(datum, Iterable) and not isinstance(datum, (str, bytes)):
                return "\n".join(datum)
            else:
                return str(datum)
        elif target is RenderTarget.pretty:
            if datum is None:
                return ""
            elif isinstance(datum, colorclass.Color):
                return datum
            elif isinstance(datum, Iterable) and not isinstance(datum, (str, bytes)):
                return "\n".join(datum)
            else:
                return str(datum)
        else:
            raise ValueError("Cannot render %r for %s" % (datum, target))

    def __repr__(self):
        return "<%s name=%s title=%r>" % (
            self.__class__.__name__,
            self.name,
            self.title,
        )


class NestedTableColumn(Column):
    def __init__(
        self, name, title=None, table=None, table_args=None, table_kwargs=None
    ):
        super(NestedTableColumn, self).__init__(name, title=title)
        self.table = table
        self.table_args = table_args
        self.table_kwargs = table_kwargs
        if table is None:
            raise ValueError("table is required.")

    def get_table(self):
        table_args = self.table_args
        if table_args is None:
            table_args = []
        table_kwargs = self.table_kwargs
        if table_kwargs is None:
            table_kwargs = {}
        return self.table(*table_args, **table_kwargs)

    def get_columns(self):
        return self.get_table().visible_columns

    def get_rows(self, target, data, duplicate=False):
        table = self.get_table()
        rows = table.get_rows(target, data)
        rows = table._filter_rows(rows)
        rows = table._compute_rows(target, rows, duplicate=duplicate)
        if len(rows) == 0:
            # Nested table column must always return one row even if its
            # an empty row.
            rows = [[" " for _ in range(len(table.visible_columns))]]
        return rows

    def render(self, target, datum):
        table = self.get_table()
        if target is RenderTarget.yaml:
            return yaml.safe_load(table.render(target, datum))
        elif target is RenderTarget.json:
            return json.loads(table.render(target, datum))
        else:
            raise ValueError(
                "Should not be called on a nested table column, the "
                "table render should handle this correctly."
            )
