using System;
using System.Threading.Tasks;

namespace CharacterEditor;
/// <summary>Cached enamel tints on the original skin shader; no per-frame compositor or added draw.</summary>
public sealed class TeethColorPresentation : IDisposable
{
	sealed class Tint
	{
		public Texture Texture;
		public Material Material;
		public int Users;
		public long LastUsed;
	}

	static Material mouth;
	static ComputeShader compositor;
	static readonly Dictionary<int, Tint> tints = new();
	static long clock;
	const int IdleCacheLimit = 16;
	SkinnedModelRenderer target;
	Model model;
	string group;
	Tint applied;
	int appliedKey = -1;
	readonly List<int> slots = new();
	public static async Task LoadAsync(CharacterAppearanceData appearance)
	{
		if (appearance?.Teeth?.CustomColor != true || mouth is not null && compositor is not null)
			return;
		mouth = await Material.LoadAsync("models/citizen_human/heads/citizen_human_malehead01/mouth.vmat");
		if (mouth is null || mouth.IsError)
			throw new InvalidOperationException("The native mouth material is unavailable.");
		compositor = new ComputeShader("shaders/characterkit/teeth_tint.shader");
		var source = mouth.GetTexture("g_tColor");
		source.MarkUsed(1024);
		for (int i = 0; !source.IsLoaded && i < 200; i++)
			await Task.Delay(10);
		if (!source.IsLoaded)
			throw new InvalidOperationException("The native mouth color map has not finished loading.");
	}

	static int Key(Color c) => ((int)MathF.Round(c.r * 255) << 16) | ((int)MathF.Round(c.g * 255) << 8) | (int)MathF.Round(c.b * 255);
	static Tint Acquire(int key)
	{
		if (!tints.TryGetValue(key, out var tint))
		{
			// Half-resolution oral albedo is 512 KB per distinct custom color.
			// The stock high-resolution normal/roughness maps remain untouched.
			var source = mouth.GetTexture("g_tColor");
			var texture = Texture.Create(512, 256).WithFormat(ImageFormat.RGBA8888).WithUAVBinding().WithName($"Shared tooth tint {key:X6}").Finish();
			var attributes = new RenderAttributes();
			attributes.Set("BaseColor", source);
			attributes.Set("Result", texture);
			attributes.Set("ToothColor", new Vector3(((key >> 16) & 255) / 255f, ((key >> 8) & 255) / 255f, (key & 255) / 255f));
			compositor.DispatchWithAttributes(attributes, 512, 256, 1);
			var material = mouth.CreateCopy($"characterkit-teeth-{key:X6}.vmat");
			material.Set("g_tColor", texture);
			tints[key] = tint = new()
			{
				Texture = texture,
				Material = material
			};
		}

		tint.Users++;
		tint.LastUsed = ++clock;
		return tint;
	}

	static void Release(Tint tint)
	{
		if (tint is null)
			return;
		tint.Users--;
		tint.LastUsed = ++clock;
		while (tints.Values.Count(x => x.Users == 0) > IdleCacheLimit)
		{
			var oldest = tints.Where(x => x.Value.Users == 0).MinBy(x => x.Value.LastUsed);
			tints.Remove(oldest.Key);
			oldest.Value.Texture.Dispose();
		}
	}

	public void Apply(SkinnedModelRenderer body, CharacterAppearanceData appearance)
	{
		if (!body.IsValid() || body.Model is null || appearance?.Teeth is null)
			return;
		if (target != body || model != body.Model || group != body.MaterialGroup)
		{
			Dispose();
			target = body;
			model = body.Model;
			group = body.MaterialGroup;
			for (int i = 0; i < body.Materials.Count; i++)
				if (body.Materials.GetOriginal(i)?.Name.EndsWith("/mouth.vmat", StringComparison.OrdinalIgnoreCase) == true)
					slots.Add(i);
		}

		if (slots.Count == 0)
			return;
		int key = appearance.Teeth.CustomColor ? Key(appearance.Teeth.Color) : -1;
		if (key == appliedKey && (key < 0 || slots.All(i => body.Materials.GetOverride(i) == applied.Material)))
			return;
		if (key >= 0 && (mouth is null || compositor is null))
			return;
		var next = key < 0 ? null : Acquire(key);
		foreach (var slot in slots)
			body.Materials.SetOverride(slot, next?.Material);
		Release(applied);
		applied = next;
		appliedKey = key;
	}

	public void Dispose()
	{
		if (target.IsValid() && target.Model == model && applied is not null)
			foreach (var slot in slots)
				if (target.Materials.GetOverride(slot) == applied.Material)
					target.Materials.SetOverride(slot, null);
		Release(applied);
		applied = null;
		appliedKey = -1;
		slots.Clear();
		target = null;
		model = null;
		group = null;
	}
}