"""Provides for the linking and/or appending of geometry nodes trees from vendored libraries included in Blender maxwell."""

import enum
from pathlib import Path

import bpy

from blender_maxwell import contracts as ct
from blender_maxwell.utils import logger

log = logger.get(__name__)


####################
# - GeoNodes Specification
####################
# GeoNodes Paths
## Internal
GN_INTERNAL_PATH = ct.addon.PATH_ASSETS / 'internal'
GN_INTERNAL_INPUTS_PATH = GN_INTERNAL_PATH / 'input'
GN_INTERNAL_SOURCES_PATH = GN_INTERNAL_PATH / 'source'
GN_INTERNAL_STRUCTURES_PATH = GN_INTERNAL_PATH / 'structure'
GN_INTERNAL_MONITORS_PATH = GN_INTERNAL_PATH / 'monitor'
GN_INTERNAL_SIMULATIONS_PATH = GN_INTERNAL_PATH / 'simulation'

## Structures
GN_STRUCTURES_PATH = ct.addon.PATH_ASSETS / 'structures'
GN_STRUCTURES_PRIMITIVES_PATH = GN_STRUCTURES_PATH / 'primitives'
GN_STRUCTURES_ARRAYS_PATH = GN_STRUCTURES_PATH / 'arrays'


class GeoNodes(enum.StrEnum):
	"""Defines the names of available GeoNodes groups vendored as part of Blender Maxwell.

	The values define both name of the .blend file containing the GeoNodes group, and the GeoNodes group itself.
	"""

	# Node Previews
	## Input
	InputConstantPhysicalPol = '_input_constant_physical_pol'
	## Source
	SourcePointDipole = '_source_point_dipole'
	SourcePlaneWave = '_source_plane_wave'
	SourceUniformCurrent = '_source_uniform_current'
	SourceTFSF = '_source_tfsf'
	SourceGaussianBeam = '_source_gaussian_beam'
	SourceAstigmaticGaussianBeam = '_source_astigmatic_gaussian_beam'
	SourceMode = '_source_mode'
	SourceEHArray = '_source_eh_array'
	SourceEHEquivArray = '_source_eh_equiv_array'
	## Structure
	StructurePrimitivePlane = '_structure_primitive_plane'
	StructurePrimitiveBox = '_structure_primitive_box'
	StructurePrimitiveSphere = '_structure_primitive_sphere'
	StructurePrimitiveCylinder = '_structure_primitive_cylinder'
	StructurePrimitiveRing = '_structure_primitive_ring'
	StructurePrimitiveCapsule = '_structure_primitive_capsule'
	StructurePrimitiveCone = '_structure_primitive_cone'
	## Monitor
	MonitorEHField = '_monitor_eh_field'
	MonitorPowerFlux = '_monitor_power_flux'
	MonitorEpsTensor = '_monitor_eps_tensor'
	MonitorDiffraction = '_monitor_diffraction'
	MonitorProjCartEHField = '_monitor_proj_eh_field'
	MonitorProjAngEHField = '_monitor_proj_ang_eh_field'
	MonitorProjKSpaceEHField = '_monitor_proj_k_space_eh_field'
	## Simulation
	SimulationSimDomain = '_simulation_sim_domain'
	SimulationBoundConds = '_simulation_bound_conds'
	SimulationBoundCondPML = '_simulation_bound_cond_pml'
	SimulationBoundCondPEC = '_simulation_bound_cond_pec'
	SimulationBoundCondPMC = '_simulation_bound_cond_pmc'
	SimulationBoundCondBloch = '_simulation_bound_cond_bloch'
	SimulationBoundCondPeriodic = '_simulation_bound_cond_periodic'
	SimulationBoundCondAbsorbing = '_simulation_bound_cond_absorbing'
	SimulationSimGrid = '_simulation_sim_grid'
	SimulationSimGridAxisAuto = '_simulation_sim_grid_axis_auto'
	SimulationSimGridAxisManual = '_simulation_sim_grid_axis_manual'
	SimulationSimGridAxisUniform = '_simulation_sim_grid_axis_uniform'
	SimulationSimGridAxisArray = '_simulation_sim_grid_axis_array'

	# Structures
	## Primitives
	PrimitiveBox = 'box'
	PrimitiveSphere = 'sphere'
	PrimitiveCylinder = 'cylinder'
	PrimitiveRing = 'ring'
	## Arrays
	ArrayRing = 'array_ring'

	@property
	def dedicated_node_type(self) -> ct.BLImportMethod:
		"""Deduces the denode type that implements a vendored GeoNodes tree (usually just "GeoNodes Structure").

		Generally, "GeoNodes Structure' is the generic triangle-mesh node that can do everything.
		Indeed, one could make perfectly useful simulations exclusively using triangle meshes.

		However, when nodes that use geometry directly supported by `Tidy3D` might be more performant, and support features (ex. differentiable parameters) critical to some simulations.

		To bridge the gap, this method provides a regularized hard-coded method of getting whichever node type is most appropriate for the given structure.

		Warnings:
			**The strings be manually checked, statically**.
			Since we don't have easy access to the `contracts` within the sim node tree, these are just plain strings.

			It's not ideal, definitely a lazy workaround, but hey.

		Returns:
			The node type (as a string) that should be created to expose the GeoNodes group.
		"""
		dedicated_node_map = {
			GeoNodes.PrimitiveBox: 'BoxStructureNodeType',
			GeoNodes.PrimitiveSphere: 'SphereStructureNodeType',
			GeoNodes.PrimitiveCylinder: 'CylinderStructureNodeType',
		}

		if dedicated_node_map.get(self) is None:
			return 'GeoNodesStructureNodeType'

		return dedicated_node_map[self]

	@property
	def import_method(self) -> ct.BLImportMethod:
		"""Deduces whether a vendored GeoNodes tree should be linked or appended.

		Currently, everything is linked.
		If the user wants to modify a group, they can always make a local copy of a linked group and just use that.

		Returns:
			Either 'link' or 'append'.
		"""
		return 'link'

	@property
	def parent_path(self) -> Path:
		"""Deduces the parent path of a vendored GeoNodes tree.

		Warnings:
			The guarantee of same-named `.blend` and tree name are not explicitly checked.
			Also, the validity of GeoNodes tree interface parameters is also unchecked; this must be manually coordinated between the driving node, and the imported tree.

		Returns:
			The parent path of the given GeoNodes tree.

			A `.blend` file of the same name as the value of the enum must be **guaranteed** to exist as a direct child of the returned path, and must be **guaranteed** to contain a GeoNodes tree of the same name.
		"""
		GN = GeoNodes
		return {
			# Node Previews
			## Input
			GN.InputConstantPhysicalPol: GN_INTERNAL_INPUTS_PATH,
			## Source
			GN.SourcePointDipole: GN_INTERNAL_SOURCES_PATH,
			GN.SourcePlaneWave: GN_INTERNAL_SOURCES_PATH,
			GN.SourceUniformCurrent: GN_INTERNAL_SOURCES_PATH,
			GN.SourceTFSF: GN_INTERNAL_SOURCES_PATH,
			GN.SourceGaussianBeam: GN_INTERNAL_SOURCES_PATH,
			GN.SourceAstigmaticGaussianBeam: GN_INTERNAL_SOURCES_PATH,
			GN.SourceMode: GN_INTERNAL_SOURCES_PATH,
			GN.SourceEHArray: GN_INTERNAL_SOURCES_PATH,
			GN.SourceEHEquivArray: GN_INTERNAL_SOURCES_PATH,
			## Structure
			GN.StructurePrimitivePlane: GN_INTERNAL_STRUCTURES_PATH,
			GN.StructurePrimitiveBox: GN_INTERNAL_STRUCTURES_PATH,
			GN.StructurePrimitiveSphere: GN_INTERNAL_STRUCTURES_PATH,
			GN.StructurePrimitiveCylinder: GN_INTERNAL_STRUCTURES_PATH,
			GN.StructurePrimitiveRing: GN_INTERNAL_STRUCTURES_PATH,
			GN.StructurePrimitiveCapsule: GN_INTERNAL_STRUCTURES_PATH,
			GN.StructurePrimitiveCone: GN_INTERNAL_STRUCTURES_PATH,
			## Monitor
			GN.MonitorEHField: GN_INTERNAL_STRUCTURES_PATH,
			GN.MonitorPowerFlux: GN_INTERNAL_STRUCTURES_PATH,
			GN.MonitorEpsTensor: GN_INTERNAL_STRUCTURES_PATH,
			GN.MonitorDiffraction: GN_INTERNAL_STRUCTURES_PATH,
			GN.MonitorProjCartEHField: GN_INTERNAL_STRUCTURES_PATH,
			GN.MonitorProjAngEHField: GN_INTERNAL_STRUCTURES_PATH,
			GN.MonitorProjKSpaceEHField: GN_INTERNAL_STRUCTURES_PATH,
			## Simulation
			GN.SimulationSimDomain: GN_INTERNAL_SIMULATIONS_PATH,
			GN.SimulationBoundConds: GN_INTERNAL_SIMULATIONS_PATH,
			GN.SimulationBoundCondPML: GN_INTERNAL_SIMULATIONS_PATH,
			GN.SimulationBoundCondPEC: GN_INTERNAL_SIMULATIONS_PATH,
			GN.SimulationBoundCondPMC: GN_INTERNAL_SIMULATIONS_PATH,
			GN.SimulationBoundCondBloch: GN_INTERNAL_SIMULATIONS_PATH,
			GN.SimulationBoundCondPeriodic: GN_INTERNAL_SIMULATIONS_PATH,
			GN.SimulationBoundCondAbsorbing: GN_INTERNAL_SIMULATIONS_PATH,
			GN.SimulationSimGrid: GN_INTERNAL_SIMULATIONS_PATH,
			GN.SimulationSimGridAxisAuto: GN_INTERNAL_SIMULATIONS_PATH,
			GN.SimulationSimGridAxisManual: GN_INTERNAL_SIMULATIONS_PATH,
			GN.SimulationSimGridAxisUniform: GN_INTERNAL_SIMULATIONS_PATH,
			GN.SimulationSimGridAxisArray: GN_INTERNAL_SIMULATIONS_PATH,
			# Structures
			## Primitives
			GN.PrimitiveBox: GN_STRUCTURES_PRIMITIVES_PATH,
			GN.PrimitiveRing: GN_STRUCTURES_PRIMITIVES_PATH,
			GN.PrimitiveSphere: GN_STRUCTURES_PRIMITIVES_PATH,
			## Arrays
			GN.ArrayRing: GN_STRUCTURES_ARRAYS_PATH,
		}[self]


