#!/usr/bin/env python3
# Copyright 2020 The Skywater PDK Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
import subprocess
import re
import argparse
import glob
from shutil import copyfile, move
from pathlib import Path
from common import convert_libname, lib_extract_from_path, version_extract_from_path, lib_extract_from_name
from verilog2full import Copyright_header
import lef_split

debug = False
debug_print = lambda x: print(x) if debug else 0

Copyright_header = """/*\n
 * Copyright 2020 The Skywater PDK Authors\n
 *\n
 * Licensed under the Apache License, Version 2.0 (the "License");\n
 * you may not use this file except in compliance with the License.\n
 * You may obtain a copy of the License at\n
 *\n
 *     https://www.apache.org/licenses/LICENSE-2.0\n
 *\n
 * Unless required by applicable law or agreed to in writing, software\n
 * distributed under the License is distributed on an "AS IS" BASIS,\n
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n
 * See the License for the specific language governing permissions and\n
 * limitations under the License.\n
 */\n
 \n
"""

def prepend_copyright(filename):
 with open(filename, 'r+') as f:
    content = f.read()
    content = content.replace(Copyright_header, '')
    f.seek(0, 0)
    f.write(lef_split.Copyright_header + content)

def split_gds_into(input_gds, input_techfile, input_celllist, output_format):

    assert output_format in ('gds', 'lef'), "Wrong output format. Supported formats are 'gds' and 'lef'"

    if os.path.isfile(input_celllist):
        with open(input_celllist, 'r') as ifile:
            contents = ifile.read()
            input_celllist = contents.splitlines()
            if len(input_celllist) == 1:
                input_celllist = input_celllist[0].split(' ')

    destdir = os.path.split(input_gds)[0]
    gdsfile = os.path.split(input_gds)[1]
    tcl_path = destdir + '/split_gds.tcl'

    with open(tcl_path, 'w') as ofile:
        print('#!/bin/env wish', file=ofile)
        print('drc off', file=ofile)
        print('gds readonly true', file=ofile)
        print('gds rescale false', file=ofile)
        print('tech unlock *', file=ofile)
        print('gds read ' + gdsfile, file=ofile)

        for cell in input_celllist:

            lib, _ = lib_extract_from_name(cell)
            new_lib = convert_libname(lib)

            if cell == 'libcell' or cell == lib:
                continue
            # known to be useless and kills magic ;(
            if any([x  in cell for x in ('gpio_opath', 'i2c_fix', 'hotswap', 'odrvr', 'opamp')]):
                continue

            if cell.find(lib) == -1:
                new_cell_name = new_lib +'__' +  cell
            else:
                new_cell_name = cell.replace(lib + '_', new_lib + '__')
            print('cellname rename ' + cell + ' ' + new_cell_name, file=ofile)
            print('load ' + new_cell_name, file=ofile)
            print(f'gds write ' + new_cell_name, file=ofile)
            print(f'lef write ' + new_cell_name, file=ofile)
            print(f'save ' + new_cell_name, file=ofile)

        print('quit -noprompt', file=ofile)

    mproc = subprocess.run(['magic', '-dnull', '-noconsole',
        '-T', input_techfile, os.path.abspath(tcl_path)],
        stdin = subprocess.DEVNULL,
        stdout = subprocess.PIPE,
        stderr = subprocess.PIPE, cwd = destdir,
        universal_newlines = True)
    if mproc.stdout:
        for line in mproc.stdout.splitlines():
            print(line)
    # if mproc.stderr:
    #     print('Error message output from magic:')
        # for line in mproc.stderr.splitlines():
        #     print(line)
        if mproc.returncode != 0:
            print('ERROR:  Magic exited with status ' + str(mproc.returncode))
    os.remove(tcl_path)

