#include "LuaCall.h"
#include "LuprexGameModeBase.h"
#include "StringDecoder.h"

#include "EdGraphSchema_K2.h"


static bool NotInitialized(const FlxStreamBuffer &sb) {
    if (sb.empty())
	{
        UE_LOG(LogBlueprint, Error, TEXT("Must use LuaCallBegin before other LuaCall steps"));
		return true;
    }
	return false;
}

static constexpr uint64_t ParseNameAsToken(std::string_view str) {
    uint64_t result = 0;
    uint64_t maxint = uint64_t(-1);

    // Leading zeros are not allowed.
    if ((!str.empty()) && (str[0]=='0')) return 0;

    for (int i = 0; i < int(str.size()); i++) {
        char c = str[i];
        uint64_t digit = 0;
        if ((c >= '0') && (c <= '9')) {
            digit = uint64_t(c - '0');
        } else if ((c >= 'a') && (c <= 'z')) {
            digit = uint64_t(c - 'a' + 10);
        } else if ((c >= 'A') && (c <= 'Z')) {
            digit = uint64_t(c - 'A' + 10);
        } else {
            return maxint;
        }
        // Multiply existing number by 36, then add the digit.
        // We have two checks to prevent integer overflow.
        if (result > (maxint / 36)) return 0;
        result *= 36;
        if (digit > (maxint - result)) return 0;
        result += digit;
    }
    return result;
}


/////////////////////////////////////////////////////////////////
//
// The lua function prototype parser.
//
/////////////////////////////////////////////////////////////////


bool FlxParsedProto::IsLiteral(const TCHAR *text) {
	return ((NextToken < Tokens.Num()) && (Tokens[NextToken] == text));
}

bool FlxParsedProto::IsIdent() {
	return ((NextToken < Tokens.Num()) && (FChar::IsAlpha(Tokens[NextToken][0])));
}
	
void FlxParsedProto::Empty() {
	ErrorMessage = TEXT("");
	Tokens.Empty();
	NextToken = 0;
	ClassName = TEXT("");
	FunctionName = TEXT("");
	Arguments.Empty();
	ReturnValues.Empty();
	ExtraReturnValues = false;
}

void FlxParsedProto::Syntax() {
	FString Message;
	if (Tokens.Num() == 0) {
		Message = TEXT("Function prototype cannot be blank");
	}
	for (int i = 0; i < Tokens.Num(); i++) {
		if (i == NextToken) {
			Message.Append(TEXT("  ?  "));
		} else {
			if ((i > 0) && (FChar::IsAlpha(Tokens[i][0])) && (FChar::IsAlpha(Tokens[i-1][0]))) {
				Message.Append(TEXT(" "));
			}
		}
		Message.Append(Tokens[i]);
	}
	Empty();
	ErrorMessage = Message;
}

void FlxParsedProto::Parse(const FString &str) {
	Empty();

	// Step one: tokenize.
	int offset = 0;
	while (offset < str.Len()) {
		TCHAR c = str[offset];
		if (FChar::IsWhitespace(c)) {
			offset++;
		} else if (FChar::IsAlpha(c)) {
			int lo = offset;
			while ((offset < str.Len()) && FChar::IsAlnum(str[offset])) offset++;
			Tokens.Add(str.Mid(lo, offset-lo));
		} else if (str.Mid(offset, 3) == TEXT("...")) {
			Tokens.Add(str.Mid(offset, 3));
			offset += 3;
		} else if (FChar::IsPunct(c)) {
			Tokens.Add(str.Mid(offset, 1));
			offset += 1;
		} else {
			Empty();
			ErrorMessage = FString::Printf(TEXT("%s  ?  %s"), *str.Mid(0, offset), *str.Mid(offset));
			return;
		}
	}
	NextToken = 0;

	// Step two: Parse.
	if (!IsLiteral(TEXT("*")) && !IsIdent()) return Syntax();
	ClassName = Tokens[NextToken++];
	if (!IsLiteral(TEXT("."))) return Syntax();
	NextToken++;
	if (!IsIdent()) return Syntax();
	FunctionName = Tokens[NextToken++];
	if (!IsLiteral(TEXT("("))) return Syntax();
	NextToken++;
	if (IsIdent()) {
		while (true) {
			if (!IsIdent()) return Syntax();
			FString Type = Tokens[NextToken++];
			if (!IsIdent()) return Syntax();
			FString Name = Tokens[NextToken++];
			Arguments.Add(Pin(Type, Name));
			if (!IsLiteral(TEXT(","))) break;
			NextToken++;
		}
	}
	if (!IsLiteral(TEXT(")"))) return Syntax();
	NextToken++;
	if (IsLiteral(TEXT(":"))) {
		NextToken++;
		while (true) {
			if (IsLiteral(TEXT("..."))) {
				ExtraReturnValues = true;
				NextToken++;
				break;
			} else if (IsIdent()) {
				FString Type = Tokens[NextToken++];
				if (!IsIdent()) return Syntax();
				FString Name = Tokens[NextToken++];
				ReturnValues.Add(Pin(Type, Name));
				if (!IsLiteral(TEXT(","))) break;
				NextToken++;
			} else {
				return Syntax();
			}
		}
	}
	if (NextToken != Tokens.Num()) return Syntax();
}

