using Editor.Prism.Compiler;
using Editor.Prism.Core;
using Editor.Prism.Model;
using Editor.Prism.Serialization;
using EngineMaterial = Sandbox.Material;

namespace Editor.Prism.Preview;

/// <summary>
/// Renders one small image per node and hands it to the graph view.
/// <para>
/// <c>INode.Thumbnail</c> has been plumbed through <c>NodeUI.OnPaint</c> in s&amp;box for years and
/// nothing has ever filled it for a shader graph. This does, and it does it for the price of a single
/// compile: the preview build of the shader already carries every node's value behind the
/// <see cref="PrismConstants.StageIdAttribute"/> switch, so a thumbnail is one attribute write and one
/// tiny offscreen render — not a compile, not a render target per node, not a shader variant per node.
/// </para>
/// <para>
/// The rig is a bare <c>SceneWorld</c> with an orthographic <c>SceneCamera</c>, a unit quad wearing
/// the preview material, and enough lighting for a lit surface shader to read. Work is budgeted per
/// frame and prioritised — nodes the user can actually see first — and results are cached against a
/// per-node subtree hash, so a node whose inputs did not change is never rendered twice.
/// </para>
/// <para>
/// <c>SceneCamera.RenderToTexture</c> is internal to the engine and unreachable from an addon, so the
/// public <c>Sandbox.Tools</c> extension <c>RenderToPixmap</c> is used instead. It performs a GPU
/// readback, which is why the per-frame budget defaults low and why nothing here is ever done for a
/// node the user cannot see.
/// </para>
/// </summary>
public sealed class NodeThumbnailService : IDisposable
{
	/// <summary>Default thumbnail edge in pixels. Matches <c>PrismTheme.ThumbnailSize</c>.</summary>
	public const int DefaultSize = 96;

	/// <summary>Thumbnails rendered per editor frame before the rest wait for the next one.</summary>
	public const int DefaultBudget = 6;

	sealed class Entry
	{
		public Pixmap Pixmap;
		public ulong Hash;
		public int StageId;
		public bool Dirty = true;
		public bool Queued;
		public long Touched;
	}

	static readonly List<NodeThumbnailService> s_live = new();
	static readonly List<NodeThumbnailService> s_frameBuffer = new();

	readonly Dictionary<NodeId, Entry> _entries = new();
	readonly List<NodeId> _queue = new();

	SceneWorld _world;
	SceneCamera _camera;
	SceneObject _quad;
	EngineMaterial _material;
	IReadOnlyDictionary<NodeId, int> _stageIds = new Dictionary<NodeId, int>();

	/// <summary>Create a service. Nothing is allocated on the GPU until the first render.</summary>
	public NodeThumbnailService( int size = DefaultSize )
	{
		Size = Math.Clamp( size, 16, 512 );

		lock ( s_live )
		{
			s_live.Add( this );
		}
	}

	/// <summary>Edge length of every generated thumbnail, in pixels.</summary>
	public int Size { get; }

	/// <summary>
	/// Whether thumbnails are produced at all. Off by default: it is a visible, opt-in feature that
	/// costs GPU readbacks, and the window persists the toggle under
	/// <see cref="PrismConstants.CookieNodePreviews"/>.
	/// </summary>
	public bool Enabled { get; set; }

	/// <summary>How many thumbnails may be rendered per editor frame.</summary>
	public int Budget { get; set; } = DefaultBudget;

	/// <summary>
	/// The most thumbnails that may be cached at once.
	/// <para>
	/// This is a real bound, not a request threshold: at 96×96 RGBA a cached image is 36 KB, so an
	/// unbounded cache on a graph that is being edited for an hour is tens of megabytes of pixmaps for
	/// nodes that no longer exist. When the cache is full the least-recently-touched entry is evicted,
	/// and <see cref="Retain"/> drops everything the graph no longer contains.
	/// </para>
	/// </summary>
	public int NodeLimit
	{
		get => _nodeLimit;
		set
		{
			_nodeLimit = Math.Clamp( value, 1, 4096 );

			Trim();
		}
	}

