#!/usr/bin/env python3

import argparse
import copy
import difflib
import json
import os
import pprint
import re
import sys
import textwrap

from pathlib import Path

import common

debug = True

CONSTANTS = {
 'bus_naming_style': '%s[%d]',
 'comp_attribute,voltage_map': {'VGND': 0.0, 'VNB': 0.0},
 'comp_attribute,capacitive_load_unit': [1.0, 'pf'],
 'comp_attribute,library_features': 'report_delay_calculation',
 'comp_attribute,technology': 'cmos',
 'current_unit': '1mA',
 'default_cell_leakage_power': 0.0,
 'default_inout_pin_cap': 0.0,
 'default_input_pin_cap': 0.0,
 'default_leakage_power_density': 0.0,
 'default_output_pin_cap': 0.0,
 'delay_model': 'table_lookup',
 'in_place_swap_mode': 'match_footprint',
 'input_threshold_pct_fall': 50.0,
 'input_threshold_pct_rise': 50.0,
 'nom_process': 1.0,
 'output_threshold_pct_fall': 50.0,
 'output_threshold_pct_rise': 50.0,
 'pulling_resistance_unit': '1kohm',
 'revision': 1.0,
 'simulation': 'true',
 'slew_lower_threshold_pct_fall': 20.0,
 'slew_lower_threshold_pct_rise': 20.0,
 'slew_upper_threshold_pct_fall': 80.0,
 'slew_upper_threshold_pct_rise': 80.0,
 'time_unit': '1ns',
 'voltage_unit': '1V',
}


def flatten(input, current, output):
    if isinstance(input, list):
        for i in range(0, len(input)):
            ik = "{}[{}]".format(current, i)
            flatten(input[i], ik, output)
        return

    if not isinstance(input, dict):
        assert current not in output, current
        output[current] = input
        return

    assert isinstance(input, dict), input
    for k in input:
        if current:
            ok = current+"."+k
        else:
            ok = k
        flatten(input[k], ok, output)


CONSTANTS_FLAT = {}
flatten(CONSTANTS, "", CONSTANTS_FLAT)


class SetEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, set):
            return list(obj)
        return json.JSONEncoder.default(self, obj)


def otemp(i):
    """

    >>> otemp(100)
    '100C'

    >>> otemp(-40)
    'n40C'

    """
    s = ''
    s = s + str(i)
    if i < 0:
        assert s.startswith('-'), s
        s = 'n' + s[1:]
    if len(s) < 3:
        s = '0'+s
    s = s + 'C'
    return s


def ovolt(f):
    """

    >>> ovolt(3.3)
    '3v30'

    >>> ovolt(33.0)
    '33v00'

    >>> ovolt(-0.96)
    'n0v96'
    """
    p = "{:1.2f}".format(f).replace('.', 'v')
    if f < 0:
        assert p.startswith('-'), p
        p = 'n'+p[1:]
    return p


IGNORING = ['lrhc', 'tpl', 'tpnl', 'io', 'pwr', 'lkg']


