using System;

namespace CharacterEditor;

public static partial class AutomaticWardrobeFit
{
	static readonly Dictionary<(Model, bool), Model> headModels = new();
	static readonly Dictionary<(Model, bool), string> headFailures = new();
	static Surface identitySurface;
	static SerializedModelMorphs.Target[] identityMorphs;
	public static int CachedHeadModels => headModels.Count;
	public static bool IsHeadFitted(Model model) => model?.Name.StartsWith("automatic-head-fit:", StringComparison.Ordinal) == true;

	static void HeadReferences()
	{
		if (identitySurface is not null) return;
		var model = Model.Load(UnifiedHuman.ModelPath);
		var buffers = ReadBuffers(model);
		int index = Enumerable.Range(0, model.MeshCount).First(i => model.MeshInfo.Meshes[i].Name.StartsWith("head", StringComparison.OrdinalIgnoreCase));
		var buffer = buffers[index];
		// The first stream is facial skin; the separate lower-neck strip is not a hair attachment surface.
		var vertices = Vertices(buffer.Vertices[0], null, buffer.Bones.Length);
		var indices = Indices(buffer.Indices[0]);
		var call = model.MeshInfo.Meshes[index].DrawCalls.First();
		var triangles = indices.Skip(call.StartIndex).Take(call.Indices).Select(i => i + call.BaseVertex).ToArray();
		var morphs = SerializedModelMorphs.Read(model, index, vertices.Length, true);
		var points = vertices.Select(v => v.Position).ToArray();
		identitySurface = new(points, points, triangles, System.Array.Empty<string>(), System.Array.Empty<(int, float)[]>());
		identityMorphs = morphs;
	}

	/// <summary>Shared identity morphs, built once for each authored style. Slider values
	/// only update native morph weights; they never create a new fitted model.</summary>
	public static Model FitHead(Model source, bool facialHair)
	{
		if (source is null || source.IsError || IsHeadFitted(source) || source.BoneCount == 0) return source;
		if (cachedRevision != Revision) { ClearCache(); cachedRevision = Revision; }
		if (headModels.TryGetValue((source, facialHair), out var ready)) return ready;
		if (headFailures.TryGetValue((source, facialHair), out var failure)) throw new InvalidOperationException(failure);
		try
		{
			var result = BuildHead(source, facialHair);
			headModels[(source, facialHair)] = result;
			return result;
		}
		catch (Exception error)
		{
			var message = "Cannot fit hair " + source.Name + ": " + error.Message;
			headFailures[(source, facialHair)] = message;
			throw new InvalidOperationException(message, error);
		}
	}