####################
# - Import GeoNodes (Link/Append)
####################
def import_geonodes(
	_geonodes: GeoNodes,
	force_append: bool = False,
) -> bpy.types.GeometryNodeGroup:
	"""Given vendored GeoNodes tree link/append and return the local datablock.

	- `GeoNodes.import_method` is used to determine whether it should be linked or imported.
	- `GeoNodes.parent_path` is used to determine

	Parameters:
		geonodes: The (name of the) GeoNodes group, which ships with Blender Maxwell.

	Returns:
		A GeoNodes group available in the current .blend file, which can ex. be attached to a 'GeoNodes Structure' node.
	"""
	# Parse Input
	geonodes = GeoNodes(_geonodes)
	import_method = geonodes.import_method

	# Linked: Don't Re-Link
	if import_method == 'link' and geonodes in bpy.data.node_groups:
		return bpy.data.node_groups[geonodes]

	filename = geonodes
	filepath = str(geonodes.parent_path / (geonodes + '.blend') / 'NodeTree' / geonodes)
	directory = filepath.removesuffix(geonodes)
	log.info(
		'%s GeoNodes (filename=%s, directory=%s, filepath=%s)',
		'Linking' if import_method == 'link' else 'Appending',
		filename,
		directory,
		filepath,
	)
	bpy.ops.wm.append(
		filepath=filepath,
		directory=directory,
		filename=filename,
		check_existing=False,
		set_fake=True,
		link=import_method == 'link',
	)

	return bpy.data.node_groups[geonodes]


