using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.Text;
using NubLang.Syntax.Binding;
using NubLang.Syntax.Binding.Node;
using NubLang.Syntax.Tokenization;

namespace NubLang.Generation.QBE;

public partial class QBEGenerator
{
    private readonly BoundSyntaxTree _syntaxTree;
    private readonly BoundDefinitionTable _definitionTable;
    private readonly QBEWriter _writer;

    private readonly List<CStringLiteral> _cStringLiterals = [];
    private readonly List<StringLiteral> _stringLiterals = [];
    private readonly Stack<string> _breakLabels = [];
    private readonly Stack<string> _continueLabels = [];
    private readonly Queue<(BoundArrowFunc Func, string Name)> _arrowFunctions = [];
    private readonly Stack<Scope> _scopes = [];
    private int _tmpIndex;
    private int _labelIndex;
    private int _arrowFuncIndex;
    private int _cStringLiteralIndex;
    private int _stringLiteralIndex;
    private bool _codeIsReachable = true;

    private Scope Scope => _scopes.Peek();

    public QBEGenerator(BoundSyntaxTree syntaxTree, BoundDefinitionTable definitionTable)
    {
        _syntaxTree = syntaxTree;
        _definitionTable = definitionTable;
        _writer = new QBEWriter();
    }

    public string Emit()
    {
        _cStringLiterals.Clear();
        _stringLiterals.Clear();
        _breakLabels.Clear();
        _continueLabels.Clear();
        _arrowFunctions.Clear();
        _scopes.Clear();
        _tmpIndex = 0;
        _labelIndex = 0;
        _arrowFuncIndex = 0;
        _cStringLiteralIndex = 0;
        _stringLiteralIndex = 0;
        _codeIsReachable = true;

        foreach (var structDef in _definitionTable.GetStructs())
        {
            EmitStructDefinition(structDef);
            _writer.NewLine();
        }

        foreach (var trait in _definitionTable.GetTraits())
        {
            EmitTraitVTable(trait);
            _writer.NewLine();
        }

        foreach (var funcDef in _syntaxTree.Definitions.OfType<BoundLocalFunc>())
        {
            EmitFuncDefinition(LocalFuncName(funcDef), funcDef.Signature.Parameters, funcDef.Signature.ReturnType, funcDef.Body);
            _writer.NewLine();
        }

        while (_arrowFunctions.TryDequeue(out var arrowFunc))
        {
            EmitFuncDefinition(arrowFunc.Name, arrowFunc.Func.Parameters, arrowFunc.Func.ReturnType, arrowFunc.Func.Body);
            _writer.NewLine();
        }

        foreach (var cStringLiteral in _cStringLiterals)
        {
            _writer.WriteLine($"data {cStringLiteral.Name} = {{ b \"{cStringLiteral.Value}\", b 0 }}");
        }

        foreach (var stringLiteral in _stringLiterals)
        {
            var bytes = Encoding.UTF8.GetBytes(stringLiteral.Value).Select(b => $"b {b}");
            _writer.WriteLine($"data {stringLiteral.Name} = {{ l {stringLiteral.Value.Length}, {string.Join(", ", bytes)} }}");
        }

        return _writer.ToString();
    }

    private static string QBEAssign(NubType type)
    {
        if (type.IsSimpleType(out var simpleType, out _))
        {
            return simpleType.StorageSize switch
            {
                StorageSize.I8 or StorageSize.U8 or StorageSize.I16 or StorageSize.U16 or StorageSize.I32 or StorageSize.U32 => "=w",
                StorageSize.I64 or StorageSize.U64 => "=l",
                StorageSize.F32 => "=s",
                StorageSize.F64 => "=d",
                _ => throw new ArgumentOutOfRangeException(nameof(simpleType.StorageSize))
            };
        }

        return "=l";
    }