	static Model BuildHead(Model source, bool facialHair)
	{
		HeadReferences();
		var buffers = ReadBuffers(source);
		var builder = Model.Builder.WithName("automatic-head-fit:" + source.Name + ":v" + Revision);
		for (int i = 0; i < source.BoneCount; i++)
		{
			var rest = source.GetBoneTransform(i); int parent = source.GetBoneParent(i);
			builder.AddBone(source.GetBoneName(i), rest.Position, rest.Rotation, parent < 0 ? null : source.GetBoneName(parent));
		}
		var originals = new List<Material>();
		var nativeVertices = source.GetVertices();
		for (int mi = 0; mi < buffers.Count; mi++)
		{
			var buffer = buffers[mi]; var info = source.MeshInfo.Meshes[mi];
			var boneMap = buffer.Bones.Select(name => Enumerable.Range(0, source.BoneCount).Single(i => source.GetBoneName(i) == name)).ToArray();
			var nativeMorphs = SerializedModelMorphs.Read(source, mi, buffer.Vertices.Sum(v => v.Count));
			int callIndex = 0;
			foreach (var call in info.DrawCalls)
			{
				int vb = buffer.Vertices.Length == 1 ? 0 : callIndex, ib = buffer.Indices.Length == 1 ? 0 : callIndex;
				int morphOffset = buffer.Vertices.Take(vb).Sum(v => v.Count);
				callIndex++;
				var vertices = Vertices(buffer.Vertices[vb], nativeVertices, buffer.Bones.Length);
				var indices = Indices(buffer.Indices[ib]).Skip(call.StartIndex).Take(call.Indices).Select(i => i + call.BaseVertex).ToArray();
				var used = indices.Distinct().ToArray();
				var lookup = used.Select((id, i) => (id, i)).ToDictionary(x => x.id, x => x.i);
				var localIndices = indices.Select(i => lookup[i]).ToArray();
				var v = used.Select(i => vertices[i]).ToArray();
				for (int i = 0; i < v.Length; i++)
				{
					var b = v[i].Indices;
					v[i].Indices = new((byte)boneMap[b.r], (byte)boneMap[b.g], (byte)boneMap[b.b], (byte)boneMap[b.a]);
				}
				var points = v.Select(x => x.Position).ToArray();
				var authoredNormals = Normals(points, localIndices);
				var authoredTangents = Tangents(v, points, authoredNormals, localIndices);
				for (int i = 0; i < v.Length; i++)
				{
					if (v[i].Normal.LengthSquared < .01f) v[i].Normal = authoredNormals[i];
					if (new Vector3(v[i].Tangent.x, v[i].Tangent.y, v[i].Tangent.z).LengthSquared < .01f) v[i].Tangent = authoredTangents[i];
				}
				var hits = points.Select(p => identitySurface.Nearest(new Vector3(p.x, p.y, MathF.Max(facialHair ? 60.1f : 64.5f, p.z)))).ToArray();
				var mesh = new Mesh(call.Material, (MeshPrimitiveType)call.PrimitiveType) { UvDensity = call.UvDensity };
				mesh.CreateVertexBuffer(v.Length, SkinLayout, v.ToList());
				mesh.CreateIndexBuffer(localIndices.Length, localIndices.ToList());
				foreach (var morph in nativeMorphs)
					mesh.AddMorph(morph.Name, used.Select((id, i) => new MorphDelta(i, morph.Positions[id + morphOffset], morph.Normals[id + morphOffset])).ToArray());
				// Native models can expose no-op controllers (for example tongueOut
				// on a moustache) without a corresponding atlas rectangle.
				if (nativeMorphs.Length > 0)
					foreach (var name in Enumerable.Range(0, source.MorphCount).Select(source.GetMorphName).Except(nativeMorphs.Select(m => m.Name)).Where(n => !string.IsNullOrWhiteSpace(n)))
						mesh.AddMorph(name, new[] { new MorphDelta(0, Vector3.Zero, Vector3.Zero) });
				var baseNormals = Normals(points, localIndices);
				foreach (var morph in identityMorphs)
				{
					var delta = hits.Select(hit =>
					{
						int t = hit.Triangle * 3; var w = hit.Barycentric;
						return morph.Positions[identitySurface.Triangles[t]] * w.x + morph.Positions[identitySurface.Triangles[t + 1]] * w.y + morph.Positions[identitySurface.Triangles[t + 2]] * w.z;
					}).ToArray();
					if (!delta.Any(d => d.LengthSquared > 1e-8f)) continue;
					var moved = points.Select((p, i) => p + delta[i]).ToArray();
					var normals = Normals(moved, localIndices);
					mesh.AddMorph(morph.Name, delta.Select((d, i) => new MorphDelta(i, d, normals[i] - baseNormals[i])).ToArray());
				}
				mesh.Bounds = new BBox(points.Aggregate((a, b) => Vector3.Min(a, b)), points.Aggregate((a, b) => Vector3.Max(a, b))).Grow(3);
				if (source.MeshInfo.LodCount == 0) builder.AddMesh(mesh);
				else for (int lod = 0; lod < source.MeshInfo.LodCount; lod++) if ((info.LodMask & (1 << lod)) != 0) builder.AddMesh(mesh, lod);
				originals.Add(call.Material);
			}
		}
		var baseMaterials = (source.MaterialGroupCount > 0 ? source.GetMaterials(0) : source.Materials).ToArray();
		for (int group = 0; group < source.MaterialGroupCount; group++)
		{
			var materials = source.GetMaterials(group).ToArray(); var output = builder.AddMaterialGroup(source.GetMaterialGroupName(group));
			foreach (var original in originals)
			{
				int index = System.Array.IndexOf(baseMaterials, original);
				output.AddMaterial(index >= 0 && index < materials.Length ? materials[index] : original);
			}
		}
		for (int lod = 1; lod < source.MeshInfo.LodSwitchDistances.Length; lod++) builder.WithLodDistance(lod, source.MeshInfo.LodSwitchDistances[lod]);
		return builder.Create();
	}
}