#!/usr/bin/env python3
# Yonghang Wang

import sys
import argparse
import os
import re
import csv
import copy
from wcwidth import wcswidth
from collections import defaultdict
from itertools import zip_longest

class SimpleTable:
    def __init__(self, data=None, header=None, cols=None, maxwidth=-1):
        self.__maxwidth = int(maxwidth)
        if cols and re.search(",", cols):
            ncols = len(header)
            xmap = [int(i) for i in re.split(r",", cols) if int(i) < ncols]
            self.__header = [header[i] for i in xmap]
            self.__data = list()
            for d in data:
                self.__data.append([d[i] for i in xmap])
        else:
            self.__header = header or list()
            self.__data = data or list()
        self.__num_of_cols = len(self.__header)

    def import_csv(self, csvfile, header=True):
        pass

    def __len__(self):
        return len(self.__data)

    def __getitem__(self, pos):
        return self.__data[pos]

    def __add__(self, other):
        if self.__num_of_cols != other.__num_of_cols:
            return None
        return SimpleTable(self.__data + other.__data, self.__header)

    def __iadd__(self, other):
        if self.__num_of_cols != other.__num_of_cols:
            return None
        self.__data += other.__data
 
    def get_json(self) :
        import json
        tbl = list()
        for r in self.__data :
            row = dict()
            for ix, val in enumerate(r) :
                row[self.__header[ix]] = val
            tbl.append(row)
        return json.dumps(tbl,indent=2)

    def get_yaml(self) :
        import yaml
        tbl = list()
        for r in self.__data :
            row = dict()
            for ix, val in enumerate(r) :
                row[self.__header[ix]] = val
            tbl.append(row)
        return yaml.safe_dump(tbl,default_flow_style=False)

    def get_csv(self) :
        import io
        si = io.StringIO()
        w = csv.writer(si)
        if self.__header :
            w.writerow(self.__header)
        for r in self.__data :
            w.writerow(r)
        return si.getvalue()

    def repr_pivot(self):
        maxwidth = max([wcswidth(h) for h in self.__header])
        fmtstr = "{:" + str(maxwidth) + "} : {}"
        res = ""
        for r, row in enumerate(self.__data):
            if r > 0:
                res += "--\n"
            for i, c in enumerate(row):
                res += fmtstr.format(self.__header[i], c) + "\n"
        return res

    def __wcswidth_x(self,s):
        res = 0
        for ln in s.splitlines() :
            wclen = wcswidth(ln)
            if wclen > res :
                res = wclen
        return res

    def __splitstring(self,s,maxwidth) :
        result = list()
        for ln in s.splitlines() :
            if wcswidth(ln) <= maxwidth :
                result.append(ln)
            else :
                left = ln
                while left :
                    start = 0
                    end = min(maxwidth,len(left))
                    while wcswidth(left[start:end]) > maxwidth :
                        end -= 1
                    result.append(left[start:end])
                    left = left[end:]
        return result

    def __splitrow(self,row) :
        if self.__maxwidth < 20 :
            return zip_longest(*[c.splitlines() for c in row],fillvalue="")
        else :
            return zip_longest(*[self.__splitstring(str(c),self.__maxwidth) for c in row],fillvalue="")

    def __repr__(self):
        width = [0 for _ in range(len(self.__header))]
        for row in self.__data + [self.__header]:
            for ix, col in enumerate(row):
                if not col :
                    col = ""
                if ix > len(width) - 1:
                    break
                wclen = self.__wcswidth_x(str(col)) 
                if wclen > width[ix]:
                    width[ix] = wclen
        if self.__maxwidth >= 20 :     
            width = [min(w,self.__maxwidth) for w in width]
        twidth = copy.copy(width)
        for ix, w in enumerate(twidth):
            twidth[ix] = w - wcswidth(str(self.__header[ix])) + len(str(self.__header[ix]))
        fmtstr = "".join(["{:" + str(l + 1) + "}" for l in twidth])
        res = fmtstr.format(*self.__header).rstrip() + "\n"
        res += ('-' * (sum(width) + len(width) - 1)) + "\n"
        for r in self.__data:
            row = [str(c) if c else "" for c in r]
            if len(row) < len(width):
                row.extend([""] * (len(width) - len(row)))
            for t in self.__splitrow(row) :
                twidth = copy.copy(width)
                for ix, w in enumerate(twidth):
                    twidth[ix] = w - wcswidth(str(t[ix])) + len(str(t[ix]))
                fmtstr = "".join(["{:" + str(l + 1) + "}" for l in twidth])
                res += fmtstr.format(*t).rstrip() + "\n"
        return res