	int _nodeLimit = 128;
	long _clock;

	/// <summary>Live uniforms pushed to the thumbnail quad. Mirror the viewport's bus into this.</summary>
	public PreviewAttributeBus Attributes { get; } = new();

	/// <summary>How many thumbnails are waiting to be rendered.</summary>
	public int PendingCount => _queue.Count;

	/// <summary>How many thumbnails are cached.</summary>
	public int Count => _entries.Count;

	/// <summary>True when there is nothing left to render.</summary>
	public bool IsIdle => _queue.Count == 0;

	/// <summary>True when a material and a stage map are both available, so thumbnails can differ per node.</summary>
	public bool CanRender => _material is not null && Enabled;

	/// <summary>Raised on the main thread once a node's thumbnail has been produced or replaced.</summary>
	public event Action<NodeId> ThumbnailReady;

	// ---- source ------------------------------------------------------------

	/// <summary>
	/// Point the service at a freshly compiled preview material. The per-node stage ids are recovered
	/// from the artifact's source map: every node that produced generated code gets the index of its
	/// first generated line, which is exactly the numbering the preview stage switch uses.
	/// </summary>
	public void SetSource( EngineMaterial material, CompileResult result )
	{
		SetSource( material, BuildStageMap( result ) );

		if ( result?.PreviewAttributes is { Count: > 0 } ) Attributes.Apply( result.PreviewAttributes );
		if ( result?.PreviewTextures is { Count: > 0 } ) Attributes.ApplyTextures( result.PreviewTextures );
	}

	/// <summary>Point the service at a material and an explicit node-to-stage-id map.</summary>
	public void SetSource( EngineMaterial material, IReadOnlyDictionary<NodeId, int> stageIds )
	{
		_material = material;
		_stageIds = stageIds ?? new Dictionary<NodeId, int>();

		if ( _quad.IsValid() && _material is not null )
		{
			PrismLog.Guard( "Applying the thumbnail material", () => _quad.SetMaterialOverride( _material ) );
		}

		// The shader changed, so every cached image is stale until its hash says otherwise.
		foreach ( var entry in _entries.Values ) entry.Dirty = true;
	}

	// ---- requests ----------------------------------------------------------

	/// <summary>The cached thumbnail for a node, or null when there is not one yet.</summary>
	public Pixmap Get( NodeId node ) => _entries.TryGetValue( node, out var entry ) ? entry.Pixmap : null;

	/// <summary>True when a node has a rendered, up-to-date thumbnail.</summary>
	public bool Has( NodeId node ) => _entries.TryGetValue( node, out var entry ) && entry.Pixmap is not null && !entry.Dirty;

	/// <summary>
	/// Ask for a node's thumbnail. <paramref name="hash"/> should be the node's dependency-subtree hash
	/// from <see cref="SubtreeHash"/>; when it matches what produced the cached image, nothing is
	/// re-rendered. Pass zero to force a render.
	/// </summary>
	public void Request( NodeId node, ulong hash = 0 )
	{
		if ( !node.IsValid ) return;

		var fresh = !_entries.TryGetValue( node, out var entry );

		if ( fresh )
		{
			entry = new Entry();
			_entries[node] = entry;
		}

		// Stamped before the trim, so the entry that was just asked for can never be the one evicted.
		entry.Touched = ++_clock;

		if ( fresh ) Trim();

		if ( hash != 0 && entry.Hash == hash && entry.Pixmap is not null )
		{
			// Same inputs, same shader path: the image cannot have changed.
			entry.Dirty = false;
			return;
		}

		entry.Hash = hash;
		entry.Dirty = true;

		Enqueue( node, entry, false );
	}