    private void EmitStore(NubType type, string value, string destination)
    {
        string store;

        if (type.IsSimpleType(out var simpleType, out _))
        {
            store = simpleType.StorageSize switch
            {
                StorageSize.I8 or StorageSize.U8 => "storeb",
                StorageSize.I16 or StorageSize.U16 => "storeh",
                StorageSize.I32 or StorageSize.U32 => "storew",
                StorageSize.I64 or StorageSize.U64 => "storel",
                StorageSize.F32 => "stores",
                StorageSize.F64 => "stored",
                _ => throw new ArgumentOutOfRangeException(nameof(simpleType.StorageSize))
            };
        }
        else
        {
            store = "storel";
        }

        _writer.Indented($"{store} {value}, {destination}");
    }

    private Val EmitLoad(NubType type, string from)
    {
        string load;

        if (type.IsSimpleType(out var simpleType, out _))
        {
            load = simpleType.StorageSize switch
            {
                StorageSize.I64 or StorageSize.U64 => "loadl",
                StorageSize.I32 or StorageSize.U32 => "loadw",
                StorageSize.I16 => "loadsh",
                StorageSize.I8 => "loadsb",
                StorageSize.U16 => "loaduh",
                StorageSize.U8 => "loadub",
                StorageSize.F64 => "loadd",
                StorageSize.F32 => "loads",
                _ => throw new ArgumentOutOfRangeException(nameof(simpleType.StorageSize))
            };
        }
        else
        {
            load = "loadl";
        }

        var into = TmpName();

        _writer.Indented($"{into} {QBEAssign(type)} {load} {from}");

        return new Val(into, type, ValKind.Direct);
    }

    private void EmitMemcpy(string source, string destination, string length)
    {
        _writer.Indented($"call $nub_memcpy(l {source}, l {destination}, l {length})");
    }

    private string EmitArraySizeInBytes(NubArrayType type, string array)
    {
        var size = TmpName();
        _writer.Indented($"{size} =l loadl {array}");
        _writer.Indented($"{size} =l mul {size}, {type.ElementType.Size(_definitionTable)}");
        _writer.Indented($"{size} =l add {size}, 8");
        return size;
    }

    private string EmitCStringSizeInBytes(string cstring)
    {
        var size = TmpName();
        _writer.Indented($"{size} =l call $nub_cstring_length(l {cstring})");
        _writer.Indented($"{size} =l add {size}, 1");
        return size;
    }

    private string EmitStringSizeInBytes(string nubstring)
    {
        var size = TmpName();
        _writer.Indented($"{size} =l loadl {nubstring}");
        _writer.Indented($"{size} =l add {size}, 8");
        return size;
    }

    private bool EmitTryMoveInto(BoundExpression source, string destinationPointer)
    {
        switch (source)
        {
            case BoundArrayInitializer arrayInitializer:
            {
                EmitStore(source.Type, EmitUnwrap(EmitArrayInitializer(arrayInitializer)), destinationPointer);
                return true;
            }
            case BoundStructInitializer structInitializer:
            {
                EmitStructInitializer(structInitializer, destinationPointer);
                return true;
            }
            case BoundLiteral { Kind: LiteralKind.String } literal:
            {
                EmitStore(source.Type, EmitUnwrap(EmitLiteral(literal)), destinationPointer);
                return true;
            }
        }

        return false;
    }

    private void EmitCopyIntoOrInitialize(BoundExpression source, string destinationPointer)
    {
        // If the source is a value which is not used yet such as an array/struct initializer or literal, we can skip copying
        if (EmitTryMoveInto(source, destinationPointer))
        {
            return;
        }

        var value = EmitUnwrap(EmitExpression(source));

        if (source.Type.IsSimpleType(out var simpleType, out var complexType))
        {
            EmitStore(simpleType, value, destinationPointer);
        }
        else
        {
            if (complexType is NubCustomType customType)
            {
                EmitMemcpy(value, destinationPointer, customType.Size(_definitionTable).ToString());
            }
            else
            {
                var size = complexType switch
                {
                    NubArrayType arrayType => EmitArraySizeInBytes(arrayType, value),
                    NubCStringType => EmitCStringSizeInBytes(value),
                    NubStringType => EmitStringSizeInBytes(value),
                    _ => throw new ArgumentOutOfRangeException(nameof(source.Type))
                };

                var buffer = TmpName();
                _writer.Indented($"{buffer} =l alloc8 {size}");
                EmitMemcpy(value, buffer, size);
                EmitStore(complexType, buffer, destinationPointer);
            }
        }
    }

