Editor/AudioTimelineView.cs
using Editor;
using Sandbox;
using System;
using System.Collections.Generic;
using System.IO;

namespace Ardi;

public class AudioTimelineView : GraphicsView
{
	private readonly AudioTools Tool;
	private readonly AudioTimeAxis TimeAxis;
	private readonly AudioScrubber Scrubber;
	private readonly AudioWaveForm WaveForm;
	private readonly List<int> splitPoints = new List<int>();
	private readonly List<int> loopPoints = new List<int>();
	private int? draggingPoint = null;
	private bool isSplitMode = true;
	private string sourceFilePath;

	public float ZoomLevel { get; set; } = 1.0f;
	public float Duration { get; private set; }
	public float Time { get; set; }
	public bool Scrubbing { get; set; }
	public string Sound { get; private set; }
	public SoundHandle SoundHandle { get; private set; }
	public Rect VisibleRect { get; private set; }

	public AudioTimelineView( AudioTools parent ) : base( parent )
	{
		Tool = parent;
		SceneRect = new( 0, Size );
		HorizontalScrollbar = ScrollbarMode.On;
		VerticalScrollbar = ScrollbarMode.Off;
		Scale = 1;
		Time = 0;

		WaveForm = new AudioWaveForm( this );
		Add( WaveForm );

		TimeAxis = new AudioTimeAxis( this );
		Add( TimeAxis );

		Scrubber = new AudioScrubber( this );
		Add( Scrubber );

		DoLayout();
	}

	protected override void OnPaint()
	{
		base.OnPaint();

		if ( Duration != 0 ) return;

		Paint.Antialiasing = true;
		Paint.ClearPen();
		Paint.SetBrush( Theme.ControlBackground );
		Paint.DrawRect( LocalRect );

		Paint.SetPen( Theme.Text.WithAlpha( 0.6f ) );
		Paint.SetDefaultFont( 16f );
		var message = "Drop an audio file here (.sound, .wav, .mp3)";
		var textRect = LocalRect;
		Paint.DrawText( textRect, message, TextFlag.Center );
	}

	protected override void DoLayout()
	{
		base.DoLayout();

		var size = Size;
		size.x = MathF.Max( Size.x, PositionFromTime( Duration ) );
		SceneRect = new( 0, size );
		TimeAxis.Size = new Vector2( size.x, Theme.RowHeight );
		Scrubber.Size = new Vector2( 9, size.y );

		var r = SceneRect;
		r.Position = Vector2.Zero.WithY( Theme.RowHeight );
		r.Size = new Vector2( PositionFromTime( Duration ), Size.y - Theme.RowHeight );
		WaveForm.SceneRect = r;
		WaveForm.Analyse();

		Scrubber.Position = Scrubber.Position.WithX( PositionFromTime( Time ) - 3 ).SnapToGrid( 1.0f );
	}

	protected override void OnResize()
	{
		base.OnResize();
		DoLayout();
	}

	public override void OnDestroyed()
	{
		base.OnDestroyed();
		SoundHandle?.Stop( 0.0f );
		SoundHandle = null;
	}

	public override void OnDragDrop( DragEvent e )
	{
		Tool.OnDragDrop( e );
	}

	public override void OnDragHover( DragEvent e )
	{
		Tool.OnDragHover( e );
	}

	protected override void OnMousePress( MouseEvent e )
	{
		if ( Duration == 0 ) return;

		if ( e.LocalPosition.y <= Theme.RowHeight )
		{
			base.OnMousePress( e );
			return;
		}

		var scenePos = ToScene( e.LocalPosition );
		var time = TimeFromPosition( scenePos.x );
		var samplePos = (int)(time / Duration * WaveForm.SampleCount);

		if ( e.Button == MouseButtons.Left )
		{
			HandleLeftClick( samplePos );
		}
		else if ( e.Button == MouseButtons.Right )
		{
			HandleRightClick( samplePos );
		}
		else if ( e.Button == MouseButtons.Middle )
		{
			MoveScrubber( PositionFromTime( time ) );
		}

		base.OnMousePress( e );
	}