	/// <summary>
	/// Drop every cached thumbnail whose node is not in <paramref name="wanted"/>. This is what keeps a
	/// long editing session from accumulating images for deleted nodes and for nodes whose preview flag
	/// the user has since turned off.
	/// </summary>
	/// <returns>How many entries were released.</returns>
	public int Retain( IEnumerable<NodeId> wanted )
	{
		if ( wanted is null ) return 0;

		var keep = new HashSet<NodeId>();

		foreach ( var node in wanted ) keep.Add( node );

		var doomed = new List<NodeId>();

		foreach ( var node in _entries.Keys )
		{
			if ( !keep.Contains( node ) ) doomed.Add( node );
		}

		foreach ( var node in doomed ) Release( node );

		return doomed.Count;
	}

	/// <summary>Evict least-recently-requested entries until the cache is inside its bound.</summary>
	void Trim()
	{
		while ( _entries.Count > _nodeLimit )
		{
			var oldest = NodeId.None;
			var oldestTouch = long.MaxValue;

			foreach ( var (node, entry) in _entries )
			{
				if ( entry.Touched >= oldestTouch ) continue;

				oldest = node;
				oldestTouch = entry.Touched;
			}

			if ( !oldest.IsValid ) break;

			Release( oldest );
		}
	}

	void Release( NodeId node )
	{
		if ( !_entries.Remove( node, out var entry ) ) return;

		_queue.Remove( node );

		entry.Queued = false;
		entry.Pixmap = null;
	}

	/// <summary>Ask for a batch of thumbnails in one go.</summary>
	public void Request( IEnumerable<NodeId> nodes )
	{
		foreach ( var node in nodes ?? Array.Empty<NodeId>() ) Request( node );
	}

	/// <summary>
	/// Ask for a batch of thumbnails, hashing each node's dependency subtree so unchanged nodes are
	/// skipped. This is the call the graph view should make after every successful compile.
	/// </summary>
	public int RequestFor( IPrismGraph graph, IEnumerable<NodeId> nodes )
	{
		if ( graph is null ) return 0;

		var requested = 0;

		foreach ( var node in nodes ?? Array.Empty<NodeId>() )
		{
			Request( node, SubtreeHash( graph, node ) );
			requested++;
		}

		return requested;
	}

	/// <summary>
	/// Move nodes to the front of the render queue. The graph view calls this with whatever is inside
	/// the viewport, so what the user is looking at fills in first.
	/// </summary>
	public void Prioritise( IEnumerable<NodeId> nodes )
	{
		if ( nodes is null ) return;

		var insert = 0;

		foreach ( var node in nodes )
		{
			var index = _queue.IndexOf( node );

			if ( index < 0 )
			{
				if ( !_entries.TryGetValue( node, out var pending ) || !pending.Dirty ) continue;

				_queue.Insert( Math.Min( insert++, _queue.Count ), node );
				pending.Queued = true;

				continue;
			}

			if ( index == insert )
			{
				insert++;
				continue;
			}

			_queue.RemoveAt( index );
			_queue.Insert( Math.Min( insert++, _queue.Count ), node );
		}
	}

	/// <summary>Mark one node's thumbnail stale.</summary>
	public void Invalidate( NodeId node )
	{
		if ( !_entries.TryGetValue( node, out var entry ) ) return;

		entry.Dirty = true;
		entry.Hash = 0;

		Enqueue( node, entry, false );
	}

	/// <summary>Mark every thumbnail stale without discarding the images, so nothing flickers to empty.</summary>
	public void InvalidateAll()
	{
		foreach ( var (node, entry) in _entries )
		{
			entry.Dirty = true;
			entry.Hash = 0;

			Enqueue( node, entry, false );
		}
	}

	/// <summary>Discard every cached image and the queue.</summary>
	public void Clear()
	{
		_entries.Clear();
		_queue.Clear();
	}

	// ---- rendering ---------------------------------------------------------

