"""PyVista-based crystal structure visualization."""

import warnings
warnings.filterwarnings("ignore", category=UserWarning, module="pyvista")
warnings.filterwarnings("ignore", category=UserWarning, module="vtk")
warnings.filterwarnings("ignore", category=UserWarning, module="pymatgen")

import gc
import numpy as np
from typing import Dict, List, Optional, Tuple

import pyvista as pv

from pymatgen.symmetry.analyzer import SpacegroupAnalyzer

from ceraviz.types import CrystalTD, LatticeTD, MaterialTD
from ceraviz.colors import (
    UNIFORM_RADIUS, hex_to_rgb, get_element_color,
)
from ceraviz.structure import (
    build_pmg_structure, make_supercell, parse_oxidation_states,
)
from ceraviz.bonds import find_bonds_crystalnn


def build_atoms(structure: CrystalTD, uniform_radius: float = UNIFORM_RADIUS) -> pv.PolyData:
    """Build PyVista PolyData for atom positions with color and radius data."""
    atoms = structure["atoms"]
    points = np.array([a["cart"] for a in atoms])
    elements = [a["element"] for a in atoms]

    poly = pv.PolyData(points)
    poly["radius"] = np.full(len(atoms), uniform_radius)
    poly["rgb"] = np.array(
        [hex_to_rgb(get_element_color(e)) for e in elements],
        dtype=np.uint8
    )
    poly["element"] = elements
    poly["atom_index"] = np.arange(len(atoms))

    return poly


def add_atoms(plotter: pv.Plotter, atom_poly: pv.PolyData) -> None:
    """Add atom spheres to a PyVista plotter."""
    spheres = atom_poly.glyph(
        scale="radius",
        geom=pv.Sphere(theta_resolution=32, phi_resolution=32),
        orient=False,
    )

    plotter.add_mesh(
        spheres,
        scalars="rgb",
        rgb=True,
        smooth_shading=True,
    )


def add_unit_cell(plotter: pv.Plotter, lattice: LatticeTD) -> None:
    """Draw unit cell wireframe edges."""
    A, B, C = lattice["vectors"]

    origin = np.zeros(3)
    corners = [
        origin,
        A, B, C,
        A + B, A + C, B + C,
        A + B + C,
    ]

    edges = [
        (0, 1), (0, 2), (0, 3),
        (1, 4), (1, 5),
        (2, 4), (2, 6),
        (3, 5), (3, 6),
        (4, 7), (5, 7), (6, 7)
    ]

    for i, j in edges:
        plotter.add_lines(
            np.vstack([corners[i], corners[j]]),
            color="black",
            width=1.5,
        )


def add_bonds(
    plotter: pv.Plotter,
    structure_td: CrystalTD,
    bonds: List[dict],
    radius: float = 0.14,
) -> None:
    """Add bonds with split-bond coloring (each half takes adjacent atom's color)."""
    atoms = structure_td["atoms"]
    lattice = np.array(structure_td["lattice"]["vectors"])

    blocks = []

    for b in bonds:
        i, j = b["i"], b["j"]
        image = np.array(b["image"], dtype=float)

        ri = np.array(atoms[i]["cart"])
        rj = np.array(atoms[j]["cart"]) + image @ lattice

        direction = rj - ri
        length = np.linalg.norm(direction)
        if length < 1e-3:
            continue

        midpoint = 0.5 * (ri + rj)
        half_length = length / 2

        ci = np.array(hex_to_rgb(get_element_color(atoms[i]["element"])), dtype=np.uint8)
        cj = np.array(hex_to_rgb(get_element_color(atoms[j]["element"])), dtype=np.uint8)

        # First half-cylinder: from atom i to midpoint
        center_i = 0.5 * (ri + midpoint)
        cyl_i = pv.Cylinder(
            center=center_i,
            direction=direction,
            radius=radius,
            height=half_length,
            resolution=24,
        )
        cyl_i.cell_data["rgb"] = np.tile(ci, (cyl_i.n_cells, 1))
        blocks.append(cyl_i)

        # Second half-cylinder: from midpoint to atom j
        center_j = 0.5 * (midpoint + rj)
        cyl_j = pv.Cylinder(
            center=center_j,
            direction=direction,
            radius=radius,
            height=half_length,
            resolution=24,
        )
        cyl_j.cell_data["rgb"] = np.tile(cj, (cyl_j.n_cells, 1))
        blocks.append(cyl_j)

    if not blocks:
        return

    bond_mesh = pv.MultiBlock(blocks).combine()

    plotter.add_mesh(
        bond_mesh,
        scalars="rgb",
        rgb=True,
        smooth_shading=True,
        specular=0.1,
    )