/////////////////////////////////////////////////////////////////
//
// General functions of the Lua Call Library.
//
/////////////////////////////////////////////////////////////////

UFunction *UlxLuaCallLibrary::GetArgumentPacker(const FString &Type)
{
	FString LType = Type.ToLower();
	FName PackerName(FString::Printf(TEXT("LuaCallArgument_%s"), *LType));
	return UlxLuaCallLibrary::StaticClass()->FindFunctionByName(PackerName);
}

UFunction *UlxLuaCallLibrary::GetReturnValueUnpacker(const FString &Type)
{
	FString LType = Type.ToLower();
	FName PackerName(FString::Printf(TEXT("LuaCallReturnValue_%s"), *LType));
	return UlxLuaCallLibrary::StaticClass()->FindFunctionByName(PackerName);
}

FString UlxLuaCallLibrary::AllFunctionsWithPrefix(const TCHAR *Prefix)
{
    FString Result;

	for ( TFieldIterator<UFunction> FIT ( UlxLuaCallLibrary::StaticClass(), EFieldIteratorFlags::ExcludeSuper); FIT; ++FIT)
	{
        UFunction* Function = *FIT;
		FString Name = Function->GetName();
		if (Name.RemoveFromStart(Prefix))
		{
			if (!Result.IsEmpty()) Result += TEXT(", ");
			Result += Name;
		}
    }
    return Result;
}

/////////////////////////////////////////////////////////////////
//
// General Lua-Callable functions of the Lua Call Library.
//
/////////////////////////////////////////////////////////////////

void UlxLuaCallLibrary::LuaCallBegin(UObject *context, const FString &cname, const FString &fname)
{
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetBuffer();
    mode->LuaCallBegin();
    sb.write_string(cname);
    sb.write_string(fname);
}


void UlxLuaCallLibrary::LuaCallInvoke(UObject *context, AActor *place)
{
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetBuffer();
    if (NotInitialized(sb)) return;
	mode->LuaCallEnd(InvocationKind::LUA_INVOKE, place);
}


bool UlxLuaCallLibrary::LuaCallProbe(UObject *context, AActor *place)
{
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetBuffer();
    if (NotInitialized(sb)) return false;
    mode->LuaCallEnd(InvocationKind::LUA_PROBE, place);
    FlxStreamBuffer &result = mode->LuaCallGetResult();
    SimpleDynamicTag tag = result.read_simple_dynamic_tag();
    if (tag != SimpleDynamicTag::STRING)
	{
		UE_LOG(LogLuprexIntegration, Error, TEXT("corruption in lua_probe"));
		return false;
	}
    FString ErrorMessage = result.read_fstring();
    if (!ErrorMessage.IsEmpty())
    {
        UE_LOG(LogLuprex, Error, TEXT("%s"), *ErrorMessage);
        return false;
    }
    return true;
}