	/// <summary>
	/// Render up to <see cref="Budget"/> queued thumbnails. Driven automatically from the editor frame
	/// hook; calling it by hand simply does the next batch early.
	/// </summary>
	public void Tick()
	{
		if ( !Enabled || _material is null ) return;
		if ( _queue.Count == 0 ) return;
		if ( !ThreadSafe.IsMainThread ) return;

		// Rendering an offscreen view from inside a render block is not legal.
		if ( PrismLog.Guard( "Checking the render state", () => Graphics.IsActive ) ) return;
		if ( !EnsureRig() ) return;

		var budget = Math.Max( 1, Budget );

		while ( budget-- > 0 && _queue.Count > 0 )
		{
			var node = _queue[0];

			_queue.RemoveAt( 0 );

			if ( !_entries.TryGetValue( node, out var entry ) ) continue;

			entry.Queued = false;

			if ( !entry.Dirty ) continue;

			if ( !Render( node, entry ) ) break;

			entry.Dirty = false;

			ThumbnailReady?.Invoke( node );
		}
	}

	/// <summary>Tear the offscreen rig down. Cached pixmaps are released with it.</summary>
	public void Dispose()
	{
		lock ( s_live )
		{
			s_live.Remove( this );
		}

		ThumbnailReady = null;
		Attributes.Target = null;

		PrismLog.Guard( "Destroying the thumbnail rig", () =>
		{
			_quad?.Delete();
			_camera?.Dispose();
			_world?.Delete();
		} );

		_quad = null;
		_camera = null;
		_world = null;
		_material = null;

		Clear();
	}

	bool Render( NodeId node, Entry entry )
	{
		entry.StageId = _stageIds is not null && _stageIds.TryGetValue( node, out var stage ) ? stage : 0;

		entry.Pixmap ??= PrismLog.Guard( "Allocating a thumbnail", () => new Pixmap( Size, Size ), null );

		if ( entry.Pixmap is null ) return false;

		return PrismLog.Guard( $"Rendering the thumbnail for {node}", () =>
		{
			var attributes = _quad.Attributes;

			attributes.Set( PrismConstants.StageIdAttribute, entry.StageId );
			attributes.Set( PrismConstants.ChannelAttribute, 0 );
			attributes.Set( PreviewViewport.TimeAttribute, RealTime.Now );

			return _camera.RenderToPixmap( entry.Pixmap );
		} );
	}

	bool EnsureRig()
	{
		if ( _world is not null && _camera is not null && _quad.IsValid() ) return true;

		return PrismLog.Guard( "Building the thumbnail rig", () =>
		{
			_world ??= new SceneWorld();

			if ( _camera is null )
			{
				var extent = PreviewMeshes.Radius;

				_camera = new SceneCamera( "Prism Thumbnails" )
				{
					World = _world,
					Ortho = true,
					OrthoHeight = extent * 2f,
					ZNear = 1f,
					ZFar = 4000f,
					Position = Vector3.Backward * 500f,
					Rotation = Rotation.LookAt( Vector3.Forward ),
					BackgroundColor = Color.Transparent,
					AmbientLightColor = Color.White * 0.05f,
					AntiAliasing = true,
					EnablePostProcessing = false
				};

				// Enough light for a lit surface shader to resolve to something other than black, and a
				// cubemap so metals and reflections read.
				new ScenePointLight( _world, new Vector3( -260f, 180f, 180f ), 1400f, Color.White * 6f )
					.ShadowsEnabled = false;

				new ScenePointLight( _world, new Vector3( -260f, -180f, -140f ), 1400f, new Color( 0.55f, 0.7f, 1f ) * 3f )
					.ShadowsEnabled = false;

				var cubemap = PrismLog.Guard( "Loading the thumbnail envmap",
					() => Texture.Load( PreviewScene.DefaultEnvmap ), null );

				if ( cubemap is not null )
				{
					_ = new SceneCubemap( _world, cubemap, BBox.FromPositionAndSize( Vector3.Zero, 20000f ) );
				}
			}

			if ( !_quad.IsValid() )
			{
				_quad = new SceneObject( _world, PreviewMeshes.Quad, Transform.Zero )
				{
					Batchable = false
				};

				_quad.Flags.CastShadows = false;

				Attributes.Target = _quad;
			}

			if ( _material is not null ) _quad.SetMaterialOverride( _material );

			return _world is not null && _camera is not null && _quad.IsValid();
		} );
	}