    private bool EmitTryCreateWithoutCopy(BoundExpression source, [NotNullWhen(true)] out string? destination)
    {
        switch (source)
        {
            case BoundArrayInitializer:
            case BoundStructInitializer:
            case BoundLiteral { Kind: LiteralKind.String }:
            {
                destination = EmitUnwrap(EmitExpression(source));
                return true;
            }
        }

        destination = null;
        return false;
    }

    private string EmitCreateCopyOrInitialize(BoundExpression source)
    {
        // If the source is a value which is not used yet such as an array/struct initializer or literal, we can skip copying
        if (EmitTryCreateWithoutCopy(source, out var uncopiedValue))
        {
            return uncopiedValue;
        }

        var value = EmitUnwrap(EmitExpression(source));

        if (source.Type.IsSimpleType(out _, out var complexType))
        {
            // Simple types are passed in registers and are therefore always copied
            return value;
        }

        var size = complexType switch
        {
            NubArrayType arrayType => EmitArraySizeInBytes(arrayType, value),
            NubCStringType => EmitCStringSizeInBytes(value),
            NubStringType => EmitStringSizeInBytes(value),
            NubCustomType customType => customType.Size(_definitionTable).ToString(),
            _ => throw new ArgumentOutOfRangeException(nameof(source.Type))
        };

        var destination = TmpName();
        _writer.Indented($"{destination} =l alloc8 {size}");
        EmitMemcpy(value, destination, size);
        return destination;
    }

    // Utility to create QBE type names for function parameters and return types
    private string FuncQBETypeName(NubType type)
    {
        if (type.IsSimpleType(out var simpleType, out var complexType))
        {
            return simpleType.StorageSize switch
            {
                StorageSize.I64 or StorageSize.U64 => "l",
                StorageSize.I32 or StorageSize.U32 => "w",
                StorageSize.I16 => "sh",
                StorageSize.I8 => "sb",
                StorageSize.U16 => "uh",
                StorageSize.U8 => "ub",
                StorageSize.F64 => "d",
                StorageSize.F32 => "s",
                _ => throw new ArgumentOutOfRangeException()
            };
        }

        if (complexType is NubCustomType customType)
        {
            return CustomTypeName(customType);
        }

        return "l";
    }

    private void EmitFuncDefinition(string name, IReadOnlyList<BoundFuncParameter> parameters, NubType returnType, BoundBlock body)
    {
        _labelIndex = 0;
        _tmpIndex = 0;

        var builder = new StringBuilder();

        builder.Append("export function ");

        if (returnType is not NubVoidType)
        {
            builder.Append(FuncQBETypeName(returnType) + ' ');
        }

        builder.Append(name);

        var parameterStrings = parameters.Select(x => FuncQBETypeName(x.Type) + $" %{x.Name}");

        builder.Append($"({string.Join(", ", parameterStrings)})");

        _writer.StartFunction(builder.ToString());

        var scope = new Scope();

        foreach (var parameter in parameters)
        {
            scope.Declare(parameter.Name, new Val("%" + parameter.Name, parameter.Type, ValKind.Direct));
        }

        EmitBlock(body, scope);

        // Implicit return for void functions if no explicit return has been set
        if (returnType is NubVoidType && body.Statements is [.., not BoundReturn])
        {
            if (returnType is NubVoidType)
            {
                _writer.Indented("ret");
            }
        }

        _writer.EndFunction();
    }