	private void HandleLeftClick( int samplePos )
	{
		if ( isSplitMode )
		{
			var nearPoint = FindNearestSplitPoint( samplePos );
			if ( nearPoint >= 0 )
			{
				draggingPoint = nearPoint;
			}
			else
			{
				splitPoints.Add( samplePos );
				splitPoints.Sort();
				WaveForm.Update();
			}
		}
		else
		{
			var nearPoint = FindNearestLoopPoint( samplePos );
			if ( nearPoint >= 0 )
			{
				draggingPoint = nearPoint;
			}
			else
			{
				if ( loopPoints.Count >= 2 )
				{
					loopPoints.Clear();
				}
				loopPoints.Add( samplePos );
				loopPoints.Sort();
				WaveForm.Update();
			}
		}
	}

	private void HandleRightClick( int samplePos )
	{
		if ( isSplitMode )
		{
			var nearPoint = FindNearestSplitPoint( samplePos );
			if ( nearPoint >= 0 && nearPoint != 0 && nearPoint != splitPoints.Count - 1 )
			{
				splitPoints.RemoveAt( nearPoint );
				WaveForm.Update();
			}
		}
		else
		{
			var nearPoint = FindNearestLoopPoint( samplePos );
			if ( nearPoint >= 0 )
			{
				loopPoints.RemoveAt( nearPoint );
				WaveForm.Update();
			}
		}
	}

	protected override void OnMouseMove( MouseEvent e )
	{
		if ( !draggingPoint.HasValue || Duration <= 0 )
		{
			base.OnMouseMove( e );
			return;
		}

		var scenePos = ToScene( e.LocalPosition );
		var time = TimeFromPosition( scenePos.x );
		var samplePos = (int)(time / Duration * WaveForm.SampleCount);

		if ( isSplitMode && draggingPoint.Value < splitPoints.Count )
		{
			splitPoints[draggingPoint.Value] = samplePos;
			splitPoints.Sort();
			WaveForm.Update();
		}
		else if ( !isSplitMode && draggingPoint.Value < loopPoints.Count )
		{
			loopPoints[draggingPoint.Value] = samplePos;
			loopPoints.Sort();
			WaveForm.Update();
		}

		base.OnMouseMove( e );
	}

	protected override void OnMouseReleased( MouseEvent e )
	{
		draggingPoint = null;
		base.OnMouseReleased( e );
	}

	private int FindNearestSplitPoint( int samplePos )
	{
		var tolerance = (int)(WaveForm.SampleCount / (Width * ZoomLevel) * 10);

		for ( int i = 0; i < splitPoints.Count; i++ )
		{
			if ( Math.Abs( splitPoints[i] - samplePos ) < tolerance )
				return i;
		}

		return -1;
	}

	private int FindNearestLoopPoint( int samplePos )
	{
		var tolerance = (int)(WaveForm.SampleCount / (Width * ZoomLevel) * 10);

		for ( int i = 0; i < loopPoints.Count; i++ )
		{
			if ( Math.Abs( loopPoints[i] - samplePos ) < tolerance )
				return i;
		}

		return -1;
	}

	public void OnFrame()
	{
		Time = Time.Clamp( 0, Duration );

		if ( !Tool.Playing )
		{
			SoundHandle?.Stop( 0.0f );
			SoundHandle = null;
		}

		VisibleRect = Rect.FromPoints( ToScene( LocalRect.TopLeft ), ToScene( LocalRect.BottomRight ) );

		if ( Tool.Playing && !Scrubbing )
		{
			Time += RealTime.Delta;
			var time = Time % Duration;
			if ( time < Time )
			{
				if ( Tool.Repeating )
				{
					Time = time;
					if ( SoundHandle.IsValid() )
					{
						SoundHandle.Time = Time;
					}
				}
				else
				{
					time = 0;
					Time = time;
					SoundHandle?.Stop( 0.0f );
					Tool.Playing = false;
				}
			}

			if ( Tool.Playing && !SoundHandle.IsValid() )
			{
				SoundHandle = EditorUtility.PlaySound( Sound, Time );
				SoundHandle.Occlusion = false;
				SoundHandle.DistanceAttenuation = false;
			}

			Scrubber.Position = Scrubber.Position.WithX( PositionFromTime( Time ) - 3 ).SnapToGrid( 1.0f );
		}

		if ( SoundHandle.IsValid() )
		{
			SoundHandle.Paused = Scrubbing;
		}

		TimeAxis.Update();
		WaveForm.Update();

		if ( Scrubbing || Tool.Playing )
		{
			CenterOn( Scrubber.Position );
		}
	}