	void Enqueue( NodeId node, Entry entry, bool front )
	{
		if ( entry.Queued ) return;

		entry.Queued = true;

		if ( front ) _queue.Insert( 0, node );
		else _queue.Add( node );
	}

	// ---- static helpers ----------------------------------------------------

	/// <summary>
	/// Recover a node-to-stage-id map from a compile result.
	/// <para>
	/// The source map records which generated line each node produced, so ordering nodes by their
	/// first generated line reproduces exactly the order the preview shader's stage switch numbers
	/// them in. Ids are one-based, because zero means "show the shader's own output".
	/// </para>
	/// </summary>
	/// <remarks>
	/// Delegated to <see cref="PreviewInstrumentation.StageMap"/> rather than reimplemented. The backend
	/// bakes these numbers into the generated switch, so two copies of the numbering rule is two copies
	/// of a contract that must agree exactly — and a drift between them would silently show the wrong
	/// node's value in the preview rather than fail.
	/// </remarks>
	public static IReadOnlyDictionary<NodeId, int> BuildStageMap( CompileResult result ) =>
		PreviewInstrumentation.StageMap( result );

	/// <summary>
	/// A stable hash of everything that can change what a node evaluates to: its own structure and
	/// values, and those of every node it depends on, plus the shape of the connections between them.
	/// Two nodes with the same hash produce the same thumbnail, which is what makes caching correct
	/// rather than merely fast.
	/// </summary>
	public static ulong SubtreeHash( IPrismGraph graph, NodeId node )
	{
		if ( graph is null || !node.IsValid ) return 0;

		return PrismLog.Guard( $"Hashing the subtree of {node}", () =>
		{
			var hash = Fnv.Offset;
			var subtree = GraphQueries.DependencySubtree( graph, node );

			foreach ( var id in subtree )
			{
				var member = graph.FindNode( id );

				if ( member is null ) continue;

				hash = Fnv.Mix( hash, member.Descriptor?.Id );
				hash = Fnv.Mix( hash, NodeProperties.StructuralHash( member ) );
				hash = Fnv.Mix( hash, NodeProperties.ValueHash( member ) );

				foreach ( var input in member.Inputs )
				{
					if ( !graph.TryGetIncomingEdge( id, input.Id, out var edge ) ) continue;

					hash = Fnv.Mix( hash, input.Id.Value );
					hash = Fnv.Mix( hash, edge.FromNode.Value );
					hash = Fnv.Mix( hash, edge.FromPort.Value );
				}
			}

			// Zero is the "no hash" sentinel, so never return it by accident.
			return hash == 0 ? 1UL : hash;
		} );
	}

	/// <summary>
	/// Drives every live service once per editor frame, whether or not a viewport is visible.
	/// <para>
	/// The snapshot is taken into a reused buffer rather than a fresh array: this runs every frame for
	/// the life of the editor, and one allocation per frame for a list that is almost always one element
	/// long is exactly the kind of garbage a tool this size should not be producing at idle.
	/// </para>
	/// </summary>
	[EditorEvent.Frame]
	static void OnEditorFrame()
	{
		lock ( s_live )
		{
			if ( s_live.Count == 0 ) return;

			s_frameBuffer.Clear();
			s_frameBuffer.AddRange( s_live );
		}

		for ( int i = 0; i < s_frameBuffer.Count; i++ )
		{
			s_frameBuffer[i].Tick();
		}

		s_frameBuffer.Clear();
	}
}