    private void EmitStructDefinition(BoundStruct structDef)
    {
        _writer.WriteLine($"type {CustomTypeName(structDef.Namespace, structDef.Name)} = {{ ");

        var types = new Dictionary<string, string>();

        foreach (var field in structDef.Fields)
        {
            types.Add(field.Name, StructDefQBEType(field));
        }

        var longest = types.Values.Max(x => x.Length);
        foreach (var (name, type) in types)
        {
            var padding = longest - type.Length;
            _writer.Indented($"{type},{new string(' ', padding)} # {name}");
        }

        _writer.WriteLine("}");
        return;

        string StructDefQBEType(BoundStructField field)
        {
            if (field.Type.IsSimpleType(out var simpleType, out var complexType))
            {
                return simpleType.StorageSize switch
                {
                    StorageSize.I64 or StorageSize.U64 => "l",
                    StorageSize.I32 or StorageSize.U32 => "w",
                    StorageSize.I16 or StorageSize.U16 => "h",
                    StorageSize.I8 or StorageSize.U8 => "b",
                    StorageSize.F64 => "d",
                    StorageSize.F32 => "s",
                    _ => throw new ArgumentOutOfRangeException()
                };
            }

            if (complexType is NubCustomType customType)
            {
                return CustomTypeName(customType);
            }

            return "l";
        }
    }

    private void EmitTraitVTable(BoundTrait traitDef)
    {
        _writer.WriteLine($"type {CustomTypeName(traitDef.Namespace, traitDef.Name)} = {{");

        foreach (var func in traitDef.Functions)
        {
            _writer.Indented($"l, # func {func.Name}({string.Join(", ", func.Signature.Parameters.Select(x => $"{x.Name}: {x.Type}"))}): {func.Signature.ReturnType}");
        }

        _writer.WriteLine("}");
    }

    private void EmitBlock(BoundBlock block, Scope? scope = null)
    {
        _scopes.Push(scope ?? Scope.SubScope());

        foreach (var statement in block.Statements)
        {
            if (_codeIsReachable)
            {
                EmitStatement(statement);
            }
        }

        _scopes.Pop();

        _codeIsReachable = true;
    }

    private string EmitUnwrap(Val val)
    {
        return val.Kind switch
        {
            ValKind.Direct => val.Name,
            ValKind.Pointer => EmitLoad(val.Type, val.Name).Name,
            _ => throw new ArgumentOutOfRangeException()
        };
    }

    private int OffsetOf(BoundStruct structDefinition, string member)
    {
        var offset = 0;

        foreach (var field in structDefinition.Fields)
        {
            if (field.Name == member)
            {
                return offset;
            }

            var fieldAlignment = field.Type.Alignment(_definitionTable);

            offset = NubType.AlignTo(offset, fieldAlignment);
            offset += field.Type.Size(_definitionTable);
        }

        throw new UnreachableException($"Member '{member}' not found in struct");
    }

    #region Naming utilities

    private string TmpName()
    {
        return $"%t{++_tmpIndex}";
    }

    private string LabelName()
    {
        return $"@l{++_labelIndex}";
    }

    private string CStringName()
    {
        return $"$cstring{++_cStringLiteralIndex}";
    }

    private string StringName()
    {
        return $"$string{++_stringLiteralIndex}";
    }

    private string LocalFuncName(BoundLocalFunc funcDef)
    {
        return $"${funcDef.Namespace}_{funcDef.Name}";
    }

    private string ExternFuncName(BoundExternFunc funcDef)
    {
        return $"${funcDef.CallName}";
    }

    private string CustomTypeName(NubCustomType customType)
    {
        return CustomTypeName(customType.Namespace, customType.Name);
    }

    private string CustomTypeName(string @namespace, string name)
    {
        return $":{@namespace}_{name}";
    }

    #endregion
}

public class StringLiteral(string value, string name)
{
    public string Value { get; } = value;
    public string Name { get; } = name;
}

public class CStringLiteral(string value, string name)
{
    public string Value { get; } = value;
    public string Name { get; } = name;
}

public record Val(string Name, NubType Type, ValKind Kind);

public class Scope(Scope? parent = null)
{
    private readonly Dictionary<string, Val> _variables = [];

    public Val Lookup(string name)
    {
        var variable = _variables.GetValueOrDefault(name);
        if (variable != null)
        {
            return variable;
        }

        return parent?.Lookup(name) ?? throw new UnreachableException($"Variable '{name}' not found");
    }

    public void Declare(string name, Val value)
    {
        _variables.Add(name, value);
    }

    public Scope SubScope()
    {
        return new Scope(this);
    }
}

public enum ValKind
{
    Pointer,
    Direct,
}