def split_vcells_gds_into(input_gds, input_techfile, input_celllist):

    if os.path.isfile(input_celllist):
        with open(input_celllist, 'r') as ifile:
            contents = ifile.read()
            input_celllist = contents.splitlines()
            if len(input_celllist) == 1:
                input_celllist = input_celllist[0].split(' ')

    destdir = os.path.split(input_gds)[0]
    gdsfile = os.path.split(input_gds)[1]
    tcl_path = destdir + '/split_gds.tcl'

    with open(tcl_path, 'w') as ofile:
        print('#!/bin/env wish', file=ofile)
        print('drc off', file=ofile)
        print('gds readonly true', file=ofile)
        print('gds rescale false', file=ofile)
        print('tech unlock *', file=ofile)
        print('gds read ' + gdsfile, file=ofile)

        for cell in input_celllist:
            if 's8rf2' in cell:
                lib = 's8rf2'
                new_lib = 'sky130_fd_pr_rf2'
            elif 's8rf' in cell:
                lib = 's8rf'
                new_lib = 'sky130_fd_pr_rf'
            else:
                lib = 's8'
                new_lib = 'sky130_fd_pr_base'

            if cell.find(lib) == -1:
                new_cell_name = new_lib +'__' +  cell
            else:
                new_cell_name = cell.replace(lib + '_', new_lib + '__')
            print('cellname rename ' + cell + ' ' + new_cell_name, file=ofile)
            print('load ' + new_cell_name, file=ofile)
            print(f'gds write ' + new_cell_name, file=ofile)
            print(f'lef write ' + new_cell_name, file=ofile)
            print(f'save ' + new_cell_name, file=ofile)

        print('quit -noprompt', file=ofile)

    mproc = subprocess.run(['magic', '-dnull', '-noconsole',
        '-T', input_techfile, os.path.abspath(tcl_path)],
        stdin = subprocess.DEVNULL,
        stdout = subprocess.PIPE,
        stderr = subprocess.PIPE, cwd = destdir,
        universal_newlines = True)
    if mproc.stdout:
        for line in mproc.stdout.splitlines():
            print(line)
    # if mproc.stderr:
    #     print('Error message output from magic:')
        # for line in mproc.stderr.splitlines():
        #     print(line)
        if mproc.returncode != 0:
            print('ERROR:  Magic exited with status ' + str(mproc.returncode))
    os.remove(tcl_path)


def create_dir_path(output_dir, mod, newlib, version):
    """
    >>> create_dir_path('output', 'top_filter_narrow', 'sky130_fd_io', 'v0.0.1')
    'output/skywater-pdk/libraries/sky130_fd_io/v0.0.1/cells/top_filter_narrow'
    >>> create_dir_path('output', 'top_filter_narrow_1', 'sky130_fd_io', 'v0.0.1')
    'output/skywater-pdk/libraries/sky130_fd_io/v0.0.1/cells/top_filter_narrow'
    >>> create_dir_path('output', 'a2bb2oi_1', 'sky130_fd_io', 'v0.0.1')
    'output/skywater-pdk/libraries/sky130_fd_io/v0.0.1/cells/a2bb2oi'
    >>> create_dir_path('output', 'aaa', 'sky130_fd_pr_base', None)
    'output/skywater-pdk/libraries/sky130_fd_pr_base/cells/aaa'

    """

    if version is None:
        new_path = f"{output_dir}/skywater-pdk/libraries/{newlib}/cells/{re.sub(r'_[0-9]{1,2}$', '',mod)}"
    else:
        new_path = f"{output_dir}/skywater-pdk/libraries/{newlib}/{version}/cells/{re.sub(r'_[0-9]{1,2}$', '',mod)}"
    return new_path

def change_gds_string(source, oldstring, newstring):
    dest = source

    sourcedir = os.path.split(source)[0]
    gdsinfile = os.path.split(source)[1]

    destdir = os.path.split(dest)[0]
    gdsoutfile = os.path.split(dest)[1]

    with open(source, 'rb') as ifile:
        gdsdata = ifile.read()

    # To be done:  Allow the user to select a specific record type or types
    # in which to restrict the string substitution.  If no restrictions are
    # specified, then substitue in library name, structure name, and strings.

    recordtypes = ['libname', 'strname', 'string']
    recordfilter = [2, 6, 25]
    bsearch = bytes(oldstring, 'ascii')
    brep = bytes(newstring, 'ascii')

    datalen = len(gdsdata)
    if debug:
        print('Original data length = ' + str(datalen))
    dataptr = 0
    while dataptr < datalen:
        # Read stream records up to any string, then search for search text.
        bheader = gdsdata[dataptr:dataptr + 2]
        reclen = int.from_bytes(bheader, 'big')
        newlen = reclen
        if newlen == 0:
            print('Error: found zero-length record at position ' + str(dataptr))
            break

        rectype = gdsdata[dataptr + 2]
        datatype = gdsdata[dataptr + 3]

        if rectype in recordfilter:
            # Datatype 6 is STRING
            if datatype == 6:
                if debug:
                    print('Record type = ' + str(rectype) + ' data type = ' + str(datatype) + ' length = ' + str(reclen))

                bstring = gdsdata[dataptr + 4: dataptr + reclen]
                repstring = bstring.replace(bsearch, brep)
                if repstring != bstring:
                    before = gdsdata[0:dataptr]
                    after = gdsdata[dataptr + reclen:]
                    newlen = len(repstring) + 4
                    # Record sizes must be even
                    if newlen % 2 != 0:
                        # Was original string padded with null byte?  If so,
                        # remove the null byte and reduce newlen.  Otherwise,
                        # add a null byte and increase newlen.
                        if bstring[-1] == 0:
                            newlen -= 1
                        else:
                            repstring += b'\x00'
                            newlen += 1

                    bnewlen = newlen.to_bytes(2, byteorder='big')
                    brectype = rectype.to_bytes(1, byteorder='big')
                    bdatatype = datatype.to_bytes(1, byteorder='big')

                    # Assemble the new record
                    newrecord = bnewlen + brectype + bdatatype + repstring
                    # Reassemble the GDS data around the new record
                    gdsdata = before + newrecord[0:newlen] + after
                    # Adjust the data end location
                    datalen += (newlen - reclen)

                    if debug:
                        print('Replaced ' + str(bstring) + ' with ' + str(repstring))

        # Advance the pointer past the data
        dataptr += newlen

    with open(dest, 'wb') as ofile:
        ofile.write(gdsdata)