####################
# - GeoNodes Asset Shelf Panel for MaxwellSimTree
####################
ASSET_PANEL_NAME: str = 'blender_maxwell__node_asset_shelf'
WM_ACTIVE_NODE_ASSETS: str = 'blender_maxwell__active_node_asset_list'
WM_ACTIVE_ASSET_IDX: str = 'blender_maxwell__active_node_asset_index'


class NodeAssetPanel(bpy.types.Panel):
	"""Provides a panel that displays vendored, usable GeoNodes trees, and enables drag-and-drop support to easily drag them into the simulation node editor."""

	## TODO: Provide an option that forces appending? So that users can modify from a baseline. Just watch out - dealing with overlaps isn't trivial.

	bl_idname = ct.PanelType.NodeAssetPanel
	bl_label = 'Node GeoNodes Asset Panel'
	bl_space_type = 'NODE_EDITOR'
	bl_region_type = 'UI'
	bl_category = 'Assets'

	@classmethod
	def poll(cls, context: bpy.types.Context) -> bool:
		"""Require window manager properties to show node assets.

		Notes:
			Run by Blender when trying to show a panel.

		Returns:
			Whether the panel can show.
		"""
		## TODO: Check that we're in a MaxwellSim node tree as well.
		wm = context.window_manager
		return hasattr(wm, WM_ACTIVE_NODE_ASSETS) and hasattr(wm, WM_ACTIVE_ASSET_IDX)

	def draw(self, context: bpy.types.Context) -> None:
		"""Draw the asset library w/drag support.

		Notes:
			Run by Blender when the panel needs to be displayed.

		Parameters:
			context: The Blender context object.
				Must contain `context.window_manager` and `context.workspace`.
		"""
		layout = self.layout

		# Draw the Asset Panel w/Drag Support
		_activate_op_props, _drag_op_props = layout.template_asset_view(
			# list_id
			## Identifies this particular UI-bound asset list.
			## Must be unique, otherwise weird things happen if invoked shown twice.
			## We use a custom name, presuming only one tree is shown at a time.
			## TODO: Allow several trees to show at once.
			ASSET_PANEL_NAME,
			# asset_library_[dataptr|propname]
			## Where to find the asset library to pull assets from.
			## We pull it from a global
			context.workspace,
			'asset_library_reference',
			# assets_[dataptr|propname]
			## The actual assets to provide user access to.
			context.window_manager,
			WM_ACTIVE_NODE_ASSETS,
			# assets_[dataptr|propname]
			## The currently selected asset to highlight.
			context.window_manager,
			WM_ACTIVE_ASSET_IDX,
			# draw_operator
			## The operator to invoke() whenever the user **starts** dragging an asset.
			drag_operator=ct.OperatorType.GeoNodesToStructureNode,
			# Other Options
			## The currently selected asset to highlight.
			# display_options={'NO_LIBRARY'},
		)


