import bpy
import bl_math
import math
from mathutils import Matrix, Vector, Quaternion
from bpy_extras import anim_utils

from ..Core import Core
from ..InputManager import InputItem
from .. import Utils as util


# ******************************************************************** #
class AnimationClip:
	"""Extracted fcurve data from an armature's action."""

	def __init__(self, armature, name=None):
		self.name = name or armature.name
		self.bone_curves = {}		# bone_name -> { prop: { array_index: fcurve } }
		self.frame_start = 0.0
		self.frame_end = 0.0
		self.frame_count = 0.0
		self._extract(armature)

	def _extract(self, armature):
		ad = armature.animation_data
		if not ad or not ad.action:
			raise ValueError(f"AnimationClip: Armature '{armature.name}' has no action")

		cb = anim_utils.action_get_channelbag_for_slot(ad.action, ad.action_slot)
		if not cb:
			raise ValueError(f"AnimationClip: No channelbag found for '{armature.name}'")

		all_frames = []
		prefix = 'pose.bones["'

		for fc in cb.fcurves:
			dp = fc.data_path
			if not dp.startswith(prefix):
				continue

			end = dp.index('"]', len(prefix))
			bone_name = dp[len(prefix):end]
			prop = dp[end + 3:]

			if bone_name not in self.bone_curves:
				self.bone_curves[bone_name] = {}
			if prop not in self.bone_curves[bone_name]:
				self.bone_curves[bone_name][prop] = {}

			self.bone_curves[bone_name][prop][fc.array_index] = fc

			if fc.keyframe_points:
				all_frames.append(fc.keyframe_points[0].co[0])
				all_frames.append(fc.keyframe_points[-1].co[0])

		if all_frames:
			self.frame_start = min(all_frames)
			self.frame_end = max(all_frames)
			self.frame_count = self.frame_end - self.frame_start

	def sample_pose(self, phase):
		frame = self.frame_start + phase * self.frame_count
		pose = {}

		for bone_name, curves in self.bone_curves.items():
			loc = None
			rot = None

			loc_curves = curves.get('location')
			if loc_curves:
				loc = Vector((
					loc_curves[0].evaluate(frame) if 0 in loc_curves else 0.0,
					loc_curves[1].evaluate(frame) if 1 in loc_curves else 0.0,
					loc_curves[2].evaluate(frame) if 2 in loc_curves else 0.0,
				))

			rot_curves = curves.get('rotation_quaternion')
			if rot_curves:
				rot = Quaternion((
					rot_curves[0].evaluate(frame) if 0 in rot_curves else 1.0,
					rot_curves[1].evaluate(frame) if 1 in rot_curves else 0.0,
					rot_curves[2].evaluate(frame) if 2 in rot_curves else 0.0,
					rot_curves[3].evaluate(frame) if 3 in rot_curves else 0.0,
				))

			pose[bone_name] = (loc, rot)

		return pose


# ******************************************************************** #
def blend_poses(pose_a, pose_b, factor):
	result = {}
	all_bones = set(pose_a.keys()) | set(pose_b.keys())

	for bone_name in all_bones:
		loc_a, rot_a = pose_a.get(bone_name, (None, None))
		loc_b, rot_b = pose_b.get(bone_name, (None, None))

		if loc_a is not None and loc_b is not None:
			loc = loc_a.lerp(loc_b, factor)
		elif loc_b is not None:
			loc = loc_b.copy()
		elif loc_a is not None:
			loc = loc_a.copy()
		else:
			loc = None

		if rot_a is not None and rot_b is not None:
			if rot_a.dot(rot_b) < 0:
				rot_b = -rot_b
			rot = rot_a.slerp(rot_b, factor)
		elif rot_b is not None:
			rot = rot_b.copy()
		elif rot_a is not None:
			rot = rot_a.copy()
		else:
			rot = None

		result[bone_name] = (loc, rot)

	return result


