import re
import textwrap
import neovim
from neovim.api import NvimError


def mapl(f, x):
    return list(map(f, x))


def line_has_sep(line):
    """Line has a `-` before a `=`
    """
    a = line.find('-')   # not header
    b = line.find('=')   # header

    if a == -1: # No `-`
        return False
    elif b == -1: # No `=`, but `-`
        return True
    else:
        return a < b


def table_has_header(chunk):
    if len(chunk) == 1:
        return False

    try:
        if line_has_sep(chunk[2]) or line_has_sep(chunk[1]):
            return False
    except IndexError:
        return True

    return True


def join_rows(rows, sep='\n'):
    """Given a list of rows (a list of lists) this function returns a
    flattened list where each the individual columns of all rows are joined
    together using the line separator.

    """
    output = []
    for row in rows:
        # grow output array, if necessary
        if len(output) <= len(row):
            for i in range(len(row) - len(output)):
                output.extend([[]])

        for i, field in enumerate(row):
            field_text = field.strip()
            if field_text:
                output[i].append(field_text)
    return mapl(lambda lines: sep.join(lines), output)


def line_is_separator(line):
    return re.match('^[\t +=-]+$', line)


def has_line_seps(raw_lines):
    for line in raw_lines:
        if line_is_separator(line):
            return True
    return False


def partition_raw_lines(raw_lines):
    """Partitions a list of raw input lines so that between each partition, a
    table row separator can be placed.

    """
    if not has_line_seps(raw_lines):
        return mapl(lambda x: [x], raw_lines)

    curr_part = []
    parts = [curr_part]
    for line in raw_lines:
        if line_is_separator(line):
            curr_part = []
            parts.append(curr_part)
        else:
            curr_part.append(line)

    # remove any empty partitions (typically the first and last ones)
    return filter(lambda x: x != [], parts)


def unify_table(table):
    """Given a list of rows (i.e. a table), this function returns a new table
    in which all rows have an equal amount of columns.  If all full column is
    empty (i.e. all rows have that field empty), the column is removed.

    """
    max_fields = max(map(lambda row: len(row), table))
    empty_cols = [True] * max_fields
    output = []
    for row in table:
        curr_len = len(row)
        if curr_len < max_fields:
            row += [''] * (max_fields - curr_len)
        output.append(row)

        # register empty columns (to be removed at the end)
        for i in range(len(row)):
            if row[i].strip():
                empty_cols[i] = False

    # remove empty columns from all rows
    table = output
    output = []
    for row in table:
        cols = []
        for i in range(len(row)):
            should_remove = empty_cols[i]
            if not should_remove:
                cols.append(row[i])
        output.append(cols)

    return output


def split_table_row(row_string):
    if row_string.find("|") >= 0:
        # first, strip off the outer table drawings
        row_string = re.sub(r'^\s*\||\|\s*$', '', row_string)
        return re.split(r'\s*\|\s*', row_string.strip())
    return re.split(r'\s\s+', row_string.rstrip())


def parse_table(raw_lines):
    row_partition = partition_raw_lines(raw_lines)
    lines = mapl(lambda row_string: join_rows(mapl(split_table_row, row_string)),
                row_partition)
    return unify_table(lines)


def table_line(widths, header=False):
    if header:
        linechar = '='
    else:
        linechar = '-'
    sep = '+'
    parts = []
    for width in widths:
        parts.append(linechar * width)
    if parts:
        parts = [''] + parts + ['']
    return sep.join(parts)


def str_width(unicode_text):
    """calc string width, support cjk characters."""
    from unicodedata import east_asian_width
    return sum(1+(east_asian_width(c) in "WF") for c in unicode_text)


def split_row_into_lines(row):
    row = mapl(lambda field: field.split('\n'), row)
    height = max(map(lambda field_lines: len(field_lines), row))
    turn_table = []
    for i in range(height):
        fields = []
        for field_lines in row:
            if i < len(field_lines):
                fields.append(field_lines[i])
            else:
                fields.append('')
        turn_table.append(fields)
    return turn_table


def get_column_widths_from_border_spec(chunk):
    border = None
    for row in chunk:
        if line_is_separator(row):
            border = row.strip()
            break

    if border is None:
        raise RuntimeError('Cannot reflow this table. Top table border not found.')

    left = right = None
    if border[0] == '+':
        left = 1
    if border[-1] == '+':
        right = -1
    return mapl(lambda drawing: max(0, len(drawing) - 2), border[left:right].split('+'))