UlxLuaValues *UlxLuaCallLibrary::LuaCallGetRest(UObject *context)
{
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetResult();
	UlxLuaValues *Values = NewObject<UlxLuaValues>(context);
	if (!Values->Initialize(sb.view()))
	{
		UE_LOG(LogBlueprint, Error, TEXT("Lua call returned corrupt data"));
		return nullptr;
	}
	return Values;
}

/////////////////////////////////////////////////////////////////
//
// Argument Packing functions
//
/////////////////////////////////////////////////////////////////

void UlxLuaCallLibrary::LuaCallArgument_string(UObject *context, const FString &pstring) {
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetBuffer();
    if (NotInitialized(sb)) return;
	sb.write_simple_dynamic_tag(SimpleDynamicTag::STRING);
	sb.write_string(pstring);
}
	
void UlxLuaCallLibrary::LuaCallArgument_name(UObject *context, const FName &pname) {
    FTCHARToUTF8 utf8str(pname.ToString());
    std::string_view namestr(utf8str.Get(), utf8str.Length());
    uint64_t tokvalue = ParseNameAsToken(namestr);
    if ((tokvalue == 0) && !namestr.empty())
	{
		UE_LOG(LogBlueprint, Error, TEXT("Names passed to lua must be short, and must contain only lowercase and digits"));
    }
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetBuffer();
    if (NotInitialized(sb)) return;
	sb.write_simple_dynamic_tag(SimpleDynamicTag::TOKEN);
	sb.write_string(namestr);
}

void UlxLuaCallLibrary::LuaCallArgument_float(UObject *context, double pfloat) {
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetBuffer();
    if (NotInitialized(sb)) return;
	sb.write_simple_dynamic_tag(SimpleDynamicTag::NUMBER);
	sb.write_double(pfloat);
}
	
void UlxLuaCallLibrary::LuaCallArgument_int(UObject *context, int value) {
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetBuffer();
    if (NotInitialized(sb)) return;
	sb.write_simple_dynamic_tag(SimpleDynamicTag::NUMBER);
	sb.write_double(value);
}

void UlxLuaCallLibrary::LuaCallArgument_vector(UObject *context, const FVector &pvector) {
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetBuffer();
    if (NotInitialized(sb)) return;
	sb.write_simple_dynamic_tag(SimpleDynamicTag::VECTOR);
	sb.write_fvector(pvector);
}

void UlxLuaCallLibrary::LuaCallArgument_vector2d(UObject *context, const FVector2D &pvector) {
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetBuffer();
    if (NotInitialized(sb)) return;
	sb.write_simple_dynamic_tag(SimpleDynamicTag::VECTOR);
	sb.write_double(pvector.X);
	sb.write_double(pvector.Y);
	sb.write_double(0.0);
}

void UlxLuaCallLibrary::LuaCallArgument_boolean(UObject *context, bool pbool) {
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetBuffer();
    if (NotInitialized(sb)) return;
	sb.write_simple_dynamic_tag(SimpleDynamicTag::BOOLEAN);
	sb.write_bool(pbool);
}


/////////////////////////////////////////////////////////////////
//
// Return Value Unpacking functions
//
/////////////////////////////////////////////////////////////////


FString UlxLuaCallLibrary::LuaCallReturnValue_string(UObject *context) {
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetResult();
    SimpleDynamicTag tag = sb.read_simple_dynamic_tag();
    if (tag != SimpleDynamicTag::STRING)
	{
		UE_LOG(LogBlueprint, Error, TEXT("expected lua to return a string"));
		return TEXT("");
	}
    return sb.read_fstring();
}

FName UlxLuaCallLibrary::LuaCallReturnValue_name(UObject *context) {
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetResult();
    SimpleDynamicTag tag = sb.read_simple_dynamic_tag();
    if (tag != SimpleDynamicTag::TOKEN)
	{
		UE_LOG(LogBlueprint, Error, TEXT("expected lua to return a name"));
		return FName();
	}
    return sb.read_fname();
}

double UlxLuaCallLibrary::LuaCallReturnValue_float(UObject *context) {
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetResult();
    SimpleDynamicTag tag = sb.read_simple_dynamic_tag();
    if (tag != SimpleDynamicTag::NUMBER)
	{
		UE_LOG(LogBlueprint, Error, TEXT("expected lua to return a float"));
		return 0.0;
	}
    return sb.read_double();
}