def apply_pose(armature, pose):
	for bone_name, (loc, rot) in pose.items():
		pb = armature.pose.bones.get(bone_name)
		if pb is None:
			continue
		if loc is not None:
			pb.location = loc
		if rot is not None:
			pb.rotation_mode = 'QUATERNION'
			pb.rotation_quaternion = rot


# ******************************************************************** #
class Module_CharacterAnimation:
	key = 'character_animation'
	display_name = 'Character Animation alpha'
	tooltip = 'Locomotion blending for character models'
	enum_order = 2100

	# ---------------------------------------------------------------- #
	def __init__(self, context, core: Core):
		settings = context.scene.omnistep_modules.character_animation
		self.module_settings = settings
		self.scale = core.scene.state.scale
		self.enable_animation = core.scene.state.enable_animation

		if not settings.idle_armature or not settings.walk_armature or not settings.run_armature:
			raise ValueError("Character Animation: Idle, Walk, and Run armatures must all be set")

		# ---- Extract Animation Clips ----
		self.clip_idle = AnimationClip(settings.idle_armature, 'idle')
		self.clip_walk = AnimationClip(settings.walk_armature, 'walk')
		self.clip_run  = AnimationClip(settings.run_armature, 'run')

		# Strafe clips (optional — both must be set to enable strafing)
		self.has_strafe = settings.strafe_left_armature is not None and settings.strafe_right_armature is not None
		if self.has_strafe:
			self.clip_strafe_left  = AnimationClip(settings.strafe_left_armature, 'strafe_left')
			self.clip_strafe_right = AnimationClip(settings.strafe_right_armature, 'strafe_right')

		# Air clip (optional)
		self.has_air = settings.idle_air_armature is not None
		if self.has_air:
			self.clip_idle_air = AnimationClip(settings.idle_air_armature, 'idle_air')

		# ---- Create or Use Output Character ----
		if settings.character_armature is not None:
			self.character = settings.character_armature
			self.owns_character = False
		else:
			self.character = util.obj.duplicate(
				settings.idle_armature,
				linked=True,
				include_children=True,
				clear_animation=True
			)
			self.character.name = 'OStep.Character'
			self.owns_character = True

		self.backup_matrix = self.character.matrix_world.copy()

		# Match the idle armature's transform so the character starts in the correct orientation
		self.character.matrix_world = settings.idle_armature.matrix_world.copy()

		if self.character.animation_data:
			self.character.animation_data_clear()

		if self.owns_character:
			for child in util.obj.get_children(self.character):
				child['ostep.no_collision'] = True
				child['ostep.force_animated'] = True
				if child.type == 'MESH':
					for mod in child.modifiers:
						if mod.type == 'ARMATURE' and mod.object == settings.idle_armature:
							mod.object = self.character

		self.character['ostep.no_collision'] = True
		self.character['ostep.force_animated'] = True

		# ---- Source Transform ----
		src_mat = self.character.matrix_world
		self.rest_rotation = src_mat.to_quaternion()
		self.rest_scale = src_mat.to_scale()

		# Forward axis offset
		forward_offsets = { '+Y': 0.0, '-Y': math.pi, '+X': -math.pi / 2.0, '-X': math.pi / 2.0 }
		self.forward_offset = forward_offsets.get(settings.forward_axis, 0.0)

		# ---- State ----
		self.locomotion_phase = 0.0
		self.idle_phase = 0.0
		self.facing_yaw = core.player.state.yaw + self.forward_offset
		self.smooth_strafe_factor = 0.0

		# ---- Animation ----
		self.recorded_frames = []

		# ---- Movement Inertia ----
		# Base speeds from player (mouse wheel will modify these via ratio tracking)
		self.base_walk_speed = core.player.walk_speed
		self.base_run_speed = core.player.run_speed
		self.last_set_walk = core.player.walk_speed
		self.last_set_run = core.player.run_speed

		self.speed_ramp = 0.0			# 0–1, start/stop inertia
		self.run_blend = 0.0			# 0–1, walk↔run smooth transition
		self.air_blend = 0.0			# 0–1, ground↔air transition
		self.idle_air_phase = 0.0
		self.last_direction = Vector((0, 0, 0))

	# ---------------------------------------------------------------- #
	def update(self, context, core: Core):
		if core.input.state['restart'].phase == InputItem.DOWN:
			self.recorded_frames.clear()

		settings = self.module_settings
		dt = core.scene.state.current_timestep

		# ---- Detect mouse wheel speed adjustments ----
		# If player speeds differ from what we set last frame, mouse wheel changed them
		if self.last_set_walk > 0.001 and abs(core.player.walk_speed - self.last_set_walk) > 0.0001:
			ratio = core.player.walk_speed / self.last_set_walk
			self.base_walk_speed *= ratio
			self.base_run_speed *= ratio

		# ---- Input ----
		input_dir = core.input.state.direction
		has_input = input_dir.length > 0.001
		wants_run = core.player.always_run ^ core.input.state['speed'].pressed

		# ---- Speed ramp (start / stop) ----
		if has_input:
			self.last_direction = input_dir.normalized()
			if settings.acceleration_time > 0.001:
				self.speed_ramp = min(1.0, self.speed_ramp + dt / settings.acceleration_time)
			else:
				self.speed_ramp = 1.0
		else:
			if settings.deceleration_time > 0.001:
				self.speed_ramp = max(0.0, self.speed_ramp - dt / settings.deceleration_time)
			else:
				self.speed_ramp = 0.0

		# ---- Run blend (walk ↔ run smooth transition) ----
		run_target = 1.0 if wants_run else 0.0
		if settings.run_transition_time > 0.001:
			if self.run_blend < run_target:
				self.run_blend = min(run_target, self.run_blend + dt / settings.run_transition_time)
			else:
				self.run_blend = max(run_target, self.run_blend - dt / settings.run_transition_time)
		else:
			self.run_blend = run_target

		# ---- Compute effective speed and overwrite player ----
		effective_speed = bl_math.lerp(self.base_walk_speed, self.base_run_speed, self.run_blend) * self.speed_ramp
		core.player.walk_speed = effective_speed
		core.player.run_speed = effective_speed

		# Track what we set so we can detect wheel changes next frame
		self.last_set_walk = effective_speed
		self.last_set_run = effective_speed

		# ---- Inject direction for deceleration coasting ----
		if not has_input and self.speed_ramp > 0.001:
			core.input.state.direction = self.last_direction

	# ---------------------------------------------------------------- #
	def late_update(self, context, core: Core):
		settings = self.module_settings
		dt = core.scene.state.current_timestep

		# ---- Player State ----
		player_matrix = core.player.state.inertia.matrix @ core.player.state.root.matrix
		player_pos = player_matrix.to_translation()
		player_pos.z -= core.player.player_radius

		# Use real_velocity (collision-resolved) for animation — stops cycle when hitting a wall
		real_vel = core.player.state.real_velocity.copy()
		real_vel.z = 0.0
		horizontal_speed = real_vel.length
		horizontal_distance = horizontal_speed * dt

		# ---- Decompose velocity into forward / lateral ----
		yaw = core.player.state.yaw
		forward_dir = Vector((-math.sin(yaw), math.cos(yaw), 0.0))
		right_dir = Vector((math.cos(yaw), math.sin(yaw), 0.0))

		forward_speed = real_vel.dot(forward_dir)
		lateral_speed = real_vel.dot(right_dir)

		abs_forward = abs(forward_speed)
		abs_lateral = abs(lateral_speed)

		raw_strafe = abs_lateral / (abs_forward + abs_lateral) if self.has_strafe and (abs_forward + abs_lateral) > 0.001 else 0.0
		self.smooth_strafe_factor = bl_math.lerp(self.smooth_strafe_factor, raw_strafe, 0.3)

		# ---- Air / Ground transition ----
		if self.has_air:
			air_target = 0.0 if core.player.state.is_grounded else 1.0
			if settings.air_transition_time > 0.001:
				if self.air_blend < air_target:
					self.air_blend = min(air_target, self.air_blend + dt / settings.air_transition_time)
				else:
					self.air_blend = max(air_target, self.air_blend - dt / settings.air_transition_time)
			else:
				self.air_blend = air_target

		# ---- Blend Weights (from ramp state, not measured speed) ----
		locomotion_weight = self.speed_ramp
		idle_weight = 1.0 - locomotion_weight
		speed_factor = self.run_blend

		# ---- Step length (blended between forward and strafe) ----
		fwd_step = bl_math.lerp(settings.walk_step_length, settings.run_step_length, speed_factor) * self.scale
		if self.has_strafe:
			strafe_step = bl_math.lerp(settings.strafe_walk_step_length, settings.strafe_run_step_length, speed_factor) * self.scale
			active_step_length = bl_math.lerp(fwd_step, strafe_step, self.smooth_strafe_factor)
		else:
			active_step_length = fwd_step

		# ---- Phase ----
		self.idle_phase += 1.0 / max(self.clip_idle.frame_count, 1.0)
		self.idle_phase %= 1.0

		if self.has_air:
			self.idle_air_phase += 1.0 / max(self.clip_idle_air.frame_count, 1.0)
			self.idle_air_phase %= 1.0

		if active_step_length > 0.001:
			self.locomotion_phase += horizontal_distance / active_step_length
			self.locomotion_phase %= 1.0

		# ---- Sample & Blend ----
		pose_idle = self.clip_idle.sample_pose(self.idle_phase)

		# Forward: walk ↔ run
		pose_walk = self.clip_walk.sample_pose(self.locomotion_phase)
		pose_run  = self.clip_run.sample_pose(self.locomotion_phase)
		pose_forward = blend_poses(pose_walk, pose_run, speed_factor)

		# Strafe: pick left or right
		if self.has_strafe and self.smooth_strafe_factor > 0.001:
			if lateral_speed >= 0:
				clip_strafe = self.clip_strafe_right
			else:
				clip_strafe = self.clip_strafe_left

			pose_strafe = clip_strafe.sample_pose(self.locomotion_phase)
			pose_locomotion = blend_poses(pose_forward, pose_strafe, self.smooth_strafe_factor)
		else:
			pose_locomotion = pose_forward

		# Idle ↔ locomotion
		if locomotion_weight > 0.001:
			ground_pose = blend_poses(pose_idle, pose_locomotion, locomotion_weight)
		else:
			ground_pose = pose_idle

		# Ground ↔ air
		if self.has_air and self.air_blend > 0.001:
			pose_air = self.clip_idle_air.sample_pose(self.idle_air_phase)
			final_pose = blend_poses(ground_pose, pose_air, self.air_blend)
		else:
			final_pose = ground_pose

		apply_pose(self.character, final_pose)

		# ---- Position & Orientation ----
		target_yaw = core.player.state.yaw + self.forward_offset

		diff = (target_yaw - self.facing_yaw + math.pi) % (2.0 * math.pi) - math.pi
		self.facing_yaw += diff * settings.facing_smoothing

		yaw_mat = Quaternion((0, 0, 1), self.facing_yaw).to_matrix().to_4x4()
		rest_mat = self.rest_rotation.to_matrix().to_4x4()
		scale_mat = Matrix.Diagonal((*self.rest_scale, 1.0))

		self.character.matrix_world = Matrix.Translation(player_pos) @ yaw_mat @ rest_mat @ scale_mat

		# ---- Record ----
		if self.enable_animation:
			self.recorded_frames.append((
				core.scene.state.current_frame,
				self.character.matrix_world.copy(),
				final_pose
			))

	# ---------------------------------------------------------------- #
	def done(self, context, canceled, core: Core):
		# ---- Restore Player Speed ----
		core.player.walk_speed = self.base_walk_speed
		core.player.run_speed = self.base_run_speed

		# ---- Write cached keyframes ----
		if self.enable_animation and self.recorded_frames:
			self.character.rotation_mode = 'QUATERNION'

			# First frame: use keyframe_insert to let Blender create the action/slot/fcurve infrastructure
			first_frame, first_mat, first_pose = self.recorded_frames[0]
			self.character.matrix_world = first_mat
			apply_pose(self.character, first_pose)
			self.character.keyframe_insert('location', frame=first_frame)
			self.character.keyframe_insert('rotation_quaternion', frame=first_frame)
			for bone_name, (bloc, brot) in first_pose.items():
				pb = self.character.pose.bones.get(bone_name)
				if pb is None:
					continue
				if bloc is not None:
					pb.keyframe_insert('location', frame=first_frame)
				if brot is not None:
					pb.keyframe_insert('rotation_quaternion', frame=first_frame)

			# Grab the channelbag Blender just created, cache fcurves
			ad = self.character.animation_data
			cb = anim_utils.action_get_channelbag_for_slot(ad.action, ad.action_slot)

			obj_loc = [cb.fcurves.find('location', index=i) for i in range(3)]
			obj_rot = [cb.fcurves.find('rotation_quaternion', index=i) for i in range(4)]

			bone_fcs = {}
			for bone_name in first_pose:
				dp_l = f'pose.bones["{bone_name}"].location'
				dp_r = f'pose.bones["{bone_name}"].rotation_quaternion'
				bone_fcs[bone_name] = (
					[cb.fcurves.find(dp_l, index=i) for i in range(3)],
					[cb.fcurves.find(dp_r, index=i) for i in range(4)],
				)

			# Remaining frames: fast insert
			for frame, mat, pose in self.recorded_frames[1:]:
				loc = mat.to_translation()
				rot = mat.to_quaternion()
				for i in range(3):
					if obj_loc[i]:
						obj_loc[i].keyframe_points.insert(frame, loc[i], options={'FAST'})
				for i in range(4):
					if obj_rot[i]:
						obj_rot[i].keyframe_points.insert(frame, rot[i], options={'FAST'})
				for bone_name, (bloc, brot) in pose.items():
					cached = bone_fcs.get(bone_name)
					if not cached:
						continue
					lf, rf = cached
					if bloc is not None:
						for i in range(3):
							if lf[i]:
								lf[i].keyframe_points.insert(frame, bloc[i], options={'FAST'})
					if brot is not None:
						for i in range(4):
							if rf[i]:
								rf[i].keyframe_points.insert(frame, brot[i], options={'FAST'})

			for fc in cb.fcurves:
				fc.update()

		# ---- Cleanup ----
		if canceled and not self.owns_character:
			self.character.matrix_world = self.backup_matrix

		if self.owns_character and not self.enable_animation:
			util.obj.delete(self.character, include_children=True)