def enable_hover(
    plotter: pv.Plotter,
    atom_poly: pv.PolyData,
    structure: CrystalTD,
    oxidation_states: Optional[Dict[str, int]] = None,
) -> None:
    """Click on atoms to show info."""
    atoms = structure["atoms"]

    def on_click(point):
        if point is None or len(point) == 0:
            return

        idx = atom_poly.find_closest_point(point)
        if idx < 0 or idx >= len(atoms):
            return

        atom = atoms[idx]
        elem = atom["element"]

        oxi_str = ""
        if oxidation_states and elem in oxidation_states:
            oxi = oxidation_states[elem]
            oxi_str = f"{oxi}+" if oxi > 0 else f"{abs(oxi)}-"

        color_hex = get_element_color(elem)

        info = (
            f"Atom #{idx}: {elem}{oxi_str}\n"
            f"Color: {color_hex}\n"
            f"Cart: ({atom['cart'][0]:.3f}, {atom['cart'][1]:.3f}, {atom['cart'][2]:.3f})\n"
            f"Frac: ({atom['frac'][0]:.3f}, {atom['frac'][1]:.3f}, {atom['frac'][2]:.3f})"
        )

        plotter.add_text(info, position="upper_left", name="hover_text", font_size=10)

    plotter.track_click_position(callback=on_click, side="left")


def add_element_legend(
    plotter: pv.Plotter,
    elements: List[str],
    oxidation_states: Optional[Dict[str, int]] = None,
) -> None:
    """Add Materials Project-style legend with horizontal pill badges."""
    import vtk

    n_elements = len(elements)
    if n_elements == 0:
        return

    badge_spacing = 0.11
    badge_y = 0.05

    total_width = (n_elements - 1) * badge_spacing
    start_x = 0.5 - total_width / 2

    renderer = plotter.renderer

    for i, elem in enumerate(elements):
        color_hex = get_element_color(elem)
        rgb = hex_to_rgb(color_hex)
        rgb_norm = (rgb[0] / 255, rgb[1] / 255, rgb[2] / 255)

        if oxidation_states and elem in oxidation_states:
            oxi = oxidation_states[elem]
            if oxi > 0:
                label = f" {elem}{oxi}+ "
            elif oxi < 0:
                label = f" {elem}{abs(oxi)}- "
            else:
                label = f" {elem} "
        else:
            label = f" {elem} "

        x_pos = start_x + i * badge_spacing

        text_actor = vtk.vtkTextActor()
        text_actor.SetInput(label)
        text_actor.GetPositionCoordinate().SetCoordinateSystemToNormalizedViewport()
        text_actor.SetPosition(x_pos, badge_y)

        text_prop = text_actor.GetTextProperty()
        text_prop.SetFontSize(18)
        text_prop.SetBold(True)
        text_prop.SetFontFamilyToArial()
        text_prop.SetJustificationToCentered()
        text_prop.SetVerticalJustificationToCentered()

        brightness = (rgb[0] * 0.299 + rgb[1] * 0.587 + rgb[2] * 0.114)
        if brightness > 150:
            text_prop.SetColor(0.1, 0.1, 0.1)
        else:
            text_prop.SetColor(1.0, 1.0, 1.0)

        text_prop.SetBackgroundColor(rgb_norm)
        text_prop.SetBackgroundOpacity(1.0)

        try:
            text_prop.SetCellOffset(8)
        except AttributeError:
            pass

        renderer.AddViewProp(text_actor)


