# SPDX-License-Identifier: GPL-2.0-or-later
from typing import Optional, Tuple, Union, cast

from pdkmaster.technology import geometry as _geo, primitive as _prm
from pdkmaster.design import circuit as _ckt, library as _lbry

from c4m.pdk import sky130

from . import frame as _frm


__all__ = ["ConnectedSRAM"]


class _io_spec:
    def __init__(self, *,
        sram_signal: Union[str, Tuple[str, str, str]],
        io_type: str, io_number: int,
    ):
        assert io_type in ("io_in", "io_out", "io_analog")
        assert (io_type != "io_inout") or isinstance(sram_signal, tuple)
        self.sram_signal = sram_signal
        self.io_type = io_type
        self.io_number = io_number

    @property
    def prefix(self) -> str:
        if self.io_type in ("io_in", "io_out", "io_inout"):
            return f"io{self.io_number}"
        elif self.io_type == "io_analog":
            return f"ioa{self.io_number}"
        elif self.io_type == "gpio":
            return f"gpio{self.io_number}"
        else:
            raise NotImplementedError(f"io_type == '{self.io_type}'")

    @property
    def toppin_name(self) -> str:
        num = self.io_number
        if self.io_type == "io_in":
            return f"io_in[{num}]"
        elif self.io_type == "io_out":
            return f"io_out[{num}]"
        elif self.io_type == "io_analog":
            return f"io_analog[{num}]"
        else:
            raise NotImplementedError(f"io_type == {self.io_type}")

    @property
    def oeb(self) -> Optional[bool]:
        if self.io_type == "io_in":
            return True
        elif self.io_type == "io_out":
            return False
        else:
            return None


io_specs = (
    _io_spec(sram_signal="a[8]", io_type="io_in", io_number=14),
    _io_spec(sram_signal="a[7]", io_type="io_in", io_number=15),
    _io_spec(sram_signal="a[6]", io_type="io_in", io_number=16),
    _io_spec(sram_signal="a[5]", io_type="io_in", io_number=17),
    _io_spec(sram_signal="a[4]", io_type="io_in", io_number=18),
    _io_spec(sram_signal="a[3]", io_type="io_in", io_number=19),
    _io_spec(sram_signal="a[2]", io_type="io_in", io_number=20),
    _io_spec(sram_signal="a[1]", io_type="io_in", io_number=21),
    _io_spec(sram_signal="a[0]", io_type="io_in", io_number=22),
    _io_spec(sram_signal="clk", io_type="io_in", io_number=23),
    _io_spec(sram_signal="q[7]", io_type="io_in", io_number=13),
    _io_spec(sram_signal="d[7]", io_type="io_out", io_number=12),
    _io_spec(sram_signal="d[6]", io_type="io_in", io_number=11),
    _io_spec(sram_signal="q[6]", io_type="io_out", io_number=10),
    _io_spec(sram_signal="q[5]", io_type="io_in", io_number=9),
    _io_spec(sram_signal="d[5]", io_type="io_out", io_number=8),
    _io_spec(sram_signal="d[4]", io_type="io_in", io_number=7),
    _io_spec(sram_signal="q[4]", io_type="io_out", io_number=6),
    _io_spec(sram_signal="q[3]", io_type="io_in", io_number=5),
    _io_spec(sram_signal="d[3]", io_type="io_out", io_number=4),
    _io_spec(sram_signal="d[2]", io_type="io_in", io_number=3),
    _io_spec(sram_signal="q[2]", io_type="io_out", io_number=2),
    _io_spec(sram_signal="q[1]", io_type="io_in", io_number=1),
    _io_spec(sram_signal="we[0]", io_type="io_in", io_number=0),
    _io_spec(sram_signal="d[1]", io_type="io_out", io_number=26),
    _io_spec(sram_signal="d[0]", io_type="io_in", io_number=25),
    _io_spec(sram_signal="q[0]", io_type="io_out", io_number=24),
    _io_spec(sram_signal="vss", io_type="io_analog", io_number=4),
    _io_spec(sram_signal="vdd", io_type="io_analog", io_number=5),
)
io_sig2spec = {
    spec.sram_signal: spec
    for spec in io_specs
}
io_pin2spec = {
    spec.toppin_name: spec
    for spec in io_specs
}