def calculate_out_filename(fname, data, ignore_ccsn=False):
    """

    >>> calculate_out_filename('/home/tim/gob/foss-eda-tools/skywater-src-nda/scs8hdll/V0.1.0/lib/scs8hdll_ff_xx_1p95lv_hv_io_100c_lrhc_tpl_PVT4.lib', {'nom_temperature': 100, 'nom_voltage': 1.20})
    'ff_100C_1v20'

    >>> calculate_out_filename('/home/tim/gob/foss-eda-tools/skywater-src-nda/scs8hs/V0.0.0/lib/scs8hs_tt_1.80v_25C_ccsnoise.lib', {'nom_temperature': 25, 'nom_voltage': 1.80, 'ccsn': True})
    'tt_025C_1v80_ccsnoise'

    >>> calculate_out_filename('scs8hs_tt_1.20v_30C_ccsnoise.lib', {'temperature': 30, 'voltage': 1.20, 'ccsn': True})
    'tt_030C_1v20_ccsnoise'

    >>> calculate_out_filename('scs8hs_tt_1.20v_30C_ccsnoise.lib', {'temperature': 30, 'voltage': 1.20, 'ccsn': True})
    'tt_030C_1v20_ccsnoise'

    >>> calculate_out_filename('scs8hs_tt_1.20v_30C.lib', {'temperature': 30, 'voltage': 1.20, 'ccsn': True})
    Traceback (most recent call last):
      ...
    SystemError: *Found* ccsn attributes (['ccsn']) but *no* ccsnoise in filename: scs8hs_tt_1.20v_30C.lib

    >>> calculate_out_filename('scs8hs_tt_1.20v_30C_ccsnoise.lib', {'temperature': 30, 'voltage': 1.20})
    Traceback (most recent call last):
      ...
    SystemError: *No* ccsn attributes found (['temperature', 'voltage']) but ccsnoise in filename: scs8hs_tt_1.20v_30C_ccsnoise.lib

    """
    assert isinstance(fname, str), repr(fname)
    assert isinstance(data, dict), repr(data)

    nom_temperature = None
    if 'nom_temperature' in data:
        nom_temperature = int(data['nom_temperature'])
    elif 'temperature' in data:
        nom_temperature = int(data['temperature'])
    assert nom_temperature is not None, data

    nom_voltage = None
    if 'nom_voltage' in data:
        nom_voltage = float(data['nom_voltage'])
    elif 'temperature' in data:
        nom_voltage = float(data['voltage'])
    assert nom_voltage is not None, data

    supplies_nom = set()
    supplies_vmap = {}
    if "comp_attribute,voltage_map" in data:
        vmap_src = data["comp_attribute,voltage_map"]
        if isinstance(vmap_src, dict):
            vmap_src = vmap_src.items()

        for k, v in vmap_src:
            if v == 0.0:
                continue

            if v == nom_voltage:
                supplies_nom.add(k)
            else:
                supplies_vmap[k] = v

    basename = os.path.basename(fname)
    bits = basename.replace('pwr_lkg', 'pwrlkg').split("_")
    if '.' in bits[-1]:
        bits[-1], fext = bits[-1].split('.', 1)
        assert fext in ('lib.json', 'lib'), (basename, fext, bits)

    flags = []
    ftypes = []
    while len(bits) > 0:
        b = bits.pop(0)
        if 'scs8' in b:
            continue

        if b in ('ccsnoise','pwrlkg'):
            flags.append(b)
            continue

        if b in ('ff', 'ss', 'tt', 'xx', 'diode'):
            ftypes.append(b)
            continue

        if b.upper().endswith('C'):
            continue

        if b.lower().endswith('v'):
            continue

        if b in IGNORING or 'PVT' in b:
            continue

        assert False, (b, bits, fname)

    if not ignore_ccsn:
        is_ccsnoise = any('ccsn' in k for k in data.keys())
        if 'ccsnoise' in flags:
            if not is_ccsnoise:
                raise SystemError('*No* ccsn attributes found ({}) but ccsnoise in filename: {}'.format(
                    [k for k in data.keys()], fname))
        else:
            if is_ccsnoise:
                raise SystemError('*Found* ccsn attributes ({}) but *no* ccsnoise in filename: {}'.format(
                    [k for k in data.keys() if 'ccsn' in k], fname))

    if 'xx' in ftypes:
        ftypes.remove('xx')

    if len(ftypes) == 2:
        ftype = 'hv{}_lv{}'.format(*ftypes)
    else:
        assert len(ftypes) == 1, (ftypes, fname)
        ftype = ftypes[0]


    o = "{}_{}_{}".format(ftype, otemp(nom_temperature), ovolt(nom_voltage))

    if supplies_vmap:
        o += "_" + "_".join("{}{}".format(k.lower().replace('pwr', ''), ovolt(v)) for k, v in supplies_vmap.items())

    if flags:
        o += "_" + "_".join(flags)

    return o


RE_CELL = re.compile('[A-Za-z_0-9]+')

def rewrite_function(f, cname):
    """

    #>>> rewrite_function('A1N|A2', 'xor')
    #'A1_N|A2'

    """
    when_changes = []
    for m in RE_CELL.finditer(f):
        new_pinname = common.convert_pinname(m.group(0), cname)
        when_changes.append(((m.start(0), m.end(0)), (m.group(0), new_pinname)))

    when_out = list(f)
    for (s, e), (old_pinname, new_pinname) in reversed(when_changes):
        assert when_out[s:e] == list(old_pinname), (when_out[s:e], old_pinname, new_pinname)
        when_out[s:e] = list(new_pinname)

    return ''.join(when_out)


POWER_PIN_OVERRIDES = {
    'VNB': {
        "pg_type": "nwell",
        "physical_connection": "device_layer",
        "voltage_name": "VNB",
    },
    'VPB': {
        "pg_type": "pwell",
        "physical_connection": "device_layer",
        "voltage_name": "VPB",
    },
    'VPWR': {
        "pg_type": "primary_power",
        "related_bias_pin": "VNB",
        "voltage_name": "VPWR",
    },
    'VGND': {
        "pg_type": "primary_ground",
        "related_bias_pin": "VPB",
        "voltage_name": "VGND",
    },
}


