#!/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 argparse
import os
import re
import simplejson as json
import subprocess
import sys

from multiprocessing import Pool
from pathlib import Path


def call_with_retcode(cmd):
    print(cmd)
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = p.communicate()

    stdout = stdout.decode("utf-8")
    stderr = stderr.decode("utf-8")

    retcode = p.wait()
    print()
    print('='*5)
    print(stdout)
    print('-'*5)
    print(stderr)
    print(retcode, '='*5)
    print()

    return retcode, stderr


def use_netlistsvg(json_file, out_file):
    """Draw module schematics using netlistsvg. Operates on json files. returns retcode"""
    cmd = ["netlistsvg", json_file, "-o", out_file]
    return call_with_retcode(cmd)


def use_yosys(command, files):
    """Execute yosys command, returns retcode"""

    defines = {
        'FUNCTIONAL':    '=1',  # Want the functional model
        'NO_PRIMITIVES': '=1',  # Yosys doesn't support primitives
        'UNIT_DELAY':    '',    # Yosys doesn't support delays
    }
    actions = " ".join([
        "verilog_defines " + " ".join("-D{}{}".format(a, b) for a, b in defines.items()) + ";",
        "read_verilog " + " ".join(files) + ";",
        "hierarchy -check -auto-top; " + command,
    ])
    cmd = ["yosys", "-f", "verilog",  "-p",  actions]
    return call_with_retcode(cmd)


def draw_schematics(verilog_path):
    """Draw module schematics using yosys+netlistsvg. Operates on verilog files.
    Output file is placed in same directory as input"""

    verilog_path = str(verilog_path)
    verilog_dir = os.path.dirname(verilog_path)
    verilog_file = os.path.basename(verilog_path)

    verilog_base = verilog_file.split('.', 1)[0]

    verilog_json = os.path.join(verilog_dir, verilog_base+".json")
    verilog_schematic = os.path.join(verilog_dir, verilog_base+".schematic.svg")

    print("Generating schematic for", str(verilog_path))

    retcode, err = use_yosys(
        #"prep; splitinout; write_json {}.json" % verilog_base,
        "prep; write_json -compat-int {}".format(verilog_json),
        [verilog_path],
    )
    if retcode != 0:
        return verilog_path, err

    json_data = ""
    with open(verilog_json, "r") as f:
        contents = f.read()
    contents = re.sub('"type": "[^/"]+__', '"type": "', contents)
    contents = re.sub('("[^"]*)\\.\\./[^/]*/skywater-pdk/libraries/[^/]*/v[^/]*/', '\\1./', contents)
    with open(verilog_json, "w") as f:
        f.write(contents)

    retcode, err = use_netlistsvg(verilog_json, verilog_schematic)
    if retcode != 0:
        return verilog_path, err

    return verilog_path, None


def files(p):
    for f in p.rglob("*.v"):
        d = str(f).split('/')

        filename = d.pop(-1)
        cell = d.pop(-1)

        x = d.pop(-1)
        if x != 'cells':
            continue
        ver = d.pop(-1)
        lib = d.pop(-1)

        o = lib + '__' + cell + '.v'
        if filename != o:
            continue
        yield str(f)


# def split_inouts(in_file, out_file = None):
#     """Split inout ports into inputs and outputs. Operates on yosys json files"""
#     if out_file == None:
#         out_file = in_file
#     with open(in_file, 'r') as in_f:
#         design = json.load(in_f)
#     top_name = split_io.find_top_module(design)
#     port_map, net_map = split_io.find_and_split_inout_ports_and_nets(design, top_name)
#     output = json.dumps(design, sort_keys=True, indent=2)
#     output = re.sub(r'"00000000000000000000000000000001"', "1", output)
#     with open(out_file, 'w') as out_f:
#        out_f.write(output)


if __name__ == "__main__":
    for a in sys.argv[1:]:
        if os.path.isfile(a):
            draw_schematics(a)
            continue
        assert os.path.exists(a), a
        p = Path(a)
        print(p)
        pool = Pool()
        for verilog_path, err in pool.imap_unordered(draw_schematics, files(p)):
            if err is not None:
                msg = err + "\nError creating schematic file for:" + verilog_path+'\n\n'
                print(msg, file=sys.stderr, flush=True)
            else:
                print("Finished creating schematic file for:", verilog_path)
        pool.close()
        pool.join()