"""
Word-Width C Code Generator.

Emits C code that uses native integer types (uint8_t..uint64_t) for
bus-width operations, working directly from AnalysisResult.
"""

from io import StringIO
from collections import defaultdict

from ..compiler.ast import PrimitiveType
from .analyzer import AnalysisResult, BusGroup, BusSource
from .graph import WireRef, OutputSink, GateNode


def _select_c_type(width: int) -> str:
    if width <= 8:
        return "uint8_t"
    elif width <= 16:
        return "uint16_t"
    elif width <= 32:
        return "uint32_t"
    else:
        return "uint64_t"


def _width_mask(width: int) -> str:
    if width >= 64:
        return "0xffffffffffffffffull"
    mask = (1 << width) - 1
    if width <= 8:
        return f"0x{mask:02x}u"
    elif width <= 16:
        return f"0x{mask:04x}u"
    elif width <= 32:
        return f"0x{mask:08x}u"
    else:
        return f"0x{mask:016x}ull"


class BusCodeGenerator:
    """Generates C code directly from AnalysisResult."""

    def __init__(self, analysis: AnalysisResult):
        self.analysis = analysis
        self.output = StringIO()
        self.indent_level = 0
        # Map group_name -> C variable name used in tick()
        self._group_vars: dict[str, str] = {}
        # Map singleton gate name -> C variable name
        self._singleton_vars: dict[str, str] = {}

    def generate(self) -> str:
        self._emit_header()
        self._emit_state_struct()
        self._emit_dut_context()
        self._emit_tick_function()
        self._emit_api_functions()
        return self.output.getvalue()

    def _w(self, line: str = "") -> None:
        if line:
            self.output.write("    " * self.indent_level + line + "\n")
        else:
            self.output.write("\n")

    def _indent(self):
        self.indent_level += 1

    def _dedent(self):
        self.indent_level = max(0, self.indent_level - 1)

    # ── Header ──

    def _emit_header(self):
        self._w("/* Generated by SHDL bus-width codegen */")
        self._w("#include <stdint.h>")
        self._w("#include <string.h>")
        self._w("#include <stdio.h>")
        self._w()

    # ── State struct (feedback groups only) ──

    def _emit_state_struct(self):
        feedback_groups = [g for g in self.analysis.bus_groups if g.is_feedback]
        self._w("typedef struct {")
        self._indent()
        if feedback_groups:
            for g in feedback_groups:
                c_type = _select_c_type(g.width)
                self._w(f"{c_type} {g.name};")
        else:
            self._w("int _dummy;")
        self._dedent()
        self._w("} State;")
        self._w()

    # ── DUT Context ──

    def _emit_dut_context(self):
        self._w("typedef struct {")
        self._indent()
        self._w("State current;")
        for name, width in self.analysis.input_ports.items():
            self._w(f"{_select_c_type(width)} input_{name};")
        for name, width in self.analysis.output_ports.items():
            self._w(f"{_select_c_type(width)} output_{name};")
        self._w("int outputs_valid;")
        self._dedent()
        self._w("} DutContext;")
        self._w()
        self._w("static DutContext dut = {0};")
        self._w()

    # ── Tick function ──

    def _emit_tick_function(self):
        self._w("static inline void tick(void) {")
        self._indent()

        # 1. Load inputs into local variables
        for name, width in self.analysis.input_ports.items():
            c_type = _select_c_type(width)
            self._w(f"{c_type} {name} = dut.input_{name};")
        if self.analysis.input_ports:
            self._w()

        feedback_groups = [g for g in self.analysis.bus_groups if g.is_feedback]
        fb_names = {g.name for g in feedback_groups}

        # Build unified evaluation order for ALL non-feedback units
        pre_fb, post_fb = self._build_eval_order(fb_names)

        # 2. Emit pre-feedback units in topological order
        for unit in pre_fb:
            self._emit_unit(unit)

        # 3. Read previous state for ALL SCC groups
        if feedback_groups:
            self._w()
            for g in feedback_groups:
                c_type = _select_c_type(g.width)
                self._w(f"{c_type} prev_{g.name} = dut.current.{g.name};")
            self._w()

            # 4. Evaluate SCC groups using prev_* for cyclic deps
            for group in feedback_groups:
                self._emit_group_eval(group, is_feedback=True)
            self._w()

            # 5. Write back SCC state
            for g in feedback_groups:
                self._w(f"dut.current.{g.name} = {g.name};")

        # 6. Emit post-feedback units (singletons/groups that depend on feedback)
        if post_fb:
            self._w()
            for unit in post_fb:
                self._emit_unit(unit)

        # 7. Store outputs
        self._w()
        self._emit_output_stores()

        self._dedent()
        self._w("}")
        self._w()

    def _emit_unit(self, unit):
        """Emit a single evaluation unit (bus group or singleton gate)."""
        if isinstance(unit, BusGroup):
            self._emit_group_eval(unit, is_feedback=False)
        else:
            var = f"s_{unit.name}"
            self._singleton_vars[unit.name] = var
            expr = self._build_singleton_expr(unit)
            self._w(f"uint8_t {var} = {expr};")

    def _build_eval_order(self, fb_names: set[str]):
        """Build unified topological order, split into pre/post feedback phases.

        Returns (pre_feedback, post_feedback) lists of units.
        """
        from collections import deque

        non_fb_groups = [g for g in self.analysis.bus_groups if not g.is_feedback]

        # Build name -> unit mapping
        units: dict[str, object] = {}
        singleton_set: set[str] = set()
        for g in non_fb_groups:
            units[g.name] = g
        for gate in self.analysis.singleton_gates:
            uid = f"s_{gate.name}"
            units[uid] = gate
            singleton_set.add(uid)

        # Build forward dependency graph AND reverse adjacency list
        deps: dict[str, set[str]] = {uid: set() for uid in units}
        rdeps: dict[str, set[str]] = {uid: set() for uid in units}  # reverse edges
        depends_on_fb: set[str] = set()

        # Helper: resolve wire to dependency
        gate_to_group = self.analysis.gate_to_group
        for g in non_fb_groups:
            for source in g.input_sources.values():
                if source.kind == "bus_group":
                    if source.ref in fb_names:
                        depends_on_fb.add(g.name)
                    elif source.ref in units:
                        deps[g.name].add(source.ref)
                        rdeps[source.ref].add(g.name)
                elif source.kind == "mixed":
                    for wire in source.per_bit:
                        if wire and wire.kind == "gate_output":
                            grp_info = gate_to_group.get(wire.name)
                            if grp_info:
                                gname = grp_info[0]
                                if gname in fb_names:
                                    depends_on_fb.add(g.name)
                                elif gname in units:
                                    deps[g.name].add(gname)
                                    rdeps[gname].add(g.name)
                            else:
                                suid = f"s_{wire.name}"
                                if suid in singleton_set:
                                    deps[g.name].add(suid)
                                    rdeps[suid].add(g.name)

        for gate in self.analysis.singleton_gates:
            uid = f"s_{gate.name}"
            for wire in gate.inputs.values():
                if wire and wire.kind == "gate_output":
                    grp_info = gate_to_group.get(wire.name)
                    if grp_info:
                        gname = grp_info[0]
                        if gname in fb_names:
                            depends_on_fb.add(uid)
                        elif gname in units:
                            deps[uid].add(gname)
                            rdeps[gname].add(uid)
                    else:
                        suid = f"s_{wire.name}"
                        if suid in singleton_set:
                            deps[uid].add(suid)
                            rdeps[suid].add(uid)

        # Propagate feedback dependency transitively via BFS on reverse edges
        fb_queue = deque(depends_on_fb)
        while fb_queue:
            uid = fb_queue.popleft()
            for dependent in rdeps.get(uid, set()):
                if dependent not in depends_on_fb:
                    depends_on_fb.add(dependent)
                    fb_queue.append(dependent)

        # Kahn's topological sort using reverse adjacency list
        in_degree = {uid: len(d) for uid, d in deps.items()}
        queue = deque(uid for uid, deg in in_degree.items() if deg == 0)
        sorted_all: list[str] = []

        while queue:
            uid = queue.popleft()
            sorted_all.append(uid)
            for dependent in rdeps.get(uid, set()):
                in_degree[dependent] -= 1
                if in_degree[dependent] == 0:
                    queue.append(dependent)

        # Add remaining (cycles — shouldn't happen for non-feedback)
        if len(sorted_all) < len(units):
            visited = set(sorted_all)
            for uid in units:
                if uid not in visited:
                    sorted_all.append(uid)

        # Split into pre and post feedback
        pre_fb = [units[uid] for uid in sorted_all if uid not in depends_on_fb]
        post_fb = [units[uid] for uid in sorted_all if uid in depends_on_fb]

        return pre_fb, post_fb

    def _emit_group_eval(self, group: BusGroup, is_feedback: bool):
        """Emit a single bus-width operation."""
        c_type = _select_c_type(group.width)
        var = group.name
        self._group_vars[group.name] = var

        ptype = PrimitiveType.from_string(group.primitive)
        mask = _width_mask(group.width)

        a_expr = self._source_to_expr(group.input_sources.get("A"), group, is_feedback)
        b_expr = self._source_to_expr(group.input_sources.get("B"), group, is_feedback)

        if ptype == PrimitiveType.NOT:
            self._w(f"{c_type} {var} = (~({a_expr})) & {mask};")
        elif ptype == PrimitiveType.AND:
            self._w(f"{c_type} {var} = ({a_expr}) & ({b_expr});")
        elif ptype == PrimitiveType.OR:
            self._w(f"{c_type} {var} = ({a_expr}) | ({b_expr});")
        elif ptype == PrimitiveType.XOR:
            self._w(f"{c_type} {var} = ({a_expr}) ^ ({b_expr});")

    def _source_to_expr(self, source: BusSource | None, group: BusGroup, is_feedback: bool) -> str:
        if source is None:
            return "0"

        if source.kind == "port_aligned":
            if source.shift > 0:
                return f"({source.ref} >> {source.shift})"
            return source.ref

        elif source.kind == "port_broadcast":
            c_type = _select_c_type(group.width)
            if source.broadcast_bit > 1:
                return f"({c_type})(-(({source.ref} >> {source.broadcast_bit - 1}) & 1u))"
            return f"({c_type})(-({source.ref} & 1u))"

        elif source.kind == "bus_group":
            # If the source group is in the same SCC, use prev_ variable
            src_group = None
            for g in self.analysis.bus_groups:
                if g.name == source.ref:
                    src_group = g
                    break
            if is_feedback and src_group and src_group.is_feedback and src_group.scc_id == group.scc_id:
                return f"prev_{source.ref}"
            return source.ref

        elif source.kind == "constant":
            if source.ref == "VCC":
                return _width_mask(group.width)
            return "0"

        elif source.kind == "mixed":
            return self._emit_mixed_gather(source, group, is_feedback)

        return "0"

    def _emit_mixed_gather(self, source: BusSource, group: BusGroup, is_feedback: bool) -> str:
        parts = []
        for i, wire in enumerate(source.per_bit):
            if wire is None:
                continue
            bit_expr = self._wire_to_bit_expr(wire, is_feedback, group)
            if i == 0:
                parts.append(f"({bit_expr} & 1u)")
            else:
                parts.append(f"(({bit_expr} & 1u) << {i})")

        if not parts:
            return "0"
        return " | ".join(parts)

    def _wire_to_bit_expr(self, wire: WireRef, is_feedback: bool, group: BusGroup) -> str:
        if wire.kind == "port_input":
            if wire.bit_index > 1:
                return f"({wire.name} >> {wire.bit_index - 1})"
            return wire.name
        elif wire.kind == "gate_output":
            grp_info = self.analysis.gate_to_group.get(wire.name)
            if grp_info:
                gname, pos = grp_info
                # Check if we need prev_ for feedback
                src_group = None
                for g in self.analysis.bus_groups:
                    if g.name == gname:
                        src_group = g
                        break
                if is_feedback and src_group and src_group.is_feedback and src_group.scc_id == group.scc_id:
                    return f"(prev_{gname} >> {pos})"
                return f"({gname} >> {pos})"
            # Singleton
            return f"s_{wire.name}"
        elif wire.kind == "constant":
            return "1" if wire.name == "VCC" else "0"
        return "0"

    def _build_singleton_expr(self, gate) -> str:
        ptype = PrimitiveType.from_string(gate.primitive)
        if ptype == PrimitiveType.NOT:
            a_expr = self._wire_to_singleton_expr(gate.inputs.get("A"))
            return f"(~({a_expr}) & 1u)"
        elif ptype == PrimitiveType.VCC:
            return "1u"
        elif ptype == PrimitiveType.GND:
            return "0u"
        else:
            a_expr = self._wire_to_singleton_expr(gate.inputs.get("A"))
            b_expr = self._wire_to_singleton_expr(gate.inputs.get("B"))
            op = ptype.c_operator
            return f"(({a_expr}) {op} ({b_expr}))"

    def _wire_to_singleton_expr(self, wire) -> str:
        if wire is None:
            return "0u"
        if wire.kind == "port_input":
            if wire.bit_index > 1:
                return f"(({wire.name} >> {wire.bit_index - 1}) & 1u)"
            return f"({wire.name} & 1u)"
        elif wire.kind == "gate_output":
            grp_info = self.analysis.gate_to_group.get(wire.name)
            if grp_info:
                group_name, pos = grp_info
                return f"(({group_name} >> {pos}) & 1u)"
            return f"s_{wire.name}"
        elif wire.kind == "constant":
            return "1u" if wire.name == "VCC" else "0u"
        return "0u"

    # ── Output stores ──

    def _emit_output_stores(self):
        """Emit code to store outputs into dut.output_* at end of tick."""
        by_port: dict[str, list[OutputSink]] = defaultdict(list)
        for sink in self.analysis.output_sinks:
            by_port[sink.port_name].append(sink)

        for port_name, width in self.analysis.output_ports.items():
            sinks = by_port.get(port_name, [])
            if not sinks:
                self._w(f"dut.output_{port_name} = 0;")
                continue

            # Check if all sinks come from one bus group, aligned
            all_from_same_group = True
            group_name = None
            for sink in sinks:
                if sink.source.kind != "gate_output":
                    all_from_same_group = False
                    break
                grp_info = self.analysis.gate_to_group.get(sink.source.name)
                if grp_info is None:
                    all_from_same_group = False
                    break
                if group_name is None:
                    group_name = grp_info[0]
                elif grp_info[0] != group_name:
                    all_from_same_group = False
                    break

            if all_from_same_group and group_name:
                # Check alignment: output port bit k (1-based) maps to
                # group position (k-1). This means sink.bit_index-1 == pos.
                group = None
                for g in self.analysis.bus_groups:
                    if g.name == group_name:
                        group = g
                        break

                if group and len(sinks) == group.width:
                    aligned = True
                    for sink in sinks:
                        grp_info = self.analysis.gate_to_group[sink.source.name]
                        pos = grp_info[1]
                        if (sink.bit_index - 1) != pos:
                            aligned = False
                            break
                    if aligned:
                        self._w(f"dut.output_{port_name} = {group_name};")
                        continue

            # Fallback: per-bit gather
            parts = []
            for sink in sorted(sinks, key=lambda s: s.bit_index):
                bit_pos = sink.bit_index - 1  # 1-based to 0-based
                expr = self._sink_to_bit_expr(sink)
                if bit_pos == 0:
                    parts.append(f"({expr} & 1u)")
                else:
                    parts.append(f"(({expr} & 1u) << {bit_pos})")

            if parts:
                self._w(f"dut.output_{port_name} = {' | '.join(parts)};")
            else:
                self._w(f"dut.output_{port_name} = 0;")

    def _sink_to_bit_expr(self, sink: OutputSink) -> str:
        wire = sink.source
        if wire.kind == "gate_output":
            grp_info = self.analysis.gate_to_group.get(wire.name)
            if grp_info:
                gname, pos = grp_info
                return f"({gname} >> {pos})"
            return f"s_{wire.name}"
        elif wire.kind == "port_input":
            if wire.bit_index > 1:
                return f"({wire.name} >> {wire.bit_index - 1})"
            return wire.name
        elif wire.kind == "constant":
            return "1" if wire.name == "VCC" else "0"
        return "0"

    # ── API functions ──

    def _emit_api_functions(self):
        self._emit_reset()
        self._emit_poke()
        self._emit_peek()
        self._emit_step()

    def _emit_reset(self):
        self._w("void reset(void) {")
        self._indent()
        self._w("memset(&dut, 0, sizeof(dut));")
        self._dedent()
        self._w("}")
        self._w()

    def _emit_poke(self):
        self._w("void poke(const char *signal, uint64_t value) {")
        self._indent()

        items = list(self.analysis.input_ports.items())
        for i, (name, width) in enumerate(items):
            c_type = _select_c_type(width)
            mask = (1 << width) - 1
            cond = "if" if i == 0 else "} else if"
            self._w(f'{cond} (strcmp(signal, "{name}") == 0) {{')
            self._indent()
            self._w(f"dut.input_{name} = ({c_type})(value & 0x{mask:x}ull);")
            self._w("dut.outputs_valid = 0;")
            self._dedent()

        if items:
            self._w("} else {")
            self._indent()
            self._w('fprintf(stderr, "Unknown signal \'%s\'\\n", signal);')
            self._dedent()
            self._w("}")
        else:
            self._w('fprintf(stderr, "Unknown signal \'%s\'\\n", signal);')

        self._dedent()
        self._w("}")
        self._w()

    def _emit_peek(self):
        self._w("uint64_t peek(const char *signal) {")
        self._indent()

        for name in self.analysis.input_ports:
            self._w(f'if (strcmp(signal, "{name}") == 0) return (uint64_t)dut.input_{name};')

        self._w()
        self._w("if (!dut.outputs_valid) {")
        self._indent()
        self._w("tick();")
        self._w("dut.outputs_valid = 1;")
        self._dedent()
        self._w("}")
        self._w()

        for name in self.analysis.output_ports:
            self._w(f'if (strcmp(signal, "{name}") == 0) return (uint64_t)dut.output_{name};')

        self._w()
        self._w('fprintf(stderr, "Unknown signal \'%s\'\\n", signal);')
        self._w("return 0ull;")

        self._dedent()
        self._w("}")
        self._w()

    def _emit_step(self):
        self._w("void step(int cycles) {")
        self._indent()
        self._w("for (int i = 0; i < cycles; ++i) {")
        self._indent()
        self._w("tick();")
        self._dedent()
        self._w("}")
        self._w("dut.outputs_valid = 1;")
        self._dedent()
        self._w("}")
        self._w()