####################
# - Append GeoNodes Operator
####################
def get_view_location(
	region: bpy.types.Region, coords: tuple[float, float], ui_scale: float
) -> tuple[float, float]:
	"""Given a pair of coordinates defined on a Blender region, project the coordinates to the corresponding `bpy.types.View2D`.

	Parameters:
		region: The region within which the coordinates are defined.
		coords: The coordinates within the region, ex. produced during a mouse click event.
		ui_scale: Scaling factor applied to pixels, which compensates for screens with differing DPIs.
			We must divide the coordinates we receive by this float to get the "real" $x,y$ coordinates.

	Returns:
		The $x,y$ coordinates within the region's `bpy.types.View2D`, correctly scaled with the current interface size/dpi.
	"""
	x, y = region.view2d.region_to_view(*coords)
	return x / ui_scale, y / ui_scale


class GeoNodesToStructureNode(bpy.types.Operator):
	"""Operator allowing the user to append a vendored GeoNodes tree for use in a simulation."""

	bl_idname = ct.OperatorType.GeoNodesToStructureNode
	bl_label = 'GeoNodes to Structure Node'
	bl_description = 'Drag-and-drop operator'
	bl_options = frozenset({'REGISTER'})

	@classmethod
	def poll(cls, context: bpy.types.Context) -> bool:
		"""Defines when the operator can be run.

		Returns:
			Whether the operator can be run.
		"""
		return context.asset is not None

	####################
	# - Properties
	####################
	_asset: bpy.types.AssetRepresentation | None = None

	####################
	# - Execution
	####################
	def invoke(
		self, context: bpy.types.Context, _: bpy.types.Event
	) -> set[ct.BLOperatorStatus]:
		"""Commences the drag-and-drop of a GeoNodes asset.

		- Starts the modal timer, which listens for a mouse-release event.
		- Changes the mouse cursor to communicate the "dragging" state to the user.

		Notes:
			Run via `NodeAssetPanel` when the user starts dragging a vendored GeoNodes asset.

			The asset being dragged can be found in `context.asset`.

		Returns:
			Indication that there is a modal running.
		"""
		self._asset = context.asset  ## Guaranteed non-None by poll()

		# Register Modal Operator
		context.window_manager.modal_handler_add(self)
		context.area.tag_redraw()

		# Set Mouse Cursor
		context.window.cursor_modal_set('CROSS')

		return {'RUNNING_MODAL'}

	def modal(
		self, context: bpy.types.Context, event: bpy.types.Event
	) -> ct.BLOperatorStatus:
		"""When LMB is released, creates a GeoNodes Structure node.

		Runs in response to events in the node editor while dragging an asset from the side panel.
		"""
		# No Asset: Do Nothing
		asset = self._asset
		if asset is None:
			return {'PASS_THROUGH'}

		# Released LMB: Add Structure Node
		## The user 'dropped' the asset!
		## Now, we create an appropriate "Structure" node corresponding to the dropped asset.
		if event.type == 'LEFTMOUSE' and event.value == 'RELEASE':
			log.debug('Dropped Asset: %s', asset.name)

			# Load Editor Region
			area = context.area
			editor_window_regions = [
				region for region in area.regions.values() if region.type == 'WINDOW'
			]
			if editor_window_regions:
				log.debug(
					'Selected Node Editor Region out of %s Options',
					str(len(editor_window_regions)),
				)
				editor_region = editor_window_regions[0]
			else:
				msg = 'No valid editor region in context where asset was dropped'
				raise RuntimeError(msg)

			# Check if Mouse Coordinates are:
			## - INSIDE of Node Editor
			## - INSIDE of Node Editor's WINDOW Region
			# Check Mouse Coordinates against Node Editor
			## The asset must be dropped inside the Maxwell Sim node editor.
			if (
				## Check: Mouse in Area (contextual)
				(event.mouse_x >= area.x and event.mouse_x < area.x + area.width)
				and (event.mouse_y >= area.y and event.mouse_y < area.y + area.height)
			) and (
				## Check: Mouse in Node Editor Region
				(
					event.mouse_x >= editor_region.x
					and event.mouse_x < editor_region.x + editor_region.width
				)
				and (
					event.mouse_y >= editor_region.y
					and event.mouse_y < editor_region.y + editor_region.height
				)
			):
				space = context.space_data
				node_tree = space.node_tree

				# Computing GeoNodes View Location
				## 1. node_tree.cursor_location gives clicked loc, not released.
				## 2. event.mouse_region_* has inverted x wrt. event.mouse_*.
				##    -> View2D.region_to_view expects the event.mouse_* order.
				##    -> Is it a bug? Who knows!
				## 3. We compute it manually, to avoid the jank.
				node_location = get_view_location(
					editor_region,
					[
						event.mouse_x - editor_region.x,
						event.mouse_y - editor_region.y,
					],
					context.preferences.system.ui_scale,
				)

				# Create GeoNodes Structure Node
				## 1. Deselect other nodes
				## 2. Select the new one
				## 3. Move it into place
				## 4. Redraw (so we see the new node right away)
				geonodes = GeoNodes(asset.name)

				log.info(
					'Creating Node "%s" at (%d, %d)',
					geonodes.dedicated_node_type,
					*tuple(node_location),
				)
				bpy.ops.node.select_all(action='DESELECT')
				node = node_tree.nodes.new(geonodes.dedicated_node_type)
				node.select = True
				node.location.x = node_location[0]
				node.location.y = node_location[1]
				context.area.tag_redraw()

				# GeoNodes Structure Specific Setup
				## Since the node doesn't itself handle the structure, we must.
				## We just import the GN tree, then attach the data block to the node.
				if geonodes.dedicated_node_type == 'GeoNodesStructureNodeType':
					geonodes_data = import_geonodes(asset.name)
					node.inputs['GeoNodes'].value = geonodes_data
					## TODO: Is this too presumptuous? Or a fine little hack?

			# Restore the Pre-Modal Mouse Cursor Shape
			context.window.cursor_modal_restore()
			return {'FINISHED'}

		return {'RUNNING_MODAL'}


