from typing import Optional, Dict, List, Tuple, Any
import networkx as nx
from rdkit import Chem
from rdkit.Chem import rdDetermineBonds

from .ascii_renderer import graph_to_ascii
from .data_loader import DATA
from .utils import graph_debug_report


def _bond_order_from_rdkit(bond) -> float:
    """Convert RDKit bond to numeric order."""
    if bond.GetIsAromatic() or bond.GetBondType() == Chem.BondType.AROMATIC:
        return 1.5
    bt = bond.GetBondType()
    if bt == Chem.BondType.SINGLE:
        return 1.0
    elif bt == Chem.BondType.DOUBLE:
        return 2.0
    elif bt == Chem.BondType.TRIPLE:
        return 3.0
    return 1.0


def xyz2mol_compare(
    atoms: List[Tuple[str, Tuple[float, float, float]]],
    charge: int = 0,
    verbose: bool = False,
    ascii: bool = False,
    ascii_scale: float = 2.0,
    ascii_include_h: bool = True,
    reference_graph: Optional[nx.Graph] = None,
) -> str:
    """
    RDKit DetermineBonds comparison / fallback diagnostic.
    verbose=True  -> include full atom neighbor table AND bond list.
    ascii=True    -> append ASCII depiction (scale, H visibility configurable).
    reference_graph -> if provided and atom counts match, ASCII can align to reference layout.
    """
    out = []
    out.append("\n" + "=" * 60)
    out.append("XYZ2MOL COMPARISON")
    out.append("=" * 60)

    # --- Build an XYZ block from atoms tuples ---
    nat = len(atoms)
    symbols = [symbol for symbol, _ in atoms]
    positions = [pos for _, pos in atoms]
    xyz_lines = [str(nat), f"Generated by xyz2mol_compare (charge={charge})"]
    for sym, (x, y, z) in zip(symbols, positions):
        xyz_lines.append(f"{sym} {x:.6f} {y:.6f} {z:.6f}")
    xyz_block = "\n".join(xyz_lines) + "\n"

    # --- Read an RDKit Mol (atoms + conformer, no bonds yet) ---
    raw_mol = Chem.MolFromXYZBlock(xyz_block)
    if raw_mol is None:
        return "# RDKit MolFromXYZBlock failed to parse XYZ\n"

    # Work on a copy we can modify
    mol = Chem.Mol(raw_mol)
    err = None
    try:
        rdDetermineBonds.DetermineBonds(mol, charge=charge, useHueckel=True)
    except Exception as e:
        err = str(e)

    if err or mol.GetNumBonds() == 0:
        lines = ["# RDKit DetermineBonds FAILED TO BUILD MOLECULE"]
        if err:
            lines.append(f"# Error: {err}")

        # Helpful hints for metals / coordination
        if any(s in DATA.metals for s in symbols):
            lines.append("# Issue: Metal atoms detected")
            lines.append("#   RDKit bond perception has limited support for coordination complexes")
        return "\n".join(lines) + "\n"

    # --- Light sanitize for properties; avoid full re-interpretation ---
    try:
        Chem.SanitizeMol(mol, sanitizeOps=Chem.SanitizeFlags.SANITIZE_PROPERTIES)
    except Exception:
        pass

    out.append(f"# RDKit DetermineBonds graph: {mol.GetNumAtoms()} atoms, {mol.GetNumBonds()} bonds (charge={charge})")

    # --- Build a NetworkX graph for ASCII and stats ---
    Gx = nx.Graph()
    for a in mol.GetAtoms():
        i = a.GetIdx()
        Gx.add_node(
            i,
            symbol=a.GetSymbol(),
            charges={},
            agg_charge=0.0,
            formal_charge=a.GetFormalCharge(),
            valence=0.0,
        )

    # Neighbor map for verbose section
    bond_map: Dict[int, List[Tuple[int, float, bool]]] = {}

    for b in mol.GetBonds():
        i = b.GetBeginAtomIdx()
        j = b.GetEndAtomIdx()
        bo = _bond_order_from_rdkit(b)
        ar = b.GetIsAromatic()

        Gx.add_edge(
            i, j,
            bond_order=bo,
            bond_type=(Gx.nodes[i]['symbol'], Gx.nodes[j]['symbol'])
        )

        bond_map.setdefault(i, []).append((j, bo, ar))
        bond_map.setdefault(j, []).append((i, bo, ar))

    for n in Gx.nodes:
        Gx.nodes[n]['valence'] = sum(Gx.edges[n, m].get('bond_order', 1.0) for m in Gx.neighbors(n))

    # --- ASCII depiction (optionally aligned to a reference layout) ---
    if ascii:
        aligned = False
        layout = None
        if reference_graph is not None and reference_graph.number_of_nodes() == Gx.number_of_nodes():
            _, layout = graph_to_ascii(
                reference_graph,
                scale=ascii_scale,
                include_h=ascii_include_h,
                return_layout=True
            )
            aligned = True

        if aligned and layout:
            ascii_txt = graph_to_ascii(Gx, scale=ascii_scale, include_h=ascii_include_h, reference_layout=layout)
            header = "# RDKit ASCII (aligned to reference)"
        else:
            ascii_txt = graph_to_ascii(Gx, scale=ascii_scale, include_h=ascii_include_h)
            header = "# RDKit ASCII"

        out.append(header)
        out.append(ascii_txt)
        out.append("")

    # --- Compare with reference graph (edge presence + bond-order diffs) ---
    if reference_graph is not None and reference_graph.number_of_nodes() == mol.GetNumAtoms():
        ref_edges = {}

        for i, j, d in reference_graph.edges(data=True):
            ref_edges[tuple(sorted((i, j)))] = float(d.get('bond_order', 1.0))

        x2m_edges = {}
        for i, j, d in Gx.edges(data=True):
            x2m_edges[tuple(sorted((i, j)))] = float(d.get('bond_order', 1.0))

        only_ref = sorted(e for e in ref_edges if e not in x2m_edges)
        only_x2 = sorted(e for e in x2m_edges if e not in ref_edges)
        shared = sorted(e for e in ref_edges if e in x2m_edges)

        bo_diffs = []
        for e in shared:
            r = ref_edges[e]
            x = x2m_edges[e]
            if abs(r - x) >= 0.25:
                bo_diffs.append((e, r, x, r - x))

        out.append("# Bond differences: only_in_native={:,}   only_in_rdkit={:,}   bond_order_diffs={:,}".format(len(only_ref), len(only_x2), len(bo_diffs)))

        if only_ref:
            out.append("#   only_in_native: " + " ".join(f"{a}-{b}" for a, b in only_ref))
        if only_x2:
            out.append("#   only_in_rdkit: " + " ".join(f"{a}-{b}" for a, b in only_x2))
        if bo_diffs:
            out.append("#   bond_order_diffs (Δ≥0.25):")
            for e, r, x, delta in bo_diffs[:40]:
                a, b = e
                out.append(f"#     {a:>3}-{b:<3}   native={r:>4.2f}   rdkit={x:>4.2f}   Δ={delta:+6.2f}")

            if len(bo_diffs) > 40:
                out.append("#     ...")

    # --- Verbose neighbor table and bonds ---
    if verbose:
        out.append(graph_debug_report(Gx, include_h=ascii_include_h))

    return "\n".join(out) + "\n"