namespace LlmPoc.Llm;

public sealed class Tensor
{
	public string Name { get; }
	public int[] Shape { get; }
	public float[] Data { get; }
	public int Rank => Shape.Length;
	public long ElementCount => Data.LongLength;

	public Tensor( string name, int[] shape, float[] data )
	{
		Name = name ?? throw new ArgumentNullException( nameof( name ) );
		Shape = shape ?? throw new ArgumentNullException( nameof( shape ) );
		Data = data ?? throw new ArgumentNullException( nameof( data ) );

		long expected = 1;
		for ( int axis = 0; axis < Shape.Length; axis++ )
		{
			if ( Shape[axis] <= 0 )
			{
				throw new ArgumentException(
					$"[LLM:ERROR] {Name} dimension {axis} must be positive, found {Shape[axis]}.",
					nameof( shape ) );
			}

			if ( expected > long.MaxValue / Shape[axis] )
			{
				throw new ArgumentException(
					$"[LLM:ERROR] {Name} shape [{string.Join( ",", Shape )}] overflows " +
					"the signed 64-bit element count.", nameof( shape ) );
			}
			expected *= Shape[axis];
		}

		if ( expected != Data.LongLength )
		{
			throw new ArgumentException(
				$"[LLM:ERROR] {Name} shape {ShapeText} requires {expected} values, " +
				$"but received {Data.LongLength}.", nameof( data ) );
		}
	}

	public string ShapeText => $"[{string.Join( ",", Shape )}]";

	public void RequireShape( params int[] expected )
	{
		if ( expected is null )
		{
			throw new ArgumentNullException( nameof( expected ) );
		}

		if ( expected.Length != Shape.Length )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] {Name} expected rank {expected.Length} and shape " +
				$"[{string.Join( ",", expected )}], found rank {Shape.Length} and shape {ShapeText}." );
		}

		for ( int axis = 0; axis < Shape.Length; axis++ )
		{
			if ( Shape[axis] != expected[axis] )
			{
				throw new InvalidOperationException(
					$"[LLM:ERROR] {Name} expected shape [{string.Join( ",", expected )}], " +
					$"found {ShapeText}; mismatch at axis {axis}." );
			}
		}
	}

	public float Get1D( int index )
	{
		if ( Rank != 1 )
		{
			throw new InvalidOperationException(
				$"[LLM:ERROR] {Name} Get1D requires rank 1, found shape {ShapeText}." );
		}

		if ( index < 0 || index >= Shape[0] )
		{
			throw new IndexOutOfRangeException(
				$"[LLM:ERROR] {Name} index {index} is outside [0,{Shape[0]})." );
		}

		return Data[index];
	}
}