using System.Diagnostics;
using System.Runtime.CompilerServices;
using JCIL.Java.Class;

namespace JCIL.CLI;

public sealed class Block
{
	public Class? CatchType { get; init; }
	public required ReadOnlyMemory<OpCode> OpCodes { get; init; }
	public IReadOnlyList<Class> Locals { get; private set; } = [];
	public IReadOnlyList<Class> Stack { get; private set; } = [];

	public int Start => OpCodes.Span[0].Position;

	public override string ToString()
	{
		return $"Block_{OpCodes.Span[0].Position:X04}";
	}
}

public sealed class CodeAnalysis
{
	public required SortedDictionary<int, Block> Blocks { get; init; }
	
	public static CodeAnalysis Analyze(TypeLoader loader, Method method)
	{
		if (!method.TryGetCode(out var code))
			throw new ArgumentException("Method does not have a body.");
		
		using var reader = code.CreateOpCodeReader();
		var opcodes = reader.ToArray();
		
		Class? catchType = null;
		var blocks = new SortedDictionary<int, Block>();
		for (int i = 0, start = 0, count = 0; i < opcodes.Length; i++, count++)
		{
			ref readonly var opcode = ref opcodes[i];
			switch (opcode.Mnemonic)
			{
				case OpCodeMnemonic.IfEq:
				case OpCodeMnemonic.IfNe:
				case OpCodeMnemonic.IfLt:
				case OpCodeMnemonic.IfGt:
				case OpCodeMnemonic.IfLe:
				case OpCodeMnemonic.IfGe:
				case OpCodeMnemonic.IfICmpEq:
				case OpCodeMnemonic.IfICmpNe:
				case OpCodeMnemonic.IfICmpLt:
				case OpCodeMnemonic.IfICmpGt:
				case OpCodeMnemonic.IfICmpLe:
				case OpCodeMnemonic.IfICmpGe:
				case OpCodeMnemonic.IfACmpEq:
				case OpCodeMnemonic.IfACmpNe:
				case OpCodeMnemonic.Goto:
				case OpCodeMnemonic.IReturn:
				case OpCodeMnemonic.LReturn:
				case OpCodeMnemonic.FReturn:
				case OpCodeMnemonic.DReturn:
				case OpCodeMnemonic.AReturn:
				case OpCodeMnemonic.Return:
				case OpCodeMnemonic.AThrow:
				case OpCodeMnemonic.IfNull:
				case OpCodeMnemonic.IfNonNull:
				case OpCodeMnemonic.TableSwitch:
				case OpCodeMnemonic.LookupSwitch:
				{
					var blockOpCodes = opcodes.AsMemory(start, ++count);
					blocks.Add(opcodes[start].Position, new Block { OpCodes = blockOpCodes, CatchType = catchType });

					if (i != opcodes.Length - 1)
					{
						var nextPosition = opcodes[i + 1].Position;
						catchType = code.ExceptionTable.FirstOrDefault(e => e.Handler == nextPosition).CatchType;
					}
					
					start = i + 1;
					count = -1;
					break;
				}
			}
		}

		Block ResolveBlock(int addr)
		{
			if (blocks.TryGetValue(addr, out var block))
				return block;

			// Split block if an instruction jumps in the middle of it
					
			foreach (var (key, nextBlock) in blocks)
			{
				if(key <= addr) block = nextBlock;
				else break;
			}

			var span = block!.OpCodes.Span;
			for (var i = 0; i < span.Length; i++)
			{
				if (span[i].Position != addr) continue;
				blocks[span[0].Position] = new Block
				{
					OpCodes = block.OpCodes[..i],
					CatchType = block.CatchType,
				};
				return blocks[addr] = new Block
				{
					OpCodes = block.OpCodes[i..],
					CatchType = block.CatchType,
				};
			}

			throw new UnreachableException();
		}
		
		foreach (ref var opcode in opcodes.AsSpan())
		{
			switch (opcode.Mnemonic)
			{
				case OpCodeMnemonic.IfEq:
				case OpCodeMnemonic.IfNe:
				case OpCodeMnemonic.IfLt:
				case OpCodeMnemonic.IfGt:
				case OpCodeMnemonic.IfLe:
				case OpCodeMnemonic.IfGe:
				case OpCodeMnemonic.IfICmpEq:
				case OpCodeMnemonic.IfICmpNe:
				case OpCodeMnemonic.IfICmpLt:
				case OpCodeMnemonic.IfICmpGt:
				case OpCodeMnemonic.IfICmpLe:
				case OpCodeMnemonic.IfICmpGe:
				case OpCodeMnemonic.IfACmpEq:
				case OpCodeMnemonic.IfACmpNe:
				case OpCodeMnemonic.Goto:
				case OpCodeMnemonic.IfNull:
				case OpCodeMnemonic.IfNonNull:
				{
					opcode.P0.Object = ResolveBlock(opcode.Position + (short)opcode.P0.UShort);
					break;
				}

				case OpCodeMnemonic.TableSwitch:
				{
					var tableSwitch = (TableSwitch) opcode.P0.Object!;
					var targets = new Block[tableSwitch.Offsets.Length + 1];
					opcode.P1.Object = targets;
					targets[0] = ResolveBlock(opcode.Position + tableSwitch.Default);
					for (var i = 0; i < tableSwitch.Offsets.Length; i++)
						targets[i + 1] = ResolveBlock(opcode.Position + tableSwitch.Offsets[i]);
					break;
				}
				
				case OpCodeMnemonic.LookupSwitch:
				{
					var tableSwitch = (LookupSwitch) opcode.P0.Object!;
					var targets = new Block[tableSwitch.Entries.Length + 1];
					opcode.P1.Object = targets;
					targets[0] = ResolveBlock(opcode.Position + tableSwitch.Default);
					for (var i = 0; i < tableSwitch.Entries.Length; i++)
						targets[i + 1] = ResolveBlock(opcode.Position + tableSwitch.Entries[i].Value);
					break;
				}
			}
		}
		
		if(code.Attributes.OfType<StackMapTable>().FirstOrDefault() is not {} stackMap)
			return new CodeAnalysis { Blocks = blocks };

		foreach (var frame in stackMap.Frames)
			ResolveBlock(frame.StartOffset);

		var stack = new List<Class>(code.MaxStack);
		var locals = new List<Class>(code.MaxLocals);
		if((method.AccessFlags & MethodAccessFlags.Static) == 0)
			locals.Add(method.DeclaringClass);
		
		Class ResolveClass(StackMapTable.VerificationTypeInfo info)
		{
			switch (info.Tag)
			{
				case StackMapTable.VerificationTypeInfoTag.Top: return locals[^1];
				case StackMapTable.VerificationTypeInfoTag.Integer: return loader.Int;
				case StackMapTable.VerificationTypeInfoTag.Float: return loader.Float;
				case StackMapTable.VerificationTypeInfoTag.Double: return loader.Double;
				case StackMapTable.VerificationTypeInfoTag.Long: return loader.Long;
				case StackMapTable.VerificationTypeInfoTag.Object: return method.DeclaringClass.GetConstant<Class>(info.Parameter); 
				case StackMapTable.VerificationTypeInfoTag.Null: 
				case StackMapTable.VerificationTypeInfoTag.UninitializedThis:
				case StackMapTable.VerificationTypeInfoTag.Uninitialized:
				default: throw new NotImplementedException(info.Tag.ToString());
			}
		}
		
		locals.AddRange(method.Signature.ParamTypes);
		foreach (var block in blocks.Values)
		{
			if(block.Start >= stackMap.Frames[0].StartOffset) break;
			LocalsRef(block) = locals.ToArray();
		}

		for (var i = 0; i < stackMap.Frames.Count; i++)
		{
			var frame = stackMap.Frames[i];
			var end = i == stackMap.Frames.Count - 1 ? code.Bytecode.Length : stackMap.Frames[i + 1].StartOffset;
			switch (frame)
			{
				case StackMapTable.SameFrame:
				case StackMapTable.SameFrameExtended:
					stack.Clear();
					break;
				
				case StackMapTable.FullFrame f:
					stack.Clear();
					locals.Clear();
					locals.AddRange(f.Locals.Select(ResolveClass));
					stack.AddRange(f.StackItems.Select(ResolveClass));
					break;
				
				case StackMapTable.ChopFrame f:
					stack.Clear();
					locals.RemoveRange(locals.Count - f.Count, f.Count);
					break;

				case StackMapTable.AppendFrame f:
					stack.Clear();
					locals.AddRange(f.Locals.Select(ResolveClass));
					break;

				case StackMapTable.OneTempFrame f:
					stack.Clear();
					stack.Add(ResolveClass(f.Entry));
					break;
				
				case StackMapTable.OneTempFrameExtended f:
					stack.Clear();
					stack.Add(ResolveClass(f.Entry));
					break;
				
				default: throw new NotImplementedException(frame.ToString());
			}
			foreach (var block in blocks.Values)
			{
				if(block.Start < frame.StartOffset) continue;
				if(block.Start >= end) break;
				LocalsRef(block) = locals.ToArray();
				StackRef(block) = stack.ToArray();
			}
		}
		
		return new CodeAnalysis { Blocks = blocks };
	}

	[UnsafeAccessor(UnsafeAccessorKind.Field, Name = "<Locals>k__BackingField")]
	private static extern ref IReadOnlyList<Class> LocalsRef(Block block);
	
	[UnsafeAccessor(UnsafeAccessorKind.Field, Name = "<Stack>k__BackingField")]
	private static extern ref IReadOnlyList<Class> StackRef(Block block);
}