namespace LlmPoc.Llm;

public sealed class SboxLlmModel
{
	private readonly Dictionary<string, Tensor> _tensors;

	public int FormatVersion { get; }
	public IReadOnlyDictionary<string, Tensor> Tensors => _tensors;
	public int TensorCount => _tensors.Count;
	public long TotalParameterCount { get; }
	public long TotalTensorBytes => TotalParameterCount * sizeof( float );

	public SboxLlmModel( int formatVersion, Dictionary<string, Tensor> tensors )
	{
		FormatVersion = formatVersion;
		_tensors = tensors ?? throw new ArgumentNullException( nameof( tensors ) );

		long total = 0;
		foreach ( Tensor tensor in _tensors.Values )
		{
			if ( total > long.MaxValue - tensor.ElementCount )
			{
				throw new LlmModelFormatException(
					"Total parameter count overflows signed 64-bit." );
			}
			total += tensor.ElementCount;
		}
		if ( total > long.MaxValue / sizeof( float ) )
		{
			throw new LlmModelFormatException(
				"Total FP32 tensor byte count overflows signed 64-bit." );
		}
		TotalParameterCount = total;
	}

	public Tensor GetRequiredTensor( string name )
	{
		if ( string.IsNullOrWhiteSpace( name ) )
		{
			throw new ArgumentException( "[LLM:ERROR] Required tensor name cannot be empty.", nameof( name ) );
		}

		if ( !_tensors.TryGetValue( name, out Tensor tensor ) )
		{
			throw new KeyNotFoundException(
				$"[LLM:ERROR] Required tensor '{name}' is missing from the loaded model." );
		}

		return tensor;
	}
}