def main(input_path, temp_dir, techfile, output):
    global debug_print
    ver = version_extract_from_path(input_path)
    if ver is not None:
        ver = "V" + ".".join([str(v) for v in ver])
    lib = lib_extract_from_path(input_path)
    new_lib = convert_libname(lib)

    if input_path.find('s8pir_10r_vcells') != -1:
        lib = '???'
        new_lib = 'sky130_fd_pr_base'
    assert lib is not None and new_lib is not None

    tmp_list = temp_dir + '/cell.list'
    tmp_gds = temp_dir + '/input.gds'
    copyfile(input_path, tmp_gds)
    # os.system(f"strings {os.path.abspath(tmp_gds)}| grep -v lbcell |sort |uniq > {os.path.abspath(tmp_list)}")
    os.system(f"strings {os.path.abspath(tmp_gds)}| grep -v libcell |sort |uniq |grep '{lib}' > {os.path.abspath(tmp_list)}")

    # Load list of all gds cells
    with open(tmp_list, 'r') as in_f:
        objects_list_from_gds = set(sorted(in_f.read().split('\n')))
    debug_print(" gds objects list")
    objects_list_from_gds.remove('')
    debug_print(objects_list_from_gds)

    #Find all cells in .lef
    lef_file = sorted(Path("/".join(input_path.split("/")[:-2])).rglob('*.lef'))[0]
    with open(lef_file, 'r') as in_f:
        lef_cont = in_f.read()
    cells_list_from_lef = set(sorted(re.findall(r'(?<=MACRO )\w+', lef_cont)))
    debug_print(" Lef based list")
    debug_print(cells_list_from_lef)

    # Find all verilog files and create celllist from it
    verilog_list = sorted(Path("/".join(input_path.split("/")[:-2])).rglob('*.v'))
    cells_list_from_verilog = []
    for x in verilog_list:
        x = str(x)
        if 'stubs' in x:
            continue
        cell_name = re.search(rf'{lib}a?_[^p][0-9a-z_]*[a-z0-9]', x)
        # should not list primitives
        """cell_name = re.search(rf'{lib}(_|a_)[A-Za-z]\w*', x)
        if cell_name is None:
            # search for primitives
            cell_name = re.search(r'(?<=/)[A-Z]\w*?(?=\.v)', x) """

        if cell_name is None:
            debug_print(f"Didn't find name in {x}")
        else:
            cells_list_from_verilog.append(cell_name[0])
    cells_list_from_verilog = [x for x in cells_list_from_verilog if 'pg' not in x and 'lpflow' not in x]
    cells_list_from_verilog = set(sorted(cells_list_from_verilog))

    if debug:
        # output .csv table
        print("Cell_name,in .LEF,in .GDS")
        for x in cells_list_from_verilog:
            exists_in = lambda x,s: "Y" if x in c else "N"
            exists_in = lambda cell,s : 'Y' if cell in s else 'N'
            print(f"{x},{exists_in(x,cells_list_from_lef)},{exists_in(x ,objects_list_from_gds)}")
        return

    split_gds_into(tmp_gds, os.path.abspath(techfile), tmp_list, 'gds')
    os.remove(tmp_list)
    os.remove(tmp_gds)

    # remove empty gds files from output
    os.system("find " + temp_dir + " -type f -size -500c -name *.gds -exec rm {} \\;")
    # move all files to proper directory
    for x in Path(temp_dir).rglob('*.mag'):
        f = str(x)
        f_gds = f.replace('.mag', '.gds')
        if not os.path.exists(f_gds):
            # if .gds does not exists it does mean that magic generates empty file for that cell
            os.remove(f)
            continue
        filename = f.split('/')[-1].split('.')[0]
        filename = filename
        mod = filename.replace(f'{new_lib}__', '')

        # replace 'input.gds` with correct file name
        # and replace `TECHNAME` with `SKY130A`
        with open(f,'r') as f_in:
            mag_contents = f_in.read()
        mag_contents = Copyright_header + mag_contents
        mag_contents = mag_contents.replace(' TECHNAME', ' SKY130A')
        mag_contents = mag_contents.replace(' input.gds', ' ' + filename + '.gds')
        with open(f, 'w') as f_out:
            f_out.write(mag_contents)
        # Skipped for now as we are releasing without .mag files
        ## try:
        ##     copyfile(f, create_dir_path(output, mod, new_lib, ver) + '/' + filename + '.mag')
        ## except FileNotFoundError:
        ##     os.remove(f_gds)
        ##     print(f" Warning:File {f} not moved to target_directory, as it does not exist!")
        ##     print(f"{create_dir_path(output, mod, new_lib, ver)}")
        os.remove(f)
    for x in Path(temp_dir).rglob('*.lef'):
        f = str(x)
        f_gds = f.replace('.lef', '.gds')
        if not os.path.exists(f_gds):
            # if .gds does not exists it does mean that magic generates empty file for that cell
            os.remove(f)
            continue
        filename = f.split('/')[-1].split('.')[0]
        mod = filename.replace(f'{new_lib}__', '')
        output_path  = create_dir_path(output, mod, new_lib, ver) + '/' + filename + '.lef'
        if os.path.exists(output_path):
            print(f"Lef already exists, skipping {f}")
        else:
            try:
                copyfile(f, output_path)
                prepend_copyright(output_path)
            except FileNotFoundError:
                os.remove(f_gds)
                print(f" Warning:File {f} not moved to target_directory, as it does not exist!")
                print(f"{create_dir_path(output, mod, new_lib, ver)}")
        os.remove(f)
    for x in Path(temp_dir).rglob('*.gds'):
        f = str(x)
        change_gds_string(f, f'hkscl5hdv1_', new_lib + '__')
        change_gds_string(f, lib + '_', new_lib + '__')
        for x in ('scs8hd', 'scs8lp', 'scs8ls', 'scs8ms', 'scs8hd'):
            change_gds_string(f, f'{x}_', convert_libname(x)+ '__')
        filename = f.split('/')[-1].split('.')[0]

        mod = filename.replace(f'{new_lib}__', '')
        copyfile(f, create_dir_path(output, mod, new_lib, ver) + '/' + filename + '.gds')
        os.remove(f)

