"""
py3Dmol visualization utilities for molecular interactions.

This module provides functions to generate py3Dmol viewer JavaScript code
for visualizing different types of molecular interactions. These functions
can be used in web applications, Jupyter notebooks, and other contexts.
"""

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from ..core.interactions import (
        HydrogenBond,
        HalogenBond,
        PiInteraction,
        PiPiStacking,
        CarbonylInteraction,
        NPiInteraction,
    )


def _escape_pdb_content(pdb_content: str) -> str:
    """Escape PDB content for JavaScript.

    :param pdb_content: Raw PDB file content
    :type pdb_content: str
    :returns: Escaped PDB content safe for JavaScript template literals
    :rtype: str
    """
    return pdb_content.replace('\\', '\\\\').replace('`', '\\`').replace('$', '\\$')


def generate_png_export_js(viewer_id: str, filename: str) -> str:
    """Generate JavaScript code to export viewer as PNG.

    :param viewer_id: Unique ID of the viewer instance
    :type viewer_id: str
    :param filename: Filename for the downloaded PNG
    :type filename: str
    :returns: JavaScript code to export viewer as PNG
    :rtype: str

    Example::

        >>> js_code = generate_png_export_js("viewer1", "interaction.png")
        >>> # In web app with nicegui:
        >>> ui.run_javascript(js_code)
    """
    return f"""
        (function() {{
            const viewer = window.{viewer_id}_instance;
            if (viewer) {{
                viewer.render();
                const imgData = viewer.pngURI();
                const link = document.createElement('a');
                link.href = imgData;
                link.download = '{filename}';
                link.click();
            }} else {{
                console.error('Viewer not found');
            }}
        }})();
    """


def generate_hydrogen_bond_viewer_js(hb: "HydrogenBond", pdb_content: str, viewer_id: str) -> str:
    """Generate JavaScript code for hydrogen bond 3D visualization.

    :param hb: Hydrogen bond interaction
    :type hb: HydrogenBond
    :param pdb_content: PDB file content
    :type pdb_content: str
    :param viewer_id: Unique ID for the viewer div
    :type viewer_id: str
    :returns: JavaScript code to initialize the viewer
    :rtype: str

    Example::

        >>> js_code = generate_hydrogen_bond_viewer_js(hb, pdb_content, "viewer1")
        >>> # In Jupyter notebook:
        >>> from IPython.display import HTML, display
        >>> display(HTML(f'<div id="viewer1" style="width:800px;height:600px;"></div>'))
        >>> display(HTML(f'<script>{js_code}</script>'))
    """
    donor_chain = hb.donor.chain_id
    donor_resi = hb.donor.res_seq
    acceptor_chain = hb.acceptor.chain_id
    acceptor_resi = hb.acceptor.res_seq

    pdb_escaped = _escape_pdb_content(pdb_content)

    javascript = f"""
    (function() {{
        // Wait for both 3Dmol to load and div to be ready
        function init3Dmol() {{
            if (typeof $3Dmol === 'undefined') {{
                setTimeout(init3Dmol, 100);
                return;
            }}

            let element = document.getElementById("{viewer_id}");
            if (!element) {{
                setTimeout(init3Dmol, 100);
                return;
            }}

            try {{
                let viewer = $3Dmol.createViewer("{viewer_id}", {{backgroundColor: 'white'}});
                window.{viewer_id}_instance = viewer;  // Store for PNG export
                let pdbData = `{pdb_escaped}`;

                viewer.addModel(pdbData, "pdb");
                viewer.setStyle({{}}, {{cartoon: {{color: 'lightgray', opacity: 0.3}}}});

                // Show donor residue
                viewer.addStyle({{chain: '{donor_chain}', resi: {donor_resi}}},
                               {{stick: {{colorscheme: 'cyanCarbon', radius: 0.25}}}});

                // Show acceptor residue
                viewer.addStyle({{chain: '{acceptor_chain}', resi: {acceptor_resi}}},
                               {{stick: {{colorscheme: 'orangeCarbon', radius: 0.25}}}});

                // Add dashed line for hydrogen bond
                viewer.addCylinder({{
                    start: {{x: {hb.hydrogen.coords.x}, y: {hb.hydrogen.coords.y}, z: {hb.hydrogen.coords.z}}},
                    end: {{x: {hb.acceptor.coords.x}, y: {hb.acceptor.coords.y}, z: {hb.acceptor.coords.z}}},
                    radius: 0.15,
                    color: 'yellow',
                    dashed: true
                }});

                // Add labels
                viewer.addLabel("{hb.get_donor_residue()}",
                               {{position: {{x: {hb.donor.coords.x}, y: {hb.donor.coords.y}, z: {hb.donor.coords.z}}},
                                backgroundColor: 'cyan', fontColor: 'black', fontSize: 14}});

                viewer.addLabel("{hb.get_acceptor_residue()}",
                               {{position: {{x: {hb.acceptor.coords.x}, y: {hb.acceptor.coords.y}, z: {hb.acceptor.coords.z}}},
                                backgroundColor: 'orange', fontColor: 'white', fontSize: 14}});

                viewer.zoomTo({{chain: ['{donor_chain}', '{acceptor_chain}'], resi: [{donor_resi}, {acceptor_resi}]}});
                viewer.render();
                viewer.zoom(1.2);

                // Force multiple resizes to ensure canvas is sized correctly
                setTimeout(function() {{
                    viewer.resize();
                    viewer.render();
                }}, 100);
                setTimeout(function() {{
                    viewer.resize();
                    viewer.render();
                }}, 300);
                setTimeout(function() {{
                    viewer.resize();
                    viewer.render();
                }}, 500);
            }} catch (error) {{
                console.error("Error creating viewer:", error);
            }}
        }}

        // Start with a delay to ensure div is in DOM and sized
        setTimeout(init3Dmol, 400);
    }})();
    """
    return javascript