####################
# - Blender Registration
####################
ASSET_LIB_POSTFIX: str = ' | BLMaxwell'
ASSET_LIB_SPECS: dict[str, Path] = {
	'Primitives': GN_STRUCTURES_PRIMITIVES_PATH,
	'Arrays': GN_STRUCTURES_ARRAYS_PATH,
}


@bpy.app.handlers.persistent
def initialize_asset_libraries(_: bpy.types.Scene):
	"""Before loading a `.blend` file, ensure that the WindowManager properties relied on by NodeAssetPanel are available.

	- Several asset libraries, defined under the global `ASSET_LIB_SPECS`, are added/replaced such that the name:path map is respected.
	- Existing asset libraries are left entirely alone.
	- The names attached to `bpy.types.WindowManager` are specifically the globals `WM_ACTIVE_NODE_ASSETS` and `WM_ACTIVE_ASSET_IDX`.

	Warnings:
		**Changing the name of an asset library will leave it dangling on all `.blend` files**.

		Therefore, an addon-specific postfix `ASSET_LIB_POSTFIX` is appended to all asset library names.
		**This postfix must never, ever change, across any and all versions of the addon**.

		Again, if it does, **all Blender installations having ever used the addon will have a dangling asset library until the user manually notices + removes it**.
	"""
	asset_libraries = bpy.context.preferences.filepaths.asset_libraries

	# Guarantee All Asset Libraries
	for _asset_lib_name, asset_lib_path in ASSET_LIB_SPECS.items():
		asset_lib_name = _asset_lib_name + ASSET_LIB_POSTFIX

		# Remove Existing Asset Library
		asset_library_idx = asset_libraries.find(asset_lib_name)
		if asset_library_idx != -1 and asset_libraries[asset_lib_name].path != str(
			asset_lib_path
		):
			log.debug('Removing Asset Library: %s', asset_lib_name)
			bpy.ops.preferences.asset_library_remove(asset_library_idx)

		# Add Asset Library
		if asset_lib_name not in asset_libraries:
			log.debug('Add Asset Library: %s', asset_lib_name)

			bpy.ops.preferences.asset_library_add()
			asset_library = asset_libraries[-1]  ## Was added to the end
			asset_library.name = asset_lib_name
			asset_library.path = str(asset_lib_path)

	# Set WindowManager Props
	## Set Active Assets Collection
	setattr(
		bpy.types.WindowManager,
		WM_ACTIVE_NODE_ASSETS,
		bpy.props.CollectionProperty(type=bpy.types.AssetHandle),
	)
	## Set Active Assets Collection
	setattr(
		bpy.types.WindowManager,
		WM_ACTIVE_ASSET_IDX,
		bpy.props.IntProperty(),
	)


bpy.app.handlers.load_pre.append(initialize_asset_libraries)

BL_REGISTER = [
	NodeAssetPanel,
	GeoNodesToStructureNode,
]

BL_HOTKEYS = []