using System;
using System.Collections.Generic;

namespace GlinerPoc.Neural;

/// <summary>
/// Phase 4.4/4.5/4.19 — FP32 weight access over the packaged GLiNER resource.
///
/// Runtime source of truth: the normal packaged s&box model resource
/// (root manifest + eagerly-materialized shard payloads). No Hugging Face or
/// development files are read.
///
/// Ownership strategy (Phase 2/4 measured + documented):
/// - Shard payload byte[]s are engine-owned via the resource BlobData and are
///   NOT duplicated here (references only).
/// - Small tensors may be decoded+cached on demand (documented per call site).
/// - The 196.6 MB word embedding table is NEVER fully decoded: row-range reads
///   go through the SBGLI1 chunk metadata and decode only the requested rows
///   (per row: 384 BitConverter.ToSingle calls straight from the shard bytes).
/// </summary>
public sealed class GlinerModelWeights
{
	private readonly GlinerPoc.Packaging.Sbgli1Manifest _manifest;
	private readonly byte[][] _shardData;
	private readonly Dictionary<string, float[]> _decodedCache = new( StringComparer.Ordinal );

	public GlinerModelWeights( GlinerPoc.Packaging.Sbgli1Manifest manifest, byte[][] shardData )
	{
		_manifest = manifest ?? throw new ArgumentNullException( nameof( manifest ) );
		_shardData = shardData ?? throw new ArgumentNullException( nameof( shardData ) );
	}

	public GlinerPoc.Packaging.Sbgli1Manifest Manifest => _manifest;

	public static GlinerModelWeights FromResource( GlinerPoc.Packaging.GlinerModelResource root )
	{
		var manifest = GlinerPoc.Packaging.Sbgli1Manifest.Parse( root.MetadataData.Bytes );
		if ( root.Shards.Count != manifest.ShardCount )
		{
			throw new InvalidOperationException( "[GLI:ERROR] Shard reference count != manifest shard count." );
		}
		var shardData = new byte[manifest.ShardCount][];
		for ( int i = 0; i < root.Shards.Count; i++ )
		{
			var bytes = root.Shards[i].Shard?.Payload?.Bytes;
			if ( bytes is not { Length: > 0 } )
			{
				throw new InvalidOperationException( $"[GLI:ERROR] Shard {i} payload not materialized." );
			}
			shardData[i] = bytes;
		}
		return new GlinerModelWeights( manifest, shardData );
	}

	/// <summary>Decode a FULL tensor (use only for small tensors; cached).</summary>
	public float[] GetTensor( string name )
	{
		if ( _decodedCache.TryGetValue( name, out var cached ) )
		{
			return cached;
		}
		var record = _manifest.GetRequiredTensor( name );
		if ( record.ByteCount > 8 * 1024 * 1024 )
		{
			throw new InvalidOperationException(
				$"[GLI:ERROR] Refusing full decode of large tensor '{name}' ({record.ByteCount} B); " +
				"use row/range access." );
		}
		byte[] bytes = _manifest.ReadTensorBytes( name, _shardData );
		var result = GlinerMath.DecodeF32( bytes, 0, (int)record.ElementCount );
		_decodedCache[name] = result;
		return result;
	}

	/// <summary>Tensor metadata lookup.</summary>
	public GlinerPoc.Packaging.Sbgli1Manifest.TensorRecord GetRecord( string name ) =>
		_manifest.GetRequiredTensor( name );

	/// <summary>
	/// Row-range read for a rank-2 tensor stored row-major: copies only the
	/// requested rows, correctly splicing across SBGLI1 row-range chunk
	/// boundaries. Returns rows×cols floats. Bit-exact source values.
	/// </summary>
	public float[] GetRows( string name, int rowStart, int rowCount )
	{
		var record = _manifest.GetRequiredTensor( name );
		if ( record.Rank != 2 )
		{
			throw new InvalidOperationException( $"[GLI:ERROR] GetRows requires rank 2, '{name}' has rank {record.Rank}." );
		}
		int cols = record.Shape[1];
		int totalRows = record.Shape[0];
		if ( rowStart < 0 || rowCount <= 0 || rowStart + rowCount > totalRows )
		{
			throw new InvalidOperationException(
				$"[GLI:ERROR] Row range [{rowStart}, {rowStart + rowCount}) outside '{name}' ({totalRows} rows)." );
		}

		var result = new float[rowCount * cols];
		int row = rowStart;
		int done = 0;
		while ( done < rowCount )
		{
			var chunk = FindChunkCoveringRow( record, row );
			int chunkRowStart = chunk.RowStart;
			int chunkRowCount = chunk.RowCount;
			int take = Math.Min( chunkRowStart + chunkRowCount - row, rowCount - done );
			byte[] shard = _shardData[chunk.ShardIndex];
			long rowBytes = cols * 4;
			long srcByte = chunk.ByteOffsetInShard + (row - chunkRowStart) * rowBytes;
			int dstBase = done * cols;
			for ( int r = 0; r < take; r++ )
			{
				long rb = srcByte + r * rowBytes;
				int db = dstBase + r * cols;
				for ( int c = 0; c < cols; c++ )
				{
					result[db + c] = BitConverter.ToSingle( shard, (int)(rb + c * 4) );
				}
			}
			row += take;
			done += take;
		}
		return result;
	}

	private GlinerPoc.Packaging.Sbgli1Manifest.ChunkRecord FindChunkCoveringRow(
		GlinerPoc.Packaging.Sbgli1Manifest.TensorRecord record, int row )
	{
		for ( int ci = record.FirstChunkIndex; ci < record.FirstChunkIndex + record.ChunkCount; ci++ )
		{
			var c = _manifest.Chunks[ci];
			if ( row >= c.RowStart && row < c.RowStart + c.RowCount )
			{
				return c;
			}
		}
		throw new InvalidOperationException(
			$"[GLI:ERROR] No chunk covers row {row} of '{record.Name}'." );
	}
}