# ******************************************************************** #
def armature_poll(self, object):
	return object.type == 'ARMATURE'

class OMNISTEP_Module_CharacterAnimation_prop(bpy.types.PropertyGroup):

	character_armature: bpy.props.PointerProperty(name="Character", description="Target armature to animate. If empty, a copy of the idle armature is created",
		type=bpy.types.Object, poll=armature_poll)

	idle_armature: bpy.props.PointerProperty(name="Idle", description="Armature with idle animation",
		type=bpy.types.Object, poll=armature_poll)

	walk_armature: bpy.props.PointerProperty(name="Walk", description="Armature with walk animation",
		type=bpy.types.Object, poll=armature_poll)

	run_armature: bpy.props.PointerProperty(name="Run", description="Armature with run animation",
		type=bpy.types.Object, poll=armature_poll)

	strafe_left_armature: bpy.props.PointerProperty(name="Strafe Left", description="Armature with left strafe animation",
		type=bpy.types.Object, poll=armature_poll)

	strafe_right_armature: bpy.props.PointerProperty(name="Strafe Right", description="Armature with right strafe animation",
		type=bpy.types.Object, poll=armature_poll)

	idle_air_armature: bpy.props.PointerProperty(name="Idle Air", description="Armature with falling/airborne animation",
		type=bpy.types.Object, poll=armature_poll)

	walk_step_length: bpy.props.FloatProperty(name="Walk Step Length",
		description="Distance covered per full walk cycle",
		min=0.01, default=1.5, unit='LENGTH')

	run_step_length: bpy.props.FloatProperty(name="Run Step Length",
		description="Distance covered per full run cycle",
		min=0.01, default=3.0, unit='LENGTH')

	strafe_walk_step_length: bpy.props.FloatProperty(name="Strafe Walk Step Length",
		description="Distance covered per full strafe walk cycle",
		min=0.01, default=1.2, unit='LENGTH')

	strafe_run_step_length: bpy.props.FloatProperty(name="Strafe Run Step Length",
		description="Distance covered per full strafe run cycle",
		min=0.01, default=2.5, unit='LENGTH')

	forward_axis: bpy.props.EnumProperty(name="Forward Axis",
		description="Which direction the character faces in its rest pose",
		items=[
			('-Y', "-Y", "Character faces -Y (Mixamo default)"),
			('+Y', "+Y", "Character faces +Y"),
			('-X', "-X", "Character faces -X"),
			('+X', "+X", "Character faces +X"),
		], default='-Y')

	facing_smoothing: bpy.props.FloatProperty(name="Facing Smoothing",
		description="How quickly the character rotates to face the movement direction",
		min=0.01, max=1.0, default=0.15)

	acceleration_time: bpy.props.FloatProperty(name="Acceleration",
		description="Time in seconds to reach full speed from standstill",
		min=0.0, max=3.0, default=0.4, unit='TIME')

	deceleration_time: bpy.props.FloatProperty(name="Deceleration",
		description="Time in seconds to come to a stop after releasing input",
		min=0.0, max=3.0, default=0.3, unit='TIME')

	run_transition_time: bpy.props.FloatProperty(name="Run Transition",
		description="Time in seconds to blend between walk and run when pressing/releasing the speed key",
		min=0.0, max=2.0, default=0.3, unit='TIME')

	air_transition_time: bpy.props.FloatProperty(name="Air Transition",
		description="Time in seconds to blend between ground and air poses",
		min=0.0, max=2.0, default=0.15, unit='TIME')