def rewrite_cell(cname, cdata, n="top"):
    if isinstance(cdata, dict):
        o = {}
        for old_k, old_v in cdata.items():
            old_k = tuple(old_k.split(','))

            # Rewrite the key
            if "pin" in old_k[0] and len(old_k) > 1:
                assert len(old_k) == 2, (old_k, old_v)

                new_pinname = common.convert_pinname(old_k[1], cname)
                new_k = (old_k[0], new_pinname)

            elif old_k[0] in ('ff','latch','statetable'):
                new_k = [old_k[0]]
                for f in old_k[1:]:
                    new_k.append(common.convert_pinname(f, cname))
                new_k = tuple(new_k)
            else:
                new_k = old_k

            if debug and old_k != new_k:
                print(n, repr(old_k), "->", repr(new_k))

            nn = '{}[{}]'.format(n, repr(new_k))

            # Rewrite the value
            if new_k[0] in ('ff', 'latch'):
                new_v = {}
                assert isinstance(old_v, dict), old_v
                for k, v in old_v.items():

                    new_v[k] = rewrite_function(v, cname)
                    if debug and old_v[k] != new_v[k]:
                        print(nn, k, repr(old_v[k]), "->", repr(new_v[k]))

            elif new_k[0] in ('pg_pin',) and new_k[1] in POWER_PIN_OVERRIDES:
                new_v = POWER_PIN_OVERRIDES[new_k[1]]

            # - 'related_ground_pin'
            # - 'related_pin'
            # - 'related_power_pin'
            # - 'input_signal_level'
            elif 'related' in new_k[0] or new_k[0] in ('input_signal_level',):
                assert isinstance(old_v, str), repr(old_v)
                new_v = common.convert_pinname(old_v, cname)

            elif '_name' in new_k[0]:
                assert isinstance(old_v, str), repr(old_v)
                new_v = common.convert_pinname(old_v, cname)

            # - 'function'
            # - 'power_down_function'
            # - 'state_function'
            # - 'internal_node'
            # - 'three_state'
            # - 'when'
            elif 'function' in new_k[0] or new_k[0] in ('internal_node', 'three_state', 'when'):
                if isinstance(old_v, int):
                    old_v = str(old_v)
                assert isinstance(old_v, str), repr(old_v)
                new_v = rewrite_function(old_v, cname)

            else:
                new_v = rewrite_cell(cname, old_v, n=nn)

            if debug and isinstance(old_v, str) and old_v != new_v:
                print(n, repr(new_k), repr(old_v), "->", repr(new_v))

            o[",".join(new_k)] = new_v
        return o
    elif isinstance(cdata, list):
        o = []
        for i, v in enumerate(cdata):
            o.append(rewrite_cell(cname, v, n='{}[{}]'.format(n, i)))
        return o
    else:
        return cdata


def json_dump(d, f):
    json.dump(d, f, sort_keys=True, indent='  ')


TO_REMOVE = [
    "def_sim_opt",
    "simulator",
    "leakage_sim_opt",
    "date",
    "comment",
]


def filemain(input_file, temp_dir=None, final_dir=None, args=None):
    if input_file.endswith('.lib.json'):
        f = input_file
    elif input_file.endswith('.lib'):
        f = input_file + '.json'
        if not os.path.exists(f):
            print("SKIPPING:", f)
            return
    else:
        raise SystemError('Input file???:' + input_file)

    oldlib, newlib, ver = common.extract_version_and_lib_from_path(f)

    print(oldlib, newlib, ver)
    print('='*50)

    data = json.load(open(f, "r"))

    assert len(data) == 1, data.keys()
    k0 = list(data.keys())[0]
    data = data[k0]

    cells = {}
    for k in list(data.keys()):
        if not k.startswith('cell,'):
            continue
        c = k[5:]
        cells[c] = data[k]

        del data[k]

    for k, v in list(data.items()):
        if k in TO_REMOVE:
            del data[k]
            continue

        if "operating_conditions" in k:
            if 'scs8' in k:
                assert "," in k, (k, str(v)[:100])
                oc_type, oc_value = k.split(",", 1)

                newoc_value = calculate_out_filename(oc_value, v, ignore_ccsn=True)

                nk = "%s,%s" % (oc_type, newoc_value)
                assert nk not in data, nk
            else:
                nk = k

            if isinstance(v, str) and 'scs8' in v:
                nv = calculate_out_filename(v, data, ignore_ccsn=True)
            else:
                nv = v

            print('Rewriting\n  ', (k, v), "\n ->\n  ", (nk, nv))
            del data[k]
            data[nk] = nv

    for k, v in list(data.items()):
        if 'scs8' in k and not k.startswith('cell,'):
            assert False, ("Removing due to scs8 in key:", k, v)
            del data[k]
            continue

        if isinstance(v, str) and 'scs8' in v:
            assert False, ("Removing due to scs8 in value:", k, v)
            del data[k]
            continue

    for k, v in list(data.items()):
        if k not in CONSTANTS_FLAT:
            continue
        assert v == CONSTANTS_FLAT[k], (v, CONSTANTS_FLAT[k])
        del data[k]

    if 'comp_attribute,voltage_map' in data:
        o = []
        for k, v in data["comp_attribute,voltage_map"]:
            o.append((common.convert_pinname(k), v))
        data["comp_attribute,voltage_map"] = o


    fout_prefix = calculate_out_filename(f, data)
    assert fout_prefix, fout_prefix

    print(fout_prefix, "="*(75-len(fout_prefix)-1))
    pprint.pprint(data)
    print("="*75)

    if temp_dir:
        assert final_dir, (input_file, temp_dir, final_dir, args)

        ofname = '{}__{}.lib.json'.format(newlib, fout_prefix)
        ofpath = os.path.join(temp_dir, ofname)

        with open(ofpath, 'w') as f:
            json_dump(data, f)

        common.copy_file_to_output(ofpath, final_dir, newlib, ver, 'liberty', filename=ofname)

    for old_cellname, cdata in cells.items():
        new_cellname = common.convert_cell_fullname(old_cellname, newlib)

        old_key = 'cell '+old_cellname
        new_key = 'cell '+new_cellname

        if debug:
            print()
            print(old_cellname)
            print('='*50)
        odata = rewrite_cell(old_cellname, cdata)
        if 'cell_footprint' in odata:
            odata['cell_footprint'] = common.convert_cell_fullname(odata['cell_footprint'], newlib)

        if debug:
            orig_str = (pprint.pformat(cdata, width=100)+'\n').splitlines(keepends=True)
            new_str = (pprint.pformat(odata, width=100)+'\n').splitlines(keepends=True)
            diff = difflib.unified_diff(
                    orig_str,
                    new_str,
                    fromfile=old_cellname+' before',
                    tofile=new_cellname+' after',
                )
            print()
            sys.stdout.writelines(diff)
            print()

        if temp_dir:
            assert final_dir, (input_file, temp_dir, final_dir, args)

            ofname = '{}__{}.lib.json'.format(new_cellname, fout_prefix)
            ofpath = os.path.join(temp_dir, ofname)

            assert not os.path.exists(ofpath), ofpath

            with open(ofpath, 'w') as f:
                json_dump(odata, f)

            common.copy_file_to_output(ofpath, final_dir, newlib, ver, new_cellname, filename=ofname)