class ConnectedSRAM(_lbry._OnDemandCell[_lbry.Library]):
    def __init__(self, *, lib: _lbry.Library):
        super().__init__(lib=lib, name="ConnectedSRAM")

    def _create_circuit(self):
        stdcells = sky130.stdcelllib.cells

        mem_fab = sky130.Sky130SP6TFactory(lib=self.lib)

        zero_cell = stdcells.zero_x1
        one_cell = stdcells.one_x1
        tie_cell = stdcells.tie_diff_w4

        sram_cell = mem_fab.block(words=512, word_size=8, we_size=1)
        self.a_bits = sram_cell.a_bits
        self.word_size = sram_cell.word_size

        ckt = self.new_circuit()

        sram = ckt.instantiate(sram_cell, name="sram")
    
        for spec in io_specs:
            prefix = spec.prefix

            sram_port = sram.ports[spec.sram_signal]
            ckt.new_net(name=spec.toppin_name, external=True, childports=sram_port)

            oeb = spec.oeb
            if oeb is not None:
                num = spec.io_number
                if oeb:
                    one = ckt.instantiate(one_cell, name=f"{prefix}one")
                    ckt.new_net(
                        name=f"io_oeb[{num}]", external=True, childports=one.ports.one,
                    )
                else:
                    zero = ckt.instantiate(zero_cell, name=f"{prefix}zero")
                    ckt.new_net(
                        name=f"io_oeb[{num}]", external=True, childports=zero.ports.zero,
                    )

            if spec.io_type != "io_analog":
                ckt.instantiate(tie_cell, name=f"{prefix}tie")

    def _create_layout(self):
        tech = sky130.tech
        prims = tech.primitives

        ckt = self.circuit
        nets = ckt.nets
        insts = ckt.instances

        li = cast(_prm.MetalWire, prims.li)
        lipin = cast(_prm.Marker, prims["li.pin"])
        m1 = cast(_prm.MetalWire, prims.m1)
        m1pin = cast(_prm.Marker, prims["m1.pin"])
        via = cast(_prm.Via, prims.via)
        m2 = cast(_prm.MetalWire, prims.m2)
        m2pin = cast(_prm.Marker, prims["m2.pin"])
        via2 = cast(_prm.Via, prims.via2)
        m3 = cast(_prm.MetalWire, prims.m3)
        m3pin = cast(_prm.Marker, prims["m3.pin"])

        layouter = self.new_circuitlayouter()
        layout = layouter.layout

        _sram_lay = layouter.inst_layout(inst=insts.sram)
        _sram_bb = _sram_lay.boundary
        assert _sram_bb is not None

        x = sky130.tech.on_grid(_frm.boundary.center.x - _sram_bb.center.x)
        y = _frm.boundary.top - 100.0 - _sram_bb.top
        sram_lay = layouter.place(_sram_lay, x=x, y=y)
        sram_bb = sram_lay.boundary
        assert sram_bb is not None

        # vss
        spec = io_sig2spec["vss"]
        net = nets[spec.toppin_name]
        toppin_bb = _frm.toppins[spec.toppin_name]
        bbs = tuple(ms.shape.bounds for ms in sram_lay.filter_polygons(
            net=net, mask=m2pin.mask, depth=1, split=True,
        ))
        right = min(bb.right for bb in bbs)
        top = max(bb.top for bb in bbs)
        bottom = min(bb.bottom for bb in bbs)
        w = 10.0
        x_via2 = right - w/2.0
        left = right - w

        for bb in bbs:
            y_via2 = bb.center.y
            layouter.add_wire(
                net=net, wire=via2, x=x_via2, y=y_via2,
                bottom_width=w, bottom_enclosure="wide",
                top_width=w, top_enclosure="wide",
            )

        shape = _geo.Rect(
            left=left, bottom=bottom, right=right, top=(top + 20.0),
        )
        layouter.add_wire(net=net, wire=m3, shape=shape)
        shape = _geo.Rect(
            left=left, bottom=(top + 10.0), right=toppin_bb.right, top=(top+20.0),
        )
        layouter.add_wire(net=net, wire=m3, shape=shape)
        shape = _geo.Rect.from_rect(rect=toppin_bb, bottom=(top + 10.0))
        layouter.add_wire(net=net, wire=m3, shape=shape)
        layouter.add_wire(net=net, wire=m3, pin=m3pin, shape=toppin_bb)

        # vdd
        spec = io_sig2spec["vdd"]
        net = nets[spec.toppin_name]
        toppin_bb = _frm.toppins[spec.toppin_name]
        bbs = tuple(ms.shape.bounds for ms in sram_lay.filter_polygons(
            net=net, mask=m2pin.mask, depth=1, split=True,
        ))
        left = max(bb.left for bb in bbs)
        top = max(bb.top for bb in bbs)
        bottom = min(bb.bottom for bb in bbs)
        w = 10.0
        x_via2 = left + w/2.0
        right = left + w

        for bb in bbs:
            y_via2 = bb.center.y
            layouter.add_wire(
                net=net, wire=via2, x=x_via2, y=y_via2,
                bottom_width=w, bottom_enclosure="wide",
                top_width=w, top_enclosure="wide",
            )

        shape = _geo.Rect(
            left=left, bottom=bottom, right=right, top=(top + 20.0),
        )
        layouter.add_wire(net=net, wire=m3, shape=shape)
        shape = _geo.Rect(
            left=toppin_bb.left, bottom=(top + 10.0), right=right, top=(top+20.0),
        )
        layouter.add_wire(net=net, wire=m3, shape=shape)
        shape = _geo.Rect.from_rect(rect=toppin_bb, bottom=(top + 10.0))
        layouter.add_wire(net=net, wire=m3, shape=shape)
        layouter.add_wire(net=net, wire=m3, pin=m3pin, shape=toppin_bb)

        # a
        col = 0
        w = 10.0
        for a_bit in reversed(range(self.a_bits)):
            spec = io_sig2spec[f"a[{a_bit}]"]
            assert spec.io_type == "io_in", "Internal error"
            net = nets[spec.toppin_name]

            sram_m1pinbb = sram_lay.bounds(mask=m1pin.mask, net=net, depth=1)
            toppin_bb = _frm.toppins[spec.toppin_name]

            right = sram_bb.left - (2*col + 1)*w
            left = right - 1
            col += 1

            x_via = right - w/2.0
            y_via = tech.on_grid(sram_m1pinbb.center.y)
            via_lay = layouter.add_wire(
                net=net, wire=via, x=x_via, y=y_via,
                bottom_width=w, bottom_enclosure="wide",
                top_width=w, top_enclosure="wide",
            )
            via_m1bb = via_lay.bounds(mask=m1.mask)
            via_m2bb = via_lay.bounds(mask=m2.mask)

            shape = _geo.Rect.from_rect(rect=sram_m1pinbb, left=via_m1bb.left)
            layouter.add_wire(net=net, wire=m1, shape=shape)

            x_via2 = x_via
            y_via2 = toppin_bb.center.y
            via2_lay = layouter.add_wire(
                net=net, wire=via2, x=x_via2, y=y_via2,
                bottom_width=w, bottom_enclosure="wide",
                top_width=w, top_enclosure="wide",
            )
            via2_m2bb = via2_lay.bounds(mask=m2.mask)
            via2_m3bb = via2_lay.bounds(mask=m3.mask)

            shape = _geo.Rect.from_rect(rect=via_m2bb, bottom=via2_m2bb.bottom)
            layouter.add_wire(net=net, wire=m2, shape=shape)
            shape = _geo.Rect.from_rect(rect=toppin_bb, right=via2_m3bb.right)
            layouter.add_wire(net=net, wire=m3, shape=shape)
            layouter.add_wire(net=net, wire=m3, pin=m3pin, shape=toppin_bb)

        # The rest of the pin on m2
        for sram_signame in (
            "clk", "we[0]",
            *(f"d[{bit}]" for bit in range(self.word_size)),
            *(f"q[{bit}]" for bit in range(self.word_size)),
        ):
            spec = io_sig2spec[sram_signame]
            net = nets[spec.toppin_name]

            sram_m2pinbb = sram_lay.bounds(mask=m2pin.mask, net=net, depth=1)
            toppin_bb = _frm.toppins[spec.toppin_name]

            _via2_lay = layouter.wire_layout(
                net=net, wire=via2, rows=6, columns=6,
            )
            _via2_m2bb = _via2_lay.bounds(mask=m2.mask)
            if toppin_bb.center.x > sram_m2pinbb.center.x:
                x_via2 = sram_m2pinbb.left - _via2_m2bb.left
            else:
                x_via2 = sram_m2pinbb.right - _via2_m2bb.right
            y_via2 = toppin_bb.center.y
            via2_lay = layouter.place(_via2_lay, x=x_via2, y=y_via2)
            via2_m2bb = via2_lay.bounds(mask=m2.mask)
            via2_m3bb = via2_lay.bounds(mask=m3.mask)

            shape = _geo.Rect.from_rect(rect=sram_m2pinbb, bottom=via2_m2bb.bottom)
            layouter.add_wire(net=net, wire=m2, shape=shape)
            if toppin_bb.center.x > sram_m2pinbb.center.x:
                shape = _geo.Rect.from_rect(rect=toppin_bb, left=via2_m3bb.left)
            else:
                shape = _geo.Rect.from_rect(rect=toppin_bb, right=via2_m3bb.right)
            layouter.add_wire(net=net, wire=m3, shape=shape)
            layouter.add_wire(net=net, wire=m3, pin=m3pin, shape=toppin_bb)

        # boundary
        layout.boundary = _frm.boundary