def tokenize(s):
    offset = 0
    tokens = list()
    while True:
        m = re.search(r"\S+\s*", s[offset:])
        if m:
            tokens.append(
                [offset + m.start(), offset + m.end(),
                 m.group().rstrip()])
            offset += m.end()
        else:
            break
    tokens[-1][1] = -1
    return tokens


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-n',
                        '--lineno',
                        dest='lineno',
                        default=False,
                        action='store_true',
                        help='show line no.')
    parser.add_argument('-H', '--header', dest='header', help='header columns')
    parser.add_argument('-f', '--infile', dest='infile', help='input file')
    parser.add_argument('-c',
                        '--column',
                        action='append',
                        dest='column',
                        help='column names. used when there'
                        're spaces within.')
    parser.add_argument(
        '-C',
        '--dump-column',
        dest='dumpcols',
        help='only print columns indentified by index numbers.')
    parser.add_argument('-s',
                        '--sortby',
                        dest='sortby',
                        help='column id starts with 0.')
    parser.add_argument('-b',
                        '--sep-char',
                        dest='sepchar',
                        default="\s+",
                        help='char to seperate columns')
    parser.add_argument('-w',
                        '--maxwidth',
                        dest='maxwidth',
                        type=int,
                        default=-1,
                        help='max col width when print in console, min==20')
    parser.add_argument(
        '-t',
        '--table',
        dest='table',
        action='store_true',
        default=False,
        help='input preformatted by spaces. header should not include spaces.')
    parser.add_argument('-p',
                        '--pivot',
                        dest='pivot',
                        action='store_true',
                        default=False,
                        help='pivot wide tables.')
    parser.add_argument('-F', '--format', dest='format', help='other than console table, json,yaml or csv')
    parser.add_argument('-X',
                        '--debug',
                        dest='debug',
                        action='store_true',
                        default=False,
                        help='debug mode')
    args = parser.parse_args()

    colsdict = defaultdict()
    colsdict_revert = defaultdict()
    if args.column:
        for c in args.column:
            oc = c
            mc = re.sub("\s", "_", oc)
            colsdict[oc] = mc
            colsdict_revert[mc] = oc
    if args.header and colsdict:
        for c, nc in colsdict.items():
            args.header = re.sub(c, nc, args.header)
    xheader = None
    header = list()
    if args.lineno:
        header.append('#')
    if args.header:
        mhdr = args.header
        if not args.table:
            header.extend(re.split(r"{}".format(args.sepchar), mhdr))
        else:
            xheader = tokenize(mhdr)
            header = [v.rstrip() for (s, t, v) in xheader]
    data = list()
    lno = 0
    INPUT = sys.stdin
    if args.infile:
        with open(args.infile, "r") as f:
            INPUT = f.readlines()
    for ln in INPUT:
        if not ln or not re.search(r"\S+",ln):
            continue
        if not header:
            for c, nc in colsdict.items():
                ln = re.sub(c, nc, ln)
        if not args.table:
            arr = re.split(r"{}".format(args.sepchar), ln.rstrip())
        else:
            if xheader:
                arr = [ln[s:t].rstrip() for (s, t, v) in xheader]
            else:
                for c, nc in colsdict.items():
                    ln = re.sub(c, nc, ln)
                xheader = tokenize(ln)
                if args.lineno:
                    header = ["#"] + [v.rstrip() for (s, t, v) in xheader]
                else:
                    header = [v.rstrip() for (s, t, v) in xheader]
                continue
        if len(header) == 0 or (args.lineno and len(header) == 1):
            if not args.table:
                header.extend(arr)
        else:
            lno += 1
            if args.lineno:
                data.append([str(lno)] + arr)
            else:
                data.append(arr)
    oheader = list()
    for h in header:
        nh = colsdict_revert.get(h, h)
        oheader.append(nh)
    if args.sortby:
        def fsort(x):
            v = list()
            for i in re.split(r",", args.sortby):
                v0 = x[int(i)] or ""
                if re.match(r"^\d+(\.\d*)*$",v0) :
                    v.append(float(v0))
                else :
                    v.append(0)
                v.append(v0)
            return v
        data = sorted(data, key=fsort)
    if args.pivot:
        print(
            SimpleTable(header=oheader, data=data,
                        cols=args.dumpcols).repr_pivot())
    else:
        t = SimpleTable(header=oheader, data=data, cols=args.dumpcols,maxwidth=args.maxwidth)
        if args.format == "json" :
            print(t.get_json())
        elif args.format == "yaml" :
            print(t.get_yaml())
        elif args.format == "csv" :
            print(t.get_csv())
        else :
            print(t,end="")


if __name__ == "__main__":
    main()
