# 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

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,
):
    counter = 0
    current_output = n(current_output)

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

    if isinstance(current_function, str):
        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):
    if port.msb is None:
        print(f"{port.direction.upper()}({n(port.name)})", file=f)
    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():
        output_hooked_to = inst.io[output]
        transform_function_rec(output_hooked_to, function, inst, f)


def netlist_to_bench(
    netlist: nl_parser.Netlist,
    cells: Dict[str, Cell],
    f: TextIOWrapper,
):
    print(f"# module {netlist.module}", file=f)
    print("# Automatically generated by nl2bench. Do not modify.", file=f)
    for port in netlist.ports.values():
        io_to_bench(port, f)
    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:
        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,
):
    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)
