from __future__ import annotations

from collections import deque

from ..backends.base import Backend
from ..ir.types import Instruction


def _build_graph(edges: list[tuple[int, int]]) -> dict[int, list[int]]:
    """Undirected adjacency from coupling_map."""
    g: dict[int, list[int]] = {}
    for a, b in edges or []:
        g.setdefault(a, []).append(b)
        g.setdefault(b, []).append(a)
    return g


def _shortest_path(g: dict[int, list[int]], src: int, dst: int) -> list[int] | None:
    """BFS shortest path in an undirected graph. Returns list of vertices [src,...,dst]."""
    if src == dst:
        return [src]
    q = deque([src])
    prev: dict[int, int | None] = {src: None}
    while q:
        u = q.popleft()
        for v in g.get(u, []):
            if v in prev:
                continue
            prev[v] = u
            if v == dst:
                # reconstruct
                path = [dst]
                while prev[path[-1]] is not None:
                    path.append(prev[path[-1]])  # type: ignore[index]
                path.reverse()
                return path
            q.append(v)
    return None


def _emit_swap_as_cx_triple(out_insts: list[Instruction], a: int, b: int):
    """Decompose SWAP(a,b) into 3 CXs on adjacent qubits a-b."""
    out_insts.append(Instruction("cx", (a, b), ()))
    out_insts.append(Instruction("cx", (b, a), ()))
    out_insts.append(Instruction("cx", (a, b), ()))


def route_on_coupling(circuit, backend: Backend):
    """
    Route CXs so they respect backend.target.coupling_map by inserting SWAP chains.
    Strategy: maintain a logical->physical mapping; 1q gates follow the mapping;
    when a CX(log_c,log_t) isn't on an edge, move the *control's physical* toward target
    along a shortest path with SWAPs (no swap-back). Update mapping as we swap.

    Returns a NEW circuit with only native ops (cx, rz, sx, x, measure),
    whose CXs are all between adjacent qubits per coupling_map.
    """
    from ..circuit.quantum_circuit import QuantumCircuit  # avoid circular import

    tgt = backend.target
    g = _build_graph(tgt.coupling_map)
    n = circuit.n_qubits

    # Start with identity layout: logical i at physical i
    # loc[logical] = physical; which_log[physical] = logical
    loc = list(range(n))
    which_log = list(range(n))

    out = QuantumCircuit(n, getattr(circuit, "n_clbits", 0))
    out_insts: list[Instruction] = out._instructions  # alias

    for inst in circuit.instructions:
        name, qs, ps, cs = inst.name, inst.qubits, inst.params, inst.clbits

        if name in ("x", "sx", "rz"):
            q = qs[0]
            out_insts.append(Instruction(name, (loc[q],), ps, cs))
            continue

        if name == "measure":
            q = qs[0]
            out_insts.append(Instruction("measure", (loc[q],), ps, cs))
            continue

        if name == "cx":
            c_log, t_log = qs
            c_phy, t_phy = loc[c_log], loc[t_log]

            # If directly connected, emit and continue
            if (c_phy, t_phy) in tgt.coupling_map or (t_phy, c_phy) in tgt.coupling_map:
                out_insts.append(Instruction("cx", (c_phy, t_phy), ()))
                continue

            # Otherwise, find path and move control toward target with SWAPs
            path = _shortest_path(g, c_phy, t_phy)
            if not path or len(path) < 2:
                raise ValueError(
                    f"No path between qubits {c_phy} and {t_phy} for routing"
                )

            # Walk path by swapping control forward one hop at a time (record hops)
            current = c_phy
            hops: list[tuple[int, int]] = []
            for next_hop in path[1:-1]:  # stop one before t_phy
                _emit_swap_as_cx_triple(out_insts, current, next_hop)
                # Update mapping
                l_a = which_log[current]
                l_b = which_log[next_hop]
                which_log[current], which_log[next_hop] = l_b, l_a
                loc[l_a], loc[l_b] = next_hop, current
                hops.append((current, next_hop))
                current = next_hop

            # Now control is adjacent to target; emit CX with current (control) -> t_phy
            c_phy = current
            if not (
                (c_phy, t_phy) in tgt.coupling_map or (t_phy, c_phy) in tgt.coupling_map
            ):
                raise ValueError("Routing logic error: not adjacent after swaps")
            out_insts.append(Instruction("cx", (c_phy, t_phy), ()))

            # Swap BACK along the reverse path to restore the original logical->physical layout
            for a, b in reversed(hops):
                _emit_swap_as_cx_triple(out_insts, a, b)
                # Undo the mapping change
                l_a = which_log[a]
                l_b = which_log[b]
                which_log[a], which_log[b] = l_b, l_a
                loc[l_a], loc[l_b] = b, a

            continue

        # Anything else: copy as-is (basis_map should have eliminated non-native)
        out_insts.append(inst)

    return out
