import sys
import os
from argparse import ArgumentParser
from pathlib import Path
import json
import iso8601
from datetime import datetime


def get_format(header: dict) -> dict:
    t = header.get("type", "general")
    if t == "general":
        return None
    elif t == "number":
        if "format" in header:
            return {"num_format": header["format"]}
    elif t == "text":
        return None
    elif t == "date":
        return {"num_format": header.get("format") or "yyyy-mm-dd"}
    elif t == "datetime":
        return {"num_format": header.get("format") or "yyyy-mm-dd hh:mm:ss"}
    elif t == "time":
        return {"num_format": header.get("format") or "hh:mm:ss"}
    elif t == "formula":
        if header.get("format"):
            return {"num_format": header.get("format")}
    return None


def write_xlsx(output_filename: str, data: dict):
    import xlsxwriter
    workbook = xlsxwriter.Workbook(output_filename, {'constant_memory': True})
    headers_format = workbook.add_format({'bold': True})
    for ws in data["worksheets"]:
        worksheet = workbook.add_worksheet(name=ws["name"])
        worksheet.remove_timezone = True
        cols_count = len(ws["columns"])
        formats = []
        row = 0
        col = 0
        # emit headers
        for header in ws["columns"]:
            worksheet.write(row, col, header["title"])
            worksheet.set_column(col, col, 20)
            f = get_format(header)
            if f:
                formats.append(workbook.add_format(f))
            else:
                formats.append(None)
            col += 1
        worksheet.set_row(0, cell_format=headers_format)

        # emit data
        row += 1
        for row_data in ws["data"]:
            if len(row_data) != cols_count:
                raise Exception(f"Row {row} must be {cols_count} element long")
            col = 0
            for cell in row_data:
                curr_col = t = ws["columns"][col]["title"]
                t = ws["columns"][col].get("type")
                ff = formats[col]
                if (not t) and ff:
                    raise Exception(f"Column [{curr_col}] error: Cannot set a format without setting the column type")

                if t == "number":
                    if ff:
                        worksheet.write_number(row, col, float(cell), ff)
                    else:
                        worksheet.write_number(row, col, float(cell))
                elif t in ("datetime", "date"):
                    if not (cell is None or cell == ""):
                        try:
                            dt = iso8601.parse_date(cell)
                            if ff:
                                worksheet.write_datetime(row, col, dt, ff)
                            else:
                                worksheet.write_datetime(row, col, dt)
                        except:
                            worksheet.write_string(row, col, "invalid")
                elif t == "time":
                    if not (cell is None or cell == ""):
                        try:
                            tm = datetime.strptime(cell, "%H:%M:%S")
                            if ff:
                                worksheet.write_datetime(row, col, tm, ff)
                            else:
                                worksheet.write_datetime(row, col, tm)
                        except:
                            worksheet.write_string(row, col, "invalid")
                elif t == "boolean":
                    if ff:
                        worksheet.write_boolean(row, col, bool(cell), ff)
                    else:
                        worksheet.write_boolean(row, col, bool(cell))
                elif t == "formula":
                    if ff:
                        worksheet.write_formula(row, col, cell, ff)
                    else:
                        worksheet.write_formula(row, col, cell)
                elif t == "text":
                    if cell is None:
                        cell = ""
                    if ff:
                        worksheet.write_string(row, col, cell, ff)
                    else:
                        worksheet.write_string(row, col, str(cell))
                else:  # default
                    worksheet.write(row, col, str(cell))
                col += 1  # for cell
            row += 1  # for row_data

    workbook.close()


def main():
    parser = ArgumentParser(prog='python jsontoexcel.py')
    parser.add_argument('jsonfile', type=str, help="The JSON file")
    params = parser.parse_args()

    datafile_path: Path = Path(params.jsonfile)
    if not datafile_path.is_file():
        raise Exception("JSONFile not found [{}]".format(params.jsonfile))

    output_filename = datafile_path.with_suffix(".xlsx")
    if output_filename.is_file():
        os.unlink(output_filename)
        # raise Exception("Output file already exists [{}]".format(output_filename))
    in_json = json.load(open(datafile_path))
    write_xlsx(output_filename, in_json)


if __name__ == '__main__':
    main()