	public float PositionFromTime( float time )
	{
		return (time / Duration).Clamp( 0, 1 ) * (Width * ZoomLevel);
	}

	public float TimeFromPosition( float position )
	{
		return (position / (Width * ZoomLevel)).Clamp( 0, 1 ) * Duration;
	}

	public void SetSourceFilePath( string filePath )
	{
		sourceFilePath = filePath;
	}

	public void SetSamples( short[] samples, float duration, string sound, int channels )
	{
		Sound = sound;
		Duration = duration;
		WaveForm.SetSamples( samples, duration );

		Time = 0;
		DoLayout();
		CenterOn( new Vector2( 0, 0 ) );
	}

	public void MoveScrubber( float position, bool centreOn = true )
	{
		Scrubber.Position = Vector2.Right * (position - 4).SnapToGrid( 1.0f );
		Time = TimeFromPosition( Scrubber.Position.x + 4 );

		if ( SoundHandle.IsValid() )
		{
			SoundHandle.Time = Time;
		}

		if ( centreOn )
		{
			CenterOn( Scrubber.Position );
			VisibleRect = Rect.FromPoints( ToScene( LocalRect.TopLeft ), ToScene( LocalRect.BottomRight ) );
		}

		WaveForm.Update();
	}

	protected override void OnWheel( WheelEvent e )
	{
		var oldZoom = ZoomLevel;
		ZoomLevel *= e.Delta > 0 ? 1.1f : 0.90f;
		ZoomLevel = ZoomLevel.Clamp( 1.0f, 20.0f );

		var mouseScenePos = ToScene( e.Position );

		DoLayout();

		var newMouseScenePos = ToScene( e.Position );

		var sceneDelta = mouseScenePos - newMouseScenePos;
		var currentCenter = ToScene( LocalRect.Center );
		var newCenter = currentCenter + sceneDelta;

		var sceneWidth = PositionFromTime( Duration );
		var viewWidth = LocalRect.Width;

		if ( sceneWidth > viewWidth )
		{
			var minCenterX = viewWidth / 2;
			var maxCenterX = sceneWidth - viewWidth / 2;
			newCenter.x = newCenter.x.Clamp( minCenterX, maxCenterX );
		}
		else
		{
			newCenter.x = sceneWidth / 2;
		}

		CenterOn( newCenter );

		TimeAxis.Update();
		WaveForm.Update();

		e.Accept();
	}

	public void SetMode( bool splitMode )
	{
		isSplitMode = splitMode;
		WaveForm.Update();
	}

	public void ResetPoints()
	{
		if ( isSplitMode )
		{
			splitPoints.Clear();
			splitPoints.Add( 0 );
			if ( WaveForm.Samples != null && WaveForm.Samples.Length > 0 )
			{
				splitPoints.Add( WaveForm.Samples.Length - 1 );
			}
		}
		else
		{
			loopPoints.Clear();
		}

		WaveForm.Update();
	}

	public void DetectLoopPoints()
	{
		if ( WaveForm.Samples == null || WaveForm.Samples.Length == 0 ) return;

		loopPoints.Clear();

		var samples = WaveForm.Samples;
		var sampleRate = WaveForm.SampleRate;

		float bpm = 120.0f;
		float beatsPerSecond = bpm / 60.0f;
		float samplesPerBeat = sampleRate / beatsPerSecond;

		int[] barLengths = { 1, 2, 4, 8, 16 };
		float bestScore = float.MaxValue;
		int bestStart = 0;
		int bestEnd = samples.Length - 1;

		foreach ( int bars in barLengths )
		{
			int samplesPerBar = (int)(samplesPerBeat * 4 * bars);

			if ( samplesPerBar > samples.Length * 0.75f ) continue;

			int searchWindow = (int)(sampleRate * 0.1f);
			int startSearch = Math.Max( 0, samples.Length / 10 - searchWindow );
			int endSearch = Math.Min( samples.Length / 10 + searchWindow, samples.Length / 4 );

			for ( int start = startSearch; start < endSearch; start += 100 )
			{
				int end = start + samplesPerBar;
				if ( end >= samples.Length ) continue;

				float score = 0;
				int compareWindow = Math.Min( 1024, samplesPerBar / 10 );

				for ( int i = 0; i < compareWindow; i++ )
				{
					float diff = Math.Abs( samples[start + i] - samples[end - compareWindow + i] );
					score += diff * diff;
				}

				score = (float)Math.Sqrt( score / compareWindow );

				if ( score < bestScore )
				{
					bestScore = score;
					bestStart = start;
					bestEnd = end;
				}
			}
		}

		loopPoints.Add( bestStart );
		loopPoints.Add( bestEnd );

		float loopDuration = (bestEnd - bestStart) / (float)sampleRate;

		WaveForm.Update();
	}