# ================================================================ #
def draw_content(layout, context):
	settings = context.scene.omnistep_modules.character_animation
	layout.use_property_split = True
	column = layout.column()

	column.label(text="Character", icon='ARMATURE_DATA')
	column.prop(settings, "character_armature")

	layout.separator()

	column = layout.column()
	column.label(text="Input Animations", icon='ACTION')
	column.prop(settings, "idle_armature")
	column.prop(settings, "walk_armature")
	column.prop(settings, "run_armature")
	column.prop(settings, "strafe_left_armature")
	column.prop(settings, "strafe_right_armature")
	column.prop(settings, "idle_air_armature")

	layout.separator()

	column = layout.column()
	column.label(text="Forward", icon='CON_FOLLOWPATH')
	column.prop(settings, "walk_step_length")
	column.prop(settings, "run_step_length")

	has_strafe = settings.strafe_left_armature is not None and settings.strafe_right_armature is not None
	if has_strafe:
		layout.separator()
		column = layout.column()
		column.label(text="Strafe", icon='CON_FOLLOWPATH')
		column.prop(settings, "strafe_walk_step_length")
		column.prop(settings, "strafe_run_step_length")

	layout.separator()

	column = layout.column()
	column.label(text="Orientation", icon='ORIENTATION_GIMBAL')
	column.prop(settings, "forward_axis")
	column.prop(settings, "facing_smoothing", slider=True)

	layout.separator()

	column = layout.column()
	column.label(text="Movement Inertia", icon='FORCE_FORCE')
	column.prop(settings, "acceleration_time")
	column.prop(settings, "deceleration_time")
	column.prop(settings, "run_transition_time")
	column.prop(settings, "air_transition_time")