def merge_common(final_dir):
    final_dir = os.path.abspath(final_dir)

    for d in Path(final_dir).rglob("*/timing"):
        libdir = os.path.dirname(str(d))
        _, newlib, ver, _ = str(d).rsplit('/', 3)
        ver = ver.upper()

        pathmap = {}
        data = {}
        for fpath in d.glob("*.lib.json"):
            fname = os.path.basename(fpath)
            pathmap[fname] = str(fpath)
            with open(fpath) as f:
                data[fname] = json.load(f)

        keys = set()
        for fname in data:
            for k in data[fname].keys():
                keys.add(k)

        common_values = {}
        for k in keys:
            values = []
            for fname, v in data.items():
                values.append(v.get(k, None))

            assert len(values) == len(data)
            if all(values[i] == values[0] for i in range(1, len(values))):
                common_values[k] = values[0]

        if common_values:
            print()
            print("Common", "-"*75)
            pprint.pprint(common_values)
            print("-"*75)
            common_path = os.path.join(libdir, 'timing', "{}__common.lib.json".format(newlib))

            # HACK: Please make better!
            common_values_json = copy.copy(CONSTANTS)
            common_values_json.update(common_values)
            for k, v in CONSTANTS['comp_attribute,voltage_map'].items():
                if k in common_values_json['comp_attribute,voltage_map']:
                    ov = common_values_json['comp_attribute,voltage_map'][k]
                    assert ov == v, (k, v, ov)
                common_values_json['comp_attribute,voltage_map'][k] = v

            with open(common_path, "w") as f:
                json_dump(common_values_json, f)

            for k in common_values:
                for fname in data:
                    del data[fname][k]

            for fname in data:
                fpath = pathmap[fname]
                with open(fpath, "w") as f:
                    json_dump(data[fname], f)
                print()
                print(fname)
                print('-'*75)
                pprint.pprint(data[fname])
                print('-'*75)
                print()


if __name__ == "__main__":
    import doctest
    fails, _ = doctest.testmod()
    if fails != 0:
        sys.exit("Some test failed")
    if len(sys.argv) == 2:
        main()
    else:
        parser = argparse.ArgumentParser()
        parser.add_argument(
            "--common",
            action='store_true',
            default=False)
        parser.add_argument(
            "input",
            help="The path to the source directory/file",
            type=Path)
        parser.add_argument(
            "output",
            help="The path to the output directory",
            type=Path,
            nargs='?')
        parser.add_argument(
            "temp",
            help="The path to the temp directory",
            type=Path,
            nargs='?')
        args = parser.parse_args()
        if not args.common:
            common.main('lib', filemain, args)
        else:
            merge_common(args.output)