def generate_halogen_bond_viewer_js(xb: "HalogenBond", pdb_content: str, viewer_id: str) -> str:
    """Generate JavaScript code for halogen bond 3D visualization.

    :param xb: Halogen bond interaction
    :type xb: HalogenBond
    :param pdb_content: PDB file content
    :type pdb_content: str
    :param viewer_id: Unique ID for the viewer div
    :type viewer_id: str
    :returns: JavaScript code to initialize the viewer
    :rtype: str
    """
    donor_chain = xb.donor_atom.chain_id
    donor_resi = xb.donor_atom.res_seq
    acceptor_chain = xb.acceptor.chain_id
    acceptor_resi = xb.acceptor.res_seq

    pdb_escaped = _escape_pdb_content(pdb_content)

    javascript = f"""
    (function() {{
        // Wait for both 3Dmol to load and div to be ready
        function init3Dmol() {{
            if (typeof $3Dmol === 'undefined') {{
                setTimeout(init3Dmol, 100);
                return;
            }}

            let element = document.getElementById("{viewer_id}");
            if (!element) {{
                setTimeout(init3Dmol, 100);
                return;
            }}

            try {{
                let viewer = $3Dmol.createViewer("{viewer_id}", {{backgroundColor: 'white'}});
                window.{viewer_id}_instance = viewer;  // Store for PNG export
                let pdbData = `{pdb_escaped}`;

                viewer.addModel(pdbData, "pdb");
                viewer.setStyle({{}}, {{cartoon: {{color: 'lightgray', opacity: 0.3}}}});

                viewer.addStyle({{chain: '{donor_chain}', resi: {donor_resi}}},
                               {{stick: {{colorscheme: 'purpleCarbon', radius: 0.25}}}});

                viewer.addStyle({{chain: '{acceptor_chain}', resi: {acceptor_resi}}},
                               {{stick: {{colorscheme: 'orangeCarbon', radius: 0.25}}}});

                viewer.addCylinder({{
                    start: {{x: {xb.halogen.coords.x}, y: {xb.halogen.coords.y}, z: {xb.halogen.coords.z}}},
                    end: {{x: {xb.acceptor.coords.x}, y: {xb.acceptor.coords.y}, z: {xb.acceptor.coords.z}}},
                    radius: 0.15,
                    color: 'orange',
                    dashed: true
                }});

                // Add labels
                viewer.addLabel("{xb.get_donor_residue()}",
                               {{position: {{x: {xb.donor_atom.coords.x}, y: {xb.donor_atom.coords.y}, z: {xb.donor_atom.coords.z}}},
                                backgroundColor: 'purple', fontColor: 'white', fontSize: 14}});

                viewer.addLabel("{xb.get_acceptor_residue()}",
                               {{position: {{x: {xb.acceptor.coords.x}, y: {xb.acceptor.coords.y}, z: {xb.acceptor.coords.z}}},
                                backgroundColor: 'orange', fontColor: 'white', fontSize: 14}});

                viewer.zoomTo({{chain: ['{donor_chain}', '{acceptor_chain}'], resi: [{donor_resi}, {acceptor_resi}]}});
                viewer.render();
                viewer.zoom(1.2);

                // Force multiple resizes to ensure canvas is sized correctly
                setTimeout(function() {{
                    viewer.resize();
                    viewer.render();
                }}, 100);
                setTimeout(function() {{
                    viewer.resize();
                    viewer.render();
                }}, 300);
                setTimeout(function() {{
                    viewer.resize();
                    viewer.render();
                }}, 500);
            }} catch (error) {{
                console.error("Error creating viewer:", error);
            }}
        }}

        // Start with a delay to ensure div is in DOM and sized
        setTimeout(init3Dmol, 400);
    }})();
    """
    return javascript


