# Copyright 2024 Mohamed Gaber

# 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

#      http://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 functools
from io import TextIOWrapper
from pathlib import Path
from typing import Iterable, Dict, Optional, Set

from . import nl_parser
from .cell import Cell


def n(identifier):
    if identifier.startswith("\\") and identifier.endswith(" "):
        identifier = identifier[1:-1]
    return identifier


def transform_function_rec(
    current_output: str,
    current_function,
    inst: nl_parser.Instance,
    f: TextIOWrapper,
    depth=0,
):
    argi = 0
    current_output = n(current_output)

    def handle_argument(arg):
        nonlocal argi
        if isinstance(arg, tuple):
            intermediate = current_output + f".{argi}"
            argi += 1
            transform_function_rec(intermediate, arg, inst, f, depth + 1)
            return intermediate
        else:
            assert (
                arg in inst.io
            ), f"Unrecognized input port {arg} in {inst.name} of kind {inst.kind}"
            return n(inst.io[arg])

    if isinstance(current_function, str):
        if current_function == "0":
            print(
                f"{current_output} = VSS()",
                file=f,
            )
        elif current_function == "1":
            print(
                f"{current_output} = VDD()",
                file=f,
            )
        else:
            print(
                f"{current_output} = BUF({handle_argument(current_function)})",
                file=f,
            )
    elif current_function[0] == "!":
        print(f"{current_output} = NOT({handle_argument(current_function[1])})", file=f)
    elif current_function[0] == "&":
        print(
            f"{current_output} = AND({handle_argument(current_function[1])} , {handle_argument(current_function[2])})",
            file=f,
        )
    elif current_function[0] == "^":
        print(
            f"{current_output} = XOR({handle_argument(current_function[1])} , {handle_argument(current_function[2])})",
            file=f,
        )
    elif current_function[0] == "+":
        print(
            f"{current_output} = OR({handle_argument(current_function[1])} , {handle_argument(current_function[2])})",
            file=f,
        )
    else:
        raise ValueError(
            f"Unknown lib function {current_function} of type {type(current_function)} in {inst}."
        )


def io_to_bench(port: nl_parser.Port, f: TextIOWrapper, msb_first: bool = False):
    if port.msb is None:
        print(f"{port.direction.upper()}({n(port.name)})", file=f)
    else:
        if msb_first:
            frm, to = max(port.msb, port.lsb), min(port.lsb, port.msb)
            i = frm
            while i >= to:
                print(f"{port.direction.upper()}({n(port.name)}[{i}])", file=f)
                i -= 1
        else:
            frm, to = min(port.msb, port.lsb), max(port.lsb, port.msb)
            i = frm
            while i <= to:
                print(f"{port.direction.upper()}({n(port.name)}[{i}])", file=f)
                i += 1


def statements_to_bench(inst: nl_parser.Instance, base: Cell, f: TextIOWrapper):
    for output, function in base.outputs.items():
        if hooked := inst.io.get(output):
            # HACK FOR SKY130
            if function is None and inst.kind.startswith("sky130_fd_sc_hd__dlclkp_"):
                function = "1"
            transform_function_rec(hooked, function, inst, f)


def netlist_to_bench(
    netlist: nl_parser.Netlist,
    cells: Dict[str, Cell],
    f: TextIOWrapper,
    bypass_ios: Optional[Set[str]] = None,
    msb_first: bool = False,
):
    bypass_ios = bypass_ios or set()
    print(f"# module {netlist.module}", file=f)
    print("# Automatically generated by nl2bench. Do not modify.", file=f)
    for port in netlist.ports.values():
        if port.name in bypass_ios:
            continue
        io_to_bench(port, f, msb_first)
    for inst in netlist.instances:
        assert (
            inst.kind in cells
        ), f"Unknown cell {inst.kind}- are you sure the design is combinational?"
        base = cells[inst.kind]
        statements_to_bench(inst, base, f)
    for asst in netlist.assignments:
        if asst[1] == 0:
            print(f"{n(asst[0])} = VSS()", file=f)
        elif asst[1] == 1:
            print(f"{n(asst[0])} = VDD()", file=f)
        else:
            print(f"{n(asst[0])} = BUF({n(asst[1])})", file=f)


def verilog_netlist_to_bench(
    netlist_in: Path,
    lib_files: Iterable[str],
    f: TextIOWrapper,
    bypass_ios: Optional[Set[str]] = None,
    msb_first: bool = False,
):
    cells: Dict[str, Cell] = functools.reduce(
        lambda acc, path: acc.update(Cell.from_lib_file(path)) or acc,
        lib_files,
        {},
    )

    netlist = nl_parser.parse(netlist_in)

    netlist_to_bench(netlist, cells, f, bypass_ios, msb_first)