def get_indent(line):
    return line[0 : len(line)-len(line.lstrip())]


def apply_indent(table, indent):
    for i in range(len(table)):
        table[i] = indent + table[i]
    return table


def get_field_width(field_text):
    return max(map(lambda s: str_width(s), field_text.split('\n')))


def get_column_widths(table):
    widths = []
    for row in table:
        num_fields = len(row)
        # dynamically grow
        if num_fields >= len(widths):
            widths.extend([0] * (num_fields - len(widths)))
        for i in range(num_fields):
            field_width = get_field_width(row[i])
            widths[i] = max(widths[i], field_width)
    return widths


def pad_fields(row, widths):
    """Pads fields of the given row, so each field lines up nicely with the
    others.
    """
    # Pad all fields using the calculated widths
    new_row = []
    for i in range(len(row)):
        unicode_len = str_width(row[i])
        col = ' ' + row[i] + ' ' * int(widths[i] - unicode_len + 1)
        new_row.append(col)
    return new_row


def wrap_text(text, width):
    """wrap text, support cjk characters."""
    lines = []
    while len(text) > 0:
        w = width
        # check 1st string, if too wide, then guess again;
        guess = textwrap.wrap(text, w)[0]
        while str_width(guess) > width:
            w -= (str_width(guess) - width + 1)/2
            guess = textwrap.wrap(text, w)[0]
        lines.append(guess)
        text = text[len(guess):].strip()
    return lines


def reflow_row_contents(row, widths):
    new_row = []
    for i, field in enumerate(row):
        wrapped_lines = wrap_text(field.replace('\n', ' '), widths[i])
        new_row.append("\n".join(wrapped_lines))
    return new_row


def draw_table(table, manual_widths=None, header=True):
    if table == []:
        return []

    if manual_widths is None:
        col_widths = get_column_widths(table)
    else:
        col_widths = manual_widths
        new_widths = get_column_widths(table)
        if len(new_widths) > len(col_widths):
            col_widths += new_widths[len(col_widths):]

    # Reserve room for the spaces
    sep_col_widths = mapl(lambda x: x + 2, col_widths)
    header_line = table_line(sep_col_widths, header=True)
    normal_line = table_line(sep_col_widths, header=False)

    output = [normal_line]
    first = True
    for row in table:

        if manual_widths:
            row = reflow_row_contents(row, manual_widths)

        row_lines = split_row_into_lines(row)

        # draw the lines (num_lines) for this row
        for row_line in row_lines:
            row_line = pad_fields(row_line, col_widths)
            output.append("|".join([''] + row_line + ['']))

        # then, draw the separator
        if first and header:
            output.append(header_line)
            first = False
        else:
            output.append(normal_line)

    return output


@neovim.plugin
class Main(object):
    def __init__(self, vim):
        self.vim = vim


    def get_table_bounds(self):
        row, col = self.vim.current.window.cursor
        upper = lower = row
        try:
            while self.vim.current.buffer[upper - 1].strip():
                upper -= 1
        except (IndexError, NvimError):
            pass
        else:
            upper += 1

        try:
            while self.vim.current.buffer[lower - 1].strip():
                lower += 1
        except (IndexError, NvimError):
            pass
        else:
            lower -= 1

        return (upper, lower)


    @neovim.command('TableRstFormat')
    def reformat_table(self):
        upper, lower = self.get_table_bounds()
        chunk = self.vim.current.buffer[upper - 1:lower]
        indent = get_indent(chunk[0])
        table = parse_table(chunk)
        has_header = table_has_header(chunk)
        chunk = draw_table(table, header=has_header)
        self.vim.current.buffer[upper - 1:lower] = apply_indent(chunk, indent)


    @neovim.command('TableRstReflow')
    def reflow_table(self):
        upper, lower = self.get_table_bounds()
        chunk = self.vim.current.buffer[upper - 1:lower]
        indent = get_indent(chunk[0])
        table = parse_table(chunk)
        widths = get_column_widths_from_border_spec(chunk)
        table = parse_table(chunk)
        has_header = table_has_header(chunk)
        chunk = draw_table(table, widths, header=has_header)
        self.vim.current.buffer[upper - 1:lower] = apply_indent(chunk, indent)