def generate_pi_interaction_viewer_js(pi: "PiInteraction", pdb_content: str, viewer_id: str) -> str:
    """Generate JavaScript code for π interaction 3D visualization.

    :param pi: π interaction
    :type pi: PiInteraction
    :param pdb_content: PDB file content
    :type pdb_content: str
    :param viewer_id: Unique ID for the viewer div
    :type viewer_id: str
    :returns: JavaScript code to initialize the viewer
    :rtype: str
    """
    donor_chain = pi.donor.chain_id
    donor_resi = pi.donor.res_seq

    # Get pi system residue
    pi_res_seq = pi.pi_atoms[0].res_seq if pi.pi_atoms else ""
    pi_chain_id = pi.pi_atoms[0].chain_id if pi.pi_atoms else ""

    pdb_escaped = _escape_pdb_content(pdb_content)

    javascript = f"""
    (function() {{
        function init3Dmol() {{
            if (typeof $3Dmol === 'undefined') {{
                setTimeout(init3Dmol, 100);
                return;
            }}

            let element = document.getElementById("{viewer_id}");
            if (!element) {{
                setTimeout(init3Dmol, 100);
                return;
            }}

            try {{
                let viewer = $3Dmol.createViewer("{viewer_id}", {{backgroundColor: 'white'}});
                window.{viewer_id}_instance = viewer;  // Store for PNG export
                let pdb_data = `{pdb_escaped}`;
                viewer.addModel(pdb_data, "pdb");

                viewer.setStyle({{}}, {{cartoon: {{color: 'lightgray', opacity: 0.3}}}});

                // Highlight donor
                viewer.setStyle({{resi: '{donor_resi}', chain: '{donor_chain}'}}, {{stick: {{colorscheme: 'cyanCarbon'}}}});

                // Highlight π system
                viewer.setStyle({{resi: '{pi_res_seq}', chain: '{pi_chain_id}'}}, {{stick: {{colorscheme: 'greenCarbon'}}}});

                // Add sphere at π center to show the interaction point
                viewer.addSphere({{
                    center: {{x: {pi.pi_center.x}, y: {pi.pi_center.y}, z: {pi.pi_center.z}}},
                    radius: 0.3,
                    color: 'green',
                    alpha: 0.7
                }});

                // Highlight the hydrogen/X-atom as a sphere
                viewer.addStyle({{serial: {pi.hydrogen.serial}}},
                               {{sphere: {{color: 'yellow', radius: 0.5}}}});

                // Add dashed line from X-atom to π center
                viewer.addCylinder({{
                    start: {{x: {pi.hydrogen.coords.x}, y: {pi.hydrogen.coords.y}, z: {pi.hydrogen.coords.z}}},
                    end: {{x: {pi.pi_center.x}, y: {pi.pi_center.y}, z: {pi.pi_center.z}}},
                    radius: 0.1,
                    color: 'yellow',
                    dashed: true
                }});

                // Add labels
                viewer.addLabel('{pi.get_donor_residue()}',
                              {{position: {{x: {pi.donor.coords.x}, y: {pi.donor.coords.y}, z: {pi.donor.coords.z}}},
                               backgroundColor: 'cyan', fontColor: 'black', fontSize: 14}});

                viewer.addLabel('{pi.get_acceptor_residue()}',
                              {{position: {{x: {pi.pi_center.x}, y: {pi.pi_center.y}, z: {pi.pi_center.z}}},
                               backgroundColor: 'green', fontColor: 'white', fontSize: 14}});

                // Distance label (at midpoint between hydrogen and pi center)
                viewer.addLabel('{pi.distance:.2f} Å',
                              {{position: {{x: ({pi.hydrogen.coords.x} + {pi.pi_center.x}) / 2,
                                          y: ({pi.hydrogen.coords.y} + {pi.pi_center.y}) / 2,
                                          z: ({pi.hydrogen.coords.z} + {pi.pi_center.z}) / 2}},
                               backgroundColor: 'black', fontColor: 'white', fontSize: 12}});

                viewer.zoomTo({{chain: ['{donor_chain}', '{pi_chain_id}'], resi: ['{donor_resi}', '{pi_res_seq}']}});
                viewer.render();

                setTimeout(function() {{ viewer.resize(); viewer.render(); }}, 100);
                setTimeout(function() {{ viewer.resize(); viewer.render(); }}, 300);
                setTimeout(function() {{ viewer.resize(); viewer.render(); }}, 500);
            }} catch (error) {{
                console.error("Error creating viewer:", error);
            }}
        }}

        setTimeout(init3Dmol, 400);
    }})();
    """
    return javascript


