import math
import sys
from pathlib import Path

import vtkmodules.vtkRenderingOpenGL2  # noqa
from vtkmodules.vtkCommonCore import vtkPoints
from vtkmodules.vtkCommonDataModel import (
    vtkCellArray, vtkPolyData, vtkPolyLine
)
from vtkmodules.vtkFiltersSources import vtkPlaneSource
from vtkmodules.vtkInteractionStyle import vtkInteractorStyleTrackballCamera
from vtkmodules.vtkRenderingCore import (
    vtkRenderer, vtkRenderWindow, vtkRenderWindowInteractor
)

from vtk import (
    vtkBoreholeRepresentation, vtkBoreholeWidget, vtkPolyDataCollection
)

from trame.app import TrameApp
from trame.decorators import change
from trame.ui.vuetify3 import SinglePageLayout
from trame.widgets import vuetify3, vtk, vtklocal, client

def make_trajectory():
    points = vtkPoints()
    points.InsertNextPoint(0.0, 0.0, 0.0)
    points.InsertNextPoint(0.2, 0.0, 10.0)
    points.InsertNextPoint(1.0, 0.5, 20.0)
    points.InsertNextPoint(2.0, 1.0, 30.0)
    points.InsertNextPoint(2.2, 1.5, 40.0)

    line = vtkPolyLine()
    line.GetPointIds().SetNumberOfIds(points.GetNumberOfPoints())
    for i in range(points.GetNumberOfPoints()):
        line.GetPointIds().SetId(i, i)

    lines = vtkCellArray()
    lines.InsertNextCell(line)

    trajectory = vtkPolyData()
    trajectory.SetPoints(points)
    trajectory.SetLines(lines)
    return trajectory


def make_structural_surface(z_base, ax, ay, wiggle, resolution):
    plane = vtkPlaneSource()
    plane.SetOrigin(-20.0, -20.0, 0.0)
    plane.SetPoint1(20.0, -20.0, 0.0)
    plane.SetPoint2(-20.0, 20.0, 0.0)
    plane.SetXResolution(resolution)
    plane.SetYResolution(resolution)
    plane.Update()

    pd = vtkPolyData()
    pd.DeepCopy(plane.GetOutput())
    pts = pd.GetPoints()
    for i in range(pts.GetNumberOfPoints()):
        p = list(pts.GetPoint(i))
        p[2] = (z_base + ax * p[0] + ay * p[1] +
                wiggle * math.sin(0.15 * p[0]) * math.cos(0.15 * p[1]))
        pts.SetPoint(i, p)
    pts.Modified()
    return pd


# -----------------------------------------------------------------------------
# Main Trame App Class
# -----------------------------------------------------------------------------

class BoreholeDemo(TrameApp):
    def __init__(self, server=None):
        super().__init__(server)
        self._setup_vtk()
        self._build_ui()

    def _setup_vtk(self):
        renderer = vtkRenderer()
        renderer.SetBackground(0.12, 0.14, 0.18)

        renderWindow = vtkRenderWindow()
        renderWindow.AddRenderer(renderer)
        renderWindow.OffScreenRenderingOn()

        interactor = vtkRenderWindowInteractor()
        interactor.SetRenderWindow(renderWindow)
        style = vtkInteractorStyleTrackballCamera()
        interactor.SetInteractorStyle(style)

        # Data generation
        trajectory = make_trajectory()
        surfaces = vtkPolyDataCollection()
        surfaces.AddItem(make_structural_surface(8.0, 0.03, -0.01, 0.8, 60))
        surfaces.AddItem(make_structural_surface(20.0, -0.02, 0.025, 0.5, 60))
        surfaces.AddItem(make_structural_surface(33.0, 0.015, 0.015, 0.7, 60))

        rep = vtkBoreholeRepresentation()
        rep.SetRenderer(renderer)
        rep.SetInputData(trajectory)
        rep.SetStructuralSurfaces(surfaces)
        rep.SetSurfaceDx(0.5)
        rep.SetSurfaceDy(0.5)
        rep.SetRadius(0.6)
        rep.SetGlyphRadius(0.5)
        rep.SetTopPosition(0.15)
        rep.SetBottomPosition(0.85)
        rep.BuildRepresentation()

        widget = vtkBoreholeWidget()
        widget.SetInteractor(interactor)
        widget.SetCurrentRenderer(renderer)
        widget.SetPriority(1.0)
        widget.SetRepresentation(rep)
        widget.SetEnabled(1)

        renderer.ResetCamera()
        renderWindow.Render()

        self.render_window = renderWindow
        self.widget = widget
        self.rep = rep

    def _build_ui(self):
        with SinglePageLayout(self.server) as layout:
            self.ui = layout

            layout.title.set_text("Borehole Widget Demo (WASM)")
            layout.icon.click = self.ctrl.view_reset_camera
            layout.toolbar.density = "compact"

            with layout.content:
                with vtklocal.LocalView(
                    self.render_window,
                    throttle_rate=20,
                    ctx_name="borehole_view",
                ) as view:
                    self.ctrl.view_update = view.update_throttle
                    self.ctrl.view_reset_camera = view.reset_camera

                    # Register VTK objects and return ID for JavaScript
                    widget_id = view.register_vtk_object(self.widget)
                    rep_id = view.register_vtk_object(self.rep)

            # JavaScript 
            client.Script(f"""
                const viewRef = trame.refs.borehole_view;
                if (viewRef && viewRef.onReady) {{
                    viewRef.onReady(() => {{
                        const widget = viewRef.getVtkObject({widget_id});
                        const representation = viewRef.getVtkObject({rep_id});
                        console.log("Borehole widget registered, id:", widget?.id);
                        console.log("Representation registered, id:", representation?.id);
                    }});
                }} else {{
                    console.warn("View not ready");
                }}
            """)

    def export(self, filename="borehole.wazex"):
        """Save scene."""
        if hasattr(self, 'ctx') and hasattr(self.ctx, 'borehole_view'):
            self.ctx.borehole_view.save(filename)
            print(f"Scene saved to {filename}")
        else:
            print("No view to export")


def main():
    app = BoreholeDemo()
    app.server.start()


if __name__ == "__main__":
    main()