int UlxLuaCallLibrary::LuaCallReturnValue_int(UObject *context) {
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetResult();
    SimpleDynamicTag tag = sb.read_simple_dynamic_tag();
    if (tag != SimpleDynamicTag::NUMBER)
	{
		UE_LOG(LogBlueprint, Error, TEXT("expected lua to return a number"));
		return 0;
	}
    return int(sb.read_double());
}

FVector UlxLuaCallLibrary::LuaCallReturnValue_vector(UObject *context) {
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetResult();
    SimpleDynamicTag tag = sb.read_simple_dynamic_tag();
    if (tag != SimpleDynamicTag::VECTOR)
	{
		UE_LOG(LogBlueprint, Error, TEXT("expected lua to return a vector"));
		return FVector();
	}
    return sb.read_fvector();
}

FVector2D UlxLuaCallLibrary::LuaCallReturnValue_vector2d(UObject *context) {
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetResult();
    SimpleDynamicTag tag = sb.read_simple_dynamic_tag();
    if (tag != SimpleDynamicTag::VECTOR)
	{
		UE_LOG(LogBlueprint, Error, TEXT("expected lua to return a vector"));
		return FVector2D();
	}
	FVector v = sb.read_fvector();
    return FVector2D(v.X, v.Y);
}

bool UlxLuaCallLibrary::LuaCallReturnValue_boolean(UObject *context) {
    ALuprexGameModeBase *mode = ALuprexGameModeBase::FromContext(context);
    FlxStreamBuffer &sb = mode->LuaCallGetResult();
    SimpleDynamicTag tag = sb.read_simple_dynamic_tag();
    if (tag != SimpleDynamicTag::BOOLEAN)
	{
		UE_LOG(LogBlueprint, Error, TEXT("expected lua to return a boolean"));
		return false;
	}
	return sb.read_bool();
}

/////////////////////////////////////////////////////////////////
//
// Returning the rest of the lua return values as an array.
//
/////////////////////////////////////////////////////////////////


void UlxLuaValues::Empty()
{
    Serialized.clear();
    Types.Empty();
    Data.Empty();
    Cursor = 0;
}

bool UlxLuaValues::Initialize(std::string_view data)
{
    Empty();
    Serialized = data;
    const char *SerializedChar = &Serialized[0];
    
    FlxStreamBuffer Decoder(Serialized);

    while (!Decoder.empty())
    {
        SimpleDynamicTag Tag = Decoder.read_simple_dynamic_tag();

        int64 Pos = Decoder.total_reads();
        ElxLuaValueType Type;
        switch (Tag)
        {
            case SimpleDynamicTag::BOOLEAN: Type=ElxLuaValueType::Boolean; Decoder.read_bool(); break;
            case SimpleDynamicTag::NUMBER:  Type=ElxLuaValueType::Float;   Decoder.read_double(); break;
            case SimpleDynamicTag::STRING:  Type=ElxLuaValueType::String;  Decoder.read_string_view(); break;
            case SimpleDynamicTag::TOKEN:   Type=ElxLuaValueType::Name;    Decoder.read_string_view(); break;
            case SimpleDynamicTag::VECTOR:  Type=ElxLuaValueType::Vector;  Decoder.read_fvector(); break;
            default: {
                Empty();
                return false;
            }
        }
        int64 Pos2 = Decoder.total_reads();
        Types.Add(Type);
        Data.Add(std::string_view(SerializedChar + Pos, Pos2 - Pos));
    }
    return true;
}

FString UlxLuaValues::DebugString() const
{
    TStringBuilder<2048> Output;
    Output << TEXT("{ ");
    for (int i = 0; i < Types.Num(); i++)
    {
        if (i > 0) Output << TEXT(", ");

        ElxLuaValueType Type = Types[i];
        FlxStreamBuffer Decoder(Data[i]);
        switch (Type)
        {
            case ElxLuaValueType::Boolean:
                Output << Decoder.read_bool();
                break;
            case ElxLuaValueType::Float:
                Output << FString::Printf(TEXT("%lf"), Decoder.read_double());
                break;
            case ElxLuaValueType::String:
                Output << TEXT("\"") << Decoder.read_fstring() << TEXT("\"");
                break;
            case ElxLuaValueType::Name:
                Output << Decoder.read_fname();
                break;
            case ElxLuaValueType::Vector: {
                FVector Vec = Decoder.read_fvector();
                Output << FString::Printf(TEXT("(%lf, %lf, %lf)"), Vec.X, Vec.Y, Vec.Z);
                break;
            }
        }
    }
    Output << TEXT(" }");
    return Output.ToString();
}