	public void ProcessAndSave( string baseFileName, string outputDir, short[] samples, SoundFile soundFile )
	{
		string actualOutputDir = outputDir;
		if ( !string.IsNullOrEmpty( sourceFilePath ) )
		{
			if ( !Path.IsPathRooted( sourceFilePath ) )
			{
				var fullSourcePath = Path.Combine( Project.Current.GetAssetsPath(), sourceFilePath );
				actualOutputDir = Path.GetDirectoryName( fullSourcePath );
			}
			else
			{
				actualOutputDir = Path.GetDirectoryName( sourceFilePath );
			}
		}

		if ( !Directory.Exists( actualOutputDir ) )
		{
			Directory.CreateDirectory( actualOutputDir );
		}

		if ( isSplitMode )
		{
			if ( splitPoints.Count < 2 ) return;

			for ( int i = 0; i < splitPoints.Count - 1; i++ )
			{
				var startIndex = splitPoints[i];
				var endIndex = splitPoints[i + 1];

				startIndex = Math.Max( 0, Math.Min( startIndex, samples.Length ) );
				endIndex = Math.Max( startIndex, Math.Min( endIndex, samples.Length ) );

				var length = endIndex - startIndex;
				var segmentSamples = new short[length * 2];

				for ( int j = 0; j < length; j++ )
				{
					segmentSamples[j * 2] = samples[j + startIndex];
					segmentSamples[j * 2 + 1] = samples[j + startIndex];
				}

				var wavPath = Path.Combine( actualOutputDir, $"{baseFileName}_part_{i + 1}.wav" );
				SaveWavFile( wavPath, segmentSamples, soundFile.Rate, 2 );
			}
		}
		else
		{
			if ( loopPoints.Count != 2 ) return;

			var startIndex = loopPoints[0];
			var endIndex = loopPoints[1];

			startIndex = Math.Max( 0, Math.Min( startIndex, samples.Length ) );
			endIndex = Math.Max( startIndex, Math.Min( endIndex, samples.Length ) );

			var length = endIndex - startIndex;
			var loopSamples = new short[length * 2];

			for ( int j = 0; j < length; j++ )
			{
				loopSamples[j * 2] = samples[j + startIndex];
				loopSamples[j * 2 + 1] = samples[j + startIndex];
			}

			var loopPath = Path.Combine( actualOutputDir, $"{baseFileName}_loop.wav" );
			SaveWavFile( loopPath, loopSamples, soundFile.Rate, 2 );
		}
	}

	private void SaveWavFile( string path, short[] samples, int sampleRate, int channels )
	{
		using ( var writer = new BinaryWriter( File.Create( path ) ) )
		{
			int bitsPerSample = 16;
			int bytesPerSample = bitsPerSample / 8;
			int dataSize = samples.Length * bytesPerSample;
			int fileSize = 36 + dataSize;

			writer.Write( System.Text.Encoding.ASCII.GetBytes( "RIFF" ) );
			writer.Write( fileSize );
			writer.Write( System.Text.Encoding.ASCII.GetBytes( "WAVE" ) );

			writer.Write( System.Text.Encoding.ASCII.GetBytes( "fmt " ) );
			writer.Write( 16 );
			writer.Write( (short)1 );
			writer.Write( (short)channels );
			writer.Write( sampleRate );
			writer.Write( sampleRate * channels * bytesPerSample );
			writer.Write( (short)(channels * bytesPerSample) );
			writer.Write( (short)bitsPerSample );

			writer.Write( System.Text.Encoding.ASCII.GetBytes( "data" ) );
			writer.Write( dataSize );

			foreach ( var sample in samples )
			{
				writer.Write( sample );
			}
		}
	}

	public List<int> SplitPoints => splitPoints;
	public List<int> LoopPoints => loopPoints;
	public bool IsSplitMode => isSplitMode;
}