def main_vcells(input_path, temp_dir, techfile, output):

    # TODO - generate list from directory
    s8_versions = [
        'V2.0.1',
        'V2.0.0',
        'V1.3.0',
        'V1.2.1',
        'V1.2.0',
        'V1.1.0',
        'V1.0.1',
        'V1.0.0']

    tmp_list = temp_dir + '/cell.list'
    tmp_gds = temp_dir + '/input.gds'
    copyfile(input_path, tmp_gds)
    copyfile(input_path.replace('.gds', '.list'), tmp_list)

    split_vcells_gds_into(tmp_gds, os.path.abspath(techfile), tmp_list)
    os.remove(tmp_list)
    os.remove(tmp_gds)

    # move all files to proper directory
    for x in Path(temp_dir).rglob('*.mag'):
        f = str(x)
        filename = f.split('/')[-1].split('.')[0]
        filename = filename
        new_lib = filename.split('__')[0]
        if new_lib == 'sky130_fd_pr_base':
            lib = ''
        elif new_lib == 'sky130_fd_pr_rf':
            lib = 's8rf'
        elif new_lib == 'sky130_fd_pr_rf2':
            lib = 's8rf2'

        mod = filename.replace(f'{new_lib}__', '')

        # replace 'input.gds` with correct file name
        # and replace `TECHNAME` with `SKY130A`
        with open(f,'r') as f_in:
            mag_contents = f_in.read()
        mag_contents = Copyright_header + mag_contents
        mag_contents = mag_contents.replace(' TECHNAME', ' SKY130A')
        mag_contents = mag_contents.replace(' input.gds', ' ' + filename + '.gds')
        with open(f, 'w') as f_out:
            f_out.write(mag_contents)
        for ver in s8_versions:
            target_dir = create_dir_path(output, mod, new_lib, ver)
            if not os.path.exists(target_dir):
                os.makedirs(target_dir)
            print(f'Moving to {target_dir}')
            # copyfile(f, target_dir + '/' + filename + '.mag')
        os.remove(f)

    for x in Path(temp_dir).rglob('*.lef'):
        f = str(x)
        f_gds = f.replace('.lef', '.gds')
        filename = f.split('/')[-1].split('.')[0]
        new_lib = filename.split('__')[0]
        if new_lib == 'sky130_fd_pr_base':
            lib = ''
        elif new_lib == 'sky130_fd_pr_rf':
            lib = 's8rf'
        elif new_lib == 'sky130_fd_pr_rf2':
            lib = 's8rf2'
        mod = filename.replace(f'{new_lib}__', '')
        for ver in s8_versions:
            target_dir = create_dir_path(output, mod, new_lib, ver)
            print(f'Moving to {target_dir}')
            target_file = target_dir + '/' + filename + '.lef'
            copyfile(f, target_file)
            prepend_copyright(target_file)
        os.remove(f)
    for x in Path(temp_dir).rglob('*.gds'):
        f = str(x)
        filename = f.split('/')[-1].split('.')[0]
        new_lib = filename.split('__')[0]
        if new_lib == 'sky130_fd_pr_base':
            lib = ''
        elif new_lib == 'sky130_fd_pr_rf':
            lib = 's8rf'
        elif new_lib == 'sky130_fd_pr_rf2':
            lib = 's8rf2'
        mod = filename.replace(f'{new_lib}__', '')
        change_gds_string(f, f'hkscl5hdv1_{mod}', new_lib + '__' + mod)
        for x in ('scs8hd', 'scs8lp', 'scs8ls', 'scs8ms'):
            change_gds_string(f, f'{x}_', convert_libname(x)+ '__')
        if lib == '':
            change_gds_string(f, mod, new_lib + '__' + mod)
        else:
            change_gds_string(f, lib + '_', new_lib + '__')
        for ver in s8_versions:
            target_dir = create_dir_path(output, mod, new_lib, ver)
            print(f'Moving to {target_dir}')
            copyfile(f, target_dir + '/' + filename + '.gds')
        os.remove(f)