ElxSuccessOrError UlxLuaValues::CheckType(ElxLuaValueType Type, ElxLuaValueType Desired)
{
    if (Type != Desired)
    {
        FString TypeName = StaticEnum<ElxLuaValueType>()->GetDisplayNameTextByValue(int64(Type)).ToString();
        FString DesiredName = StaticEnum<ElxLuaValueType>()->GetDisplayNameTextByValue(int64(Desired)).ToString();
        UE_LOG(LogBlueprint, Error, TEXT("Expected a value of type %s, but found %s instead."), *DesiredName, *TypeName);
        return ElxSuccessOrError::Error;
    }
    return ElxSuccessOrError::Success;
}

ElxLuaValueType UlxLuaValues::NextType() const
{
    if (Cursor < 0) return ElxLuaValueType::None;
    if (Cursor >= Types.Num()) return ElxLuaValueType::None;
    return Types[Cursor];
}

void UlxLuaValues::ReadString(ElxSuccessOrError &Status, FString &Result)
{
    Status = CheckType(NextType(), ElxLuaValueType::String);
    if (Status == ElxSuccessOrError::Error)
    {
        Result.Empty(); return;
    }
    Result = FlxStreamBuffer(Data[Cursor++]).read_fstring();
}

void UlxLuaValues::ReadName(ElxSuccessOrError &Status, FName &Result)
{
    Status = CheckType(NextType(), ElxLuaValueType::Name);
    if (Status == ElxSuccessOrError::Error)
    {
        Result = FName(); return;
    }
    Result = FlxStreamBuffer(Data[Cursor++]).read_fname();
}

void UlxLuaValues::ReadFloat(ElxSuccessOrError &Status, double &Result)
{
    Status = CheckType(NextType(), ElxLuaValueType::Float);
    if (Status == ElxSuccessOrError::Error)
    {
        Result = 0.0; return;
    }
    Result = FlxStreamBuffer(Data[Cursor++]).read_double();
}

void UlxLuaValues::ReadInt(ElxSuccessOrError &Status, int &Result)
{
    Status = CheckType(NextType(), ElxLuaValueType::Float);
    if (Status == ElxSuccessOrError::Error)
    {
        Result = 0.0; return;
    }
    double dvalue = FlxStreamBuffer(Data[Cursor++]).read_double();
    Result = int(dvalue);
    if (double(Result) != dvalue)
    {
        Result = 0; Status = ElxSuccessOrError::Error; return;
    }
}

void UlxLuaValues::ReadVector(ElxSuccessOrError &Status, FVector &Result)
{
    Status = CheckType(NextType(), ElxLuaValueType::Vector);
    if (Status == ElxSuccessOrError::Error)
    {
        Result = FVector(); return;
    }
    Result = FlxStreamBuffer(Data[Cursor++]).read_fvector();
}

void UlxLuaValues::ReadVector2D(ElxSuccessOrError &Status, FVector2D &Result)
{
    Status = CheckType(NextType(), ElxLuaValueType::Vector);
    if (Status == ElxSuccessOrError::Error)
    {
        Result = FVector2D(); return;
    }
    FVector VValue = FlxStreamBuffer(Data[Cursor++]).read_fvector();
    Result = FVector2D(VValue.X, VValue.Y);
}

void UlxLuaValues::ReadBoolean(ElxSuccessOrError &Status, bool &Result)
{
    Status = CheckType(NextType(), ElxLuaValueType::Boolean);
    if (Status == ElxSuccessOrError::Error)
    {
        Result = false; return;
    }
    Result = FlxStreamBuffer(Data[Cursor++]).read_bool();
}