def generate_pi_pi_stacking_viewer_js(pi_pi: "PiPiStacking", pdb_content: str, viewer_id: str) -> str:
    """Generate JavaScript code for π-π stacking 3D visualization.

    :param pi_pi: π-π stacking interaction
    :type pi_pi: PiPiStacking
    :param pdb_content: PDB file content
    :type pdb_content: str
    :param viewer_id: Unique ID for the viewer div
    :type viewer_id: str
    :returns: JavaScript code to initialize the viewer
    :rtype: str
    """
    # Get residue info from ring atoms
    ring1_res = pi_pi.ring1_atoms[0].res_seq if pi_pi.ring1_atoms else ""
    ring1_chain = pi_pi.ring1_atoms[0].chain_id if pi_pi.ring1_atoms else ""
    ring2_res = pi_pi.ring2_atoms[0].res_seq if pi_pi.ring2_atoms else ""
    ring2_chain = pi_pi.ring2_atoms[0].chain_id if pi_pi.ring2_atoms else ""

    pdb_escaped = _escape_pdb_content(pdb_content)

    javascript = f"""
    (function() {{
        function init3Dmol() {{
            if (typeof $3Dmol === 'undefined') {{
                setTimeout(init3Dmol, 100);
                return;
            }}

            let element = document.getElementById("{viewer_id}");
            if (!element) {{
                setTimeout(init3Dmol, 100);
                return;
            }}

            try {{
                let viewer = $3Dmol.createViewer("{viewer_id}", {{backgroundColor: 'white'}});
                window.{viewer_id}_instance = viewer;  // Store for PNG export

                let pdb_data = `{pdb_escaped}`;
                viewer.addModel(pdb_data, "pdb");

                viewer.setStyle({{}}, {{cartoon: {{color: 'lightgray', opacity: 0.3}}}});

                // Highlight both aromatic rings
                viewer.setStyle({{resi: '{ring1_res}', chain: '{ring1_chain}'}}, {{stick: {{colorscheme: 'cyanCarbon'}}}});
                viewer.setStyle({{resi: '{ring2_res}', chain: '{ring2_chain}'}}, {{stick: {{colorscheme: 'magentaCarbon'}}}});

                // Add residue labels
                viewer.addLabel('{pi_pi.ring1_residue}',
                              {{position: {{x: {pi_pi.ring1_center.x}, y: {pi_pi.ring1_center.y}, z: {pi_pi.ring1_center.z}}},
                               backgroundColor: 'cyan', fontColor: 'black', fontSize: 14}});

                viewer.addLabel('{pi_pi.ring2_residue}',
                              {{position: {{x: {pi_pi.ring2_center.x}, y: {pi_pi.ring2_center.y}, z: {pi_pi.ring2_center.z}}},
                               backgroundColor: 'magenta', fontColor: 'white', fontSize: 14}});

                // Add distance label
                viewer.addLabel('{pi_pi._distance:.2f} Å ({pi_pi.stacking_type})',
                              {{position: {{x: {pi_pi.midpoint.x}, y: {pi_pi.midpoint.y}, z: {pi_pi.midpoint.z}}},
                               backgroundColor: 'black', fontColor: 'white', fontSize: 12}});

                // Add line between ring centers
                viewer.addCylinder({{
                    start: {{x: {pi_pi.ring1_center.x}, y: {pi_pi.ring1_center.y}, z: {pi_pi.ring1_center.z}}},
                    end: {{x: {pi_pi.ring2_center.x}, y: {pi_pi.ring2_center.y}, z: {pi_pi.ring2_center.z}}},
                    radius: 0.1,
                    color: 'purple',
                    dashed: true
                }});

                viewer.zoomTo({{chain: ['{ring1_chain}', '{ring2_chain}'], resi: ['{ring1_res}', '{ring2_res}']}});
                viewer.render();

                setTimeout(function() {{ viewer.resize(); viewer.render(); }}, 100);
                setTimeout(function() {{ viewer.resize(); viewer.render(); }}, 300);
                setTimeout(function() {{ viewer.resize(); viewer.render(); }}, 500);
            }} catch (error) {{
                console.error("Error creating viewer:", error);
            }}
        }}

        setTimeout(init3Dmol, 400);
    }})();
    """
    return javascript