if __name__ == "__main__":
    import doctest
    fails, _ = doctest.testmod()
    if fails != 0:
        sys.exit("Some test failed")
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "input",
        help="The path to the source directory/file",
        type=str)
    parser.add_argument(
        "output",
        help="The path to the output directory",
        type=Path)
    parser.add_argument(
        "techfile",
        help="Full path to the techfile",
        type=str)
    parser.add_argument(
        "temp",
        help="The path to the temp directory",
        type=Path)
    args = parser.parse_args()
    temp_dir = str(args.temp / 'gds_split')

    try:
        os.makedirs(temp_dir)
    except FileExistsError:
        pass

    input_files = [
        'scs8hs/V0.0.0/gds/scs8hs.gds',
        'scs8hvl/V0.0.0/gds/scs8hvl.gds',
        'scs8hvl/V0.0.1/gds/scs8hvl.gds',
        'scs8lp/V0.0.0/gds/scs8lp.gds',
        'scs8ls/V0.1.0/gds/scs8ls.gds',
        'scs8ms/V0.0.0/gds/scs8ms.gds',
        'scs8hd/V0.0.1/gds/scs8hd.gds',
        'scs8hdll/V0.1.0/gds/scs8hdll.gds'
        # 's8iom0s8/V0.2.0/gds/s8iom0s8.gds',
        # 's8iom0s8/V0.1.0/gds/s8iom0s8.gds',
        # 's8iom0s8/V0.2.1/gds/s8iom0s8.gds'
        # 's8iom0s8/V0.0.0/gds/s8iom0s8.gds',
        ]
    for x in input_files:
        path = args.input.rstrip('/') + '/' + x
        print(path)
        main(path, temp_dir, args.techfile, str(args.output))
    main_vcells(args.input.rstrip('/') + '/' + 'vcells/s8pir_10r_vcells_lvs.gds', temp_dir, args.techfile, str(args.output))
    os.rmdir(temp_dir)