def view_crystal(
    material: MaterialTD,
    mpdoc=None,
    supercell: Tuple[int, int, int] = (1, 1, 1),
    plotter: Optional[pv.Plotter] = None,
    formula: Optional[str] = None,
) -> None:
    """Visualize crystal structure with conventional cell and optional supercell expansion.

    Args:
        material: Typed material dict with structure and metadata.
        mpdoc: Optional Materials Project document (for oxidation states).
        supercell: Tuple of (nx, ny, nz) for supercell expansion.
        plotter: Optional existing PyVista plotter. Created if None.
        formula: Optional formula string to display in the viewer.
    """
    pmg_struct = build_pmg_structure(material["structure"])

    # Convert to Conventional Standard Structure
    sga = SpacegroupAnalyzer(pmg_struct)
    conv_struct = sga.get_conventional_standard_structure()

    conv_lattice = {
        "abc": tuple(conv_struct.lattice.abc),
        "angles": tuple(conv_struct.lattice.angles),
        "vectors": np.array(conv_struct.lattice.matrix),
        "volume": conv_struct.lattice.volume,
        "pbc": tuple(conv_struct.lattice.pbc),
    }

    conv_atoms = [
        {
            "element": site.specie.symbol,
            "frac": np.array(site.frac_coords),
            "cart": np.array(site.coords),
        }
        for site in conv_struct.sites
    ]

    calc_structure = {
        "lattice": conv_lattice,
        "atoms": conv_atoms,
        "nsites": len(conv_atoms),
        "elements": [str(e) for e in conv_struct.composition.elements],
    }

    # Apply supercell if requested
    if supercell != (1, 1, 1):
        calc_structure = make_supercell(calc_structure, *supercell)

    viz_structure = calc_structure

    own_plotter = plotter is None
    if own_plotter:
        plotter = pv.Plotter()
        plotter.set_background("white")

    oxi_map = {}
    if mpdoc is not None:
        oxi_map = parse_oxidation_states(mpdoc.possible_species)

    # Calculate bonds
    bonds = find_bonds_crystalnn(calc_structure, oxidation_states=oxi_map)

    # Build atom position lookup for bond filtering
    structure_atoms = calc_structure["atoms"]
    lattice_vecs = np.array(calc_structure["lattice"]["vectors"])

    CART_PRECISION = 2
    atom_position_lookup = {}
    for idx, atom in enumerate(structure_atoms):
        cart = np.array(atom["cart"])
        key = tuple(np.round(cart, CART_PRECISION))
        atom_position_lookup[key] = idx

    # Filter bonds: keep only bonds where both endpoints match structure atoms
    active_bonds = []
    for b in bonds:
        i, j = b["i"], b["j"]
        image = np.array(b["image"], dtype=float)

        ri = np.array(structure_atoms[i]["cart"])
        rj = np.array(structure_atoms[j]["cart"]) + image @ lattice_vecs

        key_i = tuple(np.round(ri, CART_PRECISION))
        key_j = tuple(np.round(rj, CART_PRECISION))

        if key_i in atom_position_lookup and key_j in atom_position_lookup:
            active_bonds.append(b)

    # Build and add atoms
    atom_poly = build_atoms(viz_structure)
    add_atoms(plotter, atom_poly)

    # Add bonds
    add_bonds(plotter, calc_structure, active_bonds)

    add_unit_cell(plotter, calc_structure["lattice"])
    enable_hover(plotter, atom_poly, viz_structure, oxidation_states=oxi_map)
    add_element_legend(plotter, viz_structure["elements"], oxidation_states=oxi_map)

    # Add compound name in top right corner
    if formula:
        plotter.add_text(
            formula,
            position="upper_right",
            font_size=14,
            color="black",
            name="compound_name",
        )

    # Add hint about closing the window
    plotter.add_text(
        "Press 'q' to close",
        position="lower_left",
        font_size=10,
        color="gray",
        name="close_hint",
    )

    plotter.show_axes()
    plotter.enable_anti_aliasing("ssaa")
    plotter.show()

    # Aggressive cleanup to ensure window closes on macOS
    try:
        if hasattr(plotter, 'ren_win') and plotter.ren_win is not None:
            plotter.ren_win.Finalize()
    except Exception:
        pass
    try:
        plotter.deep_clean()
    except Exception:
        pass
    try:
        if plotter.iren is not None:
            plotter.iren.TerminateApp()
            plotter.iren.GetRenderWindow().Finalize()
    except Exception:
        pass
    plotter.close()
    del plotter
    gc.collect()