def generate_carbonyl_interaction_viewer_js(carbonyl: "CarbonylInteraction", pdb_content: str, viewer_id: str) -> str:
    """Generate JavaScript code for carbonyl n→π* interaction 3D visualization.

    :param carbonyl: Carbonyl interaction
    :type carbonyl: CarbonylInteraction
    :param pdb_content: PDB file content
    :type pdb_content: str
    :param viewer_id: Unique ID for the viewer div
    :type viewer_id: str
    :returns: JavaScript code to initialize the viewer
    :rtype: str
    """
    pdb_escaped = _escape_pdb_content(pdb_content)

    javascript = f"""
    (function() {{
        function init3Dmol() {{
            if (typeof $3Dmol === 'undefined') {{
                setTimeout(init3Dmol, 100);
                return;
            }}

            let element = document.getElementById("{viewer_id}");
            if (!element) {{
                setTimeout(init3Dmol, 100);
                return;
            }}

            try {{
                let viewer = $3Dmol.createViewer("{viewer_id}", {{backgroundColor: 'white'}});
                window.{viewer_id}_instance = viewer;  // Store for PNG export

                let pdb_data = `{pdb_escaped}`;
                viewer.addModel(pdb_data, "pdb");

                viewer.setStyle({{}}, {{cartoon: {{color: 'lightgray', opacity: 0.3}}}});

                // Highlight donor and acceptor residues
                viewer.setStyle({{resi: '{carbonyl.donor_oxygen.res_seq}', chain: '{carbonyl.donor_oxygen.chain_id}'}},
                              {{stick: {{colorscheme: 'redCarbon'}}}});
                viewer.setStyle({{resi: '{carbonyl.acceptor_carbon.res_seq}', chain: '{carbonyl.acceptor_carbon.chain_id}'}},
                              {{stick: {{colorscheme: 'blueCarbon'}}}});

                // Add residue labels
                viewer.addLabel('{carbonyl.get_donor_residue()}',
                              {{position: {{x: {carbonyl.donor_oxygen.coords.x}, y: {carbonyl.donor_oxygen.coords.y}, z: {carbonyl.donor_oxygen.coords.z}}},
                               backgroundColor: 'red', fontColor: 'white', fontSize: 14}});

                viewer.addLabel('{carbonyl.get_acceptor_residue()}',
                              {{position: {{x: {carbonyl.acceptor_carbon.coords.x}, y: {carbonyl.acceptor_carbon.coords.y}, z: {carbonyl.acceptor_carbon.coords.z}}},
                               backgroundColor: 'blue', fontColor: 'white', fontSize: 14}});

                // Add distance label (at midpoint)
                let carbonyl_midpoint = {{
                    x: ({carbonyl.donor_oxygen.coords.x} + {carbonyl.acceptor_carbon.coords.x}) / 2,
                    y: ({carbonyl.donor_oxygen.coords.y} + {carbonyl.acceptor_carbon.coords.y}) / 2,
                    z: ({carbonyl.donor_oxygen.coords.z} + {carbonyl.acceptor_carbon.coords.z}) / 2
                }};
                viewer.addLabel('{carbonyl.distance:.2f} Å',
                              {{position: carbonyl_midpoint,
                               backgroundColor: 'black', fontColor: 'white', fontSize: 12}});

                // Add line from donor O to acceptor C
                viewer.addCylinder({{
                    start: {{x: {carbonyl.donor_oxygen.coords.x}, y: {carbonyl.donor_oxygen.coords.y}, z: {carbonyl.donor_oxygen.coords.z}}},
                    end: {{x: {carbonyl.acceptor_carbon.coords.x}, y: {carbonyl.acceptor_carbon.coords.y}, z: {carbonyl.acceptor_carbon.coords.z}}},
                    radius: 0.1,
                    color: 'orange',
                    dashed: true
                }});

                viewer.zoomTo({{chain: ['{carbonyl.donor_oxygen.chain_id}', '{carbonyl.acceptor_carbon.chain_id}'], resi: ['{carbonyl.donor_oxygen.res_seq}', '{carbonyl.acceptor_carbon.res_seq}']}});
                viewer.render();

                setTimeout(function() {{ viewer.resize(); viewer.render(); }}, 100);
                setTimeout(function() {{ viewer.resize(); viewer.render(); }}, 300);
                setTimeout(function() {{ viewer.resize(); viewer.render(); }}, 500);
            }} catch (error) {{
                console.error("Error creating viewer:", error);
            }}
        }}

        setTimeout(init3Dmol, 400);
    }})();
    """
    return javascript


def generate_n_pi_interaction_viewer_js(n_pi: "NPiInteraction", pdb_content: str, viewer_id: str) -> str:
    """Generate JavaScript code for n→π* interaction 3D visualization.

    :param n_pi: n→π* interaction
    :type n_pi: NPiInteraction
    :param pdb_content: PDB file content
    :type pdb_content: str
    :param viewer_id: Unique ID for the viewer div
    :type viewer_id: str
    :returns: JavaScript code to initialize the viewer
    :rtype: str
    """
    # Get the π system residue sequence number from the first pi atom
    pi_res_seq = n_pi.pi_atoms[0].res_seq if n_pi.pi_atoms else ""
    pi_chain_id = n_pi.pi_atoms[0].chain_id if n_pi.pi_atoms else ""

    pdb_escaped = _escape_pdb_content(pdb_content)

    javascript = f"""
    (function() {{
        function init3Dmol() {{
            if (typeof $3Dmol === 'undefined') {{
                setTimeout(init3Dmol, 100);
                return;
            }}

            let element = document.getElementById("{viewer_id}");
            if (!element) {{
                setTimeout(init3Dmol, 100);
                return;
            }}

            try {{
                let viewer = $3Dmol.createViewer("{viewer_id}", {{backgroundColor: 'white'}});
                window.{viewer_id}_instance = viewer;  // Store for PNG export

                let pdb_data = `{pdb_escaped}`;
                viewer.addModel(pdb_data, "pdb");

                viewer.setStyle({{}}, {{cartoon: {{color: 'lightgray', opacity: 0.3}}}});

                // Highlight lone pair atom residue
                viewer.setStyle({{resi: '{n_pi.lone_pair_atom.res_seq}', chain: '{n_pi.lone_pair_atom.chain_id}'}},
                              {{stick: {{colorscheme: 'orangeCarbon'}}}});

                // Highlight π system residue
                viewer.setStyle({{resi: '{pi_res_seq}', chain: '{pi_chain_id}'}},
                              {{stick: {{colorscheme: 'tealCarbon'}}}});

                // Add residue labels
                viewer.addLabel('{n_pi.get_donor_residue()}',
                              {{position: {{x: {n_pi.lone_pair_atom.coords.x}, y: {n_pi.lone_pair_atom.coords.y}, z: {n_pi.lone_pair_atom.coords.z}}},
                               backgroundColor: 'orange', fontColor: 'white', fontSize: 14}});

                viewer.addLabel('{n_pi.get_acceptor_residue()}',
                              {{position: {{x: {n_pi.pi_center.x}, y: {n_pi.pi_center.y}, z: {n_pi.pi_center.z}}},
                               backgroundColor: 'teal', fontColor: 'white', fontSize: 14}});

                // Add distance label (at midpoint)
                let npi_midpoint = {{
                    x: ({n_pi.lone_pair_atom.coords.x} + {n_pi.pi_center.x}) / 2,
                    y: ({n_pi.lone_pair_atom.coords.y} + {n_pi.pi_center.y}) / 2,
                    z: ({n_pi.lone_pair_atom.coords.z} + {n_pi.pi_center.z}) / 2
                }};
                viewer.addLabel('{n_pi.distance:.2f} Å',
                              {{position: npi_midpoint,
                               backgroundColor: 'black', fontColor: 'white', fontSize: 12}});

                // Add line from lone pair atom to π center
                viewer.addCylinder({{
                    start: {{x: {n_pi.lone_pair_atom.coords.x}, y: {n_pi.lone_pair_atom.coords.y}, z: {n_pi.lone_pair_atom.coords.z}}},
                    end: {{x: {n_pi.pi_center.x}, y: {n_pi.pi_center.y}, z: {n_pi.pi_center.z}}},
                    radius: 0.1,
                    color: 'green',
                    dashed: true
                }});

                viewer.zoomTo({{chain: ['{n_pi.lone_pair_atom.chain_id}', '{pi_chain_id}'], resi: ['{n_pi.lone_pair_atom.res_seq}', '{pi_res_seq}']}});
                viewer.render();

                setTimeout(function() {{ viewer.resize(); viewer.render(); }}, 100);
                setTimeout(function() {{ viewer.resize(); viewer.render(); }}, 300);
                setTimeout(function() {{ viewer.resize(); viewer.render(); }}, 500);
            }} catch (error) {{
                console.error("Error creating viewer:", error);
            }}
        }}

        setTimeout(init3Dmol, 400);
    }})();
    """
    return javascript
