/*++

Copyright (c) Microsoft. All rights reserved.

Module Name:

    WSLCCLIArgumentUnitTests.cpp

Abstract:

    This file contains unit tests for WSLC CLI argument parsing and validation.

--*/

#include "precomp.h"
#include "windows/Common.h"
#include "WSLCCLITestHelpers.h"

#include "Argument.h"
#include "ArgumentTypes.h"
#include "ArgumentValidation.h"
#include "Exceptions.h"
#include <wslc.h>

using namespace wsl::windows::wslc;
using namespace wsl::windows::wslc::argument;

using namespace WSLCTestHelpers;
using namespace WEX::Logging;
using namespace WEX::Common;
using namespace WEX::TestExecution;

namespace WSLCCLIArgumentUnitTests {
class WSLCCLIArgumentUnitTests
{
    WSLC_TEST_CLASS(WSLCCLIArgumentUnitTests)

    TEST_CLASS_SETUP(TestClassSetup)
    {
        // Add any necessary setup for argument tests
        return true;
    }

    TEST_CLASS_CLEANUP(TestClassCleanup)
    {
        // Add any necessary cleanup for argument tests
        return true;
    }

    // Test: Verify Argument::Create() successfully creates arguments for all ArgType enum values
    TEST_METHOD(ArgumentCreate_AllArguments)
    {
        // ArgMap is the container for processed args.
        ArgMap args;

        // Iterate through all ArgType enum values except Max
        auto allArgTypes = std::vector<ArgType>{};
        for (int i = 0; i < static_cast<int>(ArgType::Max); ++i)
        {
            ArgType argType = static_cast<ArgType>(i);

            // Create argument using Create
            Argument arg = Argument::Create(argType);

            // Verify the argument was created successfully by checking its type matches
            VERIFY_ARE_EQUAL(static_cast<int>(arg.Type()), i);

            // Verify the argument has basic properties set
            // (Name should not be empty for valid argument types)
            VERIFY_IS_FALSE(arg.Name().empty());
            LogComment(L"Verified Argument::Create() creates argument with name: " + arg.Name());

            // Add the argument to the ArgMap with a test value based on its type.
            VERIFY_IS_FALSE(args.Contains(argType));
            switch (arg.Kind())
            {
            case Kind::Value:
            case Kind::Positional:
                args.Add(argType, std::wstring(L"test"));
                break;
            case Kind::Forward:
                args.Add(argType, std::vector<std::wstring>{L"forward1", L"forward2"});
                break;
            case Kind::Flag:
                args.Add(argType, true);
                break;
            default:
                VERIFY_FAIL(L"Unhandled ValueType in test");
            }

            allArgTypes.push_back(argType);
            VERIFY_IS_TRUE(args.Contains(argType));
        }

        // We do not have a runtime Get for argument values, so we will instead use the keys
        // in the argmap. The fact that the keys exist and can be used to retrieve values
        // verifies that Argument::Create() created arguments that are compatible with ArgMap.
        // Verify all created argument types are in the ArgMap keys
        auto argMapKeys = args.GetKeys();
        VERIFY_ARE_EQUAL(argMapKeys.size(), allArgTypes.size());
        for (const auto& argType : allArgTypes)
        {
            VERIFY_IS_TRUE(std::find(argMapKeys.begin(), argMapKeys.end(), argType) != argMapKeys.end());
        }
    }

    // Test: Verify Argument::Create() successfully creates arguments for all ArgType enum values
    TEST_METHOD(ArgumentValidation_ValueValidation)
    {
        // Verify integer conversion for supported types.
        auto longlong = validation::GetIntegerFromString<LONGLONG>(L"1234567890123");
        VERIFY_ARE_EQUAL(longlong, 1234567890123LL);
        VERIFY_THROWS(validation::GetIntegerFromString<LONGLONG>(L"abc"), ArgumentException);                      // Not a number
        VERIFY_THROWS(validation::GetIntegerFromString<LONGLONG>(L"-92233720369999854775808"), ArgumentException); // Out of range
        VERIFY_NO_THROW(validation::ValidateIntegerFromString<LONGLONG>({L"1234", L"-1234567890123"}, L"testArg"));
        VERIFY_THROWS(validation::ValidateIntegerFromString<LONGLONG>({L"1234", L"-92233720369999854775808"}, L"testArg"), ArgumentException);

        // Verify --tail validation rejects 0 (mirrors ArgType::Tail validation)
        VERIFY_THROWS(validation::ValidateIntegerFromString<ULONGLONG>({L"0"}, L"tail", [](auto value) { return value != 0; }), ArgumentException);
        VERIFY_NO_THROW(validation::ValidateIntegerFromString<ULONGLONG>({L"10"}, L"tail", [](auto value) { return value != 0; }));
        VERIFY_NO_THROW(validation::ValidateIntegerFromString<ULONGLONG>({L"1"}, L"tail", [](auto value) { return value != 0; }));

        // Verify WSLCSignal conversion
        auto validSignal = validation::GetWSLCSignalFromString(L"SIGTERM");
        VERIFY_ARE_EQUAL(validSignal, WSLCSignalSIGTERM);
        validSignal = validation::GetWSLCSignalFromString(L"TERM"); // No prefix
        VERIFY_ARE_EQUAL(validSignal, WSLCSignalSIGTERM);
        validSignal = validation::GetWSLCSignalFromString(L"sIgTerm"); // Case-insensitive
        VERIFY_ARE_EQUAL(validSignal, WSLCSignalSIGTERM);
        validSignal = validation::GetWSLCSignalFromString(L"term"); // Case-insensitive no prefix
        VERIFY_ARE_EQUAL(validSignal, WSLCSignalSIGTERM);
        VERIFY_THROWS(validation::GetWSLCSignalFromString(L"INVALID_SIGNAL"), ArgumentException);
        validSignal = validation::GetWSLCSignalFromString(L"15"); // SIGTERM is 15
        VERIFY_ARE_EQUAL(validSignal, WSLCSignalSIGTERM);
        VERIFY_THROWS(validation::GetWSLCSignalFromString(L"999"), ArgumentException); // Out of range
        VERIFY_NO_THROW(validation::ValidateWSLCSignalFromString({L"HUP", L"9", L"SIGKILL", L"stop"}, L"signalArg"));
        VERIFY_THROWS(validation::ValidateWSLCSignalFromString({L"SIGHUP", L"999"}, L"signalArg"), ArgumentException); // 999 is out of range

        // Verify format type
        auto format = validation::GetFormatTypeFromString(L"json");
        VERIFY_ARE_EQUAL(format, FormatType::Json);
        format = validation::GetFormatTypeFromString(L"table");
        VERIFY_ARE_EQUAL(format, FormatType::Table);
        VERIFY_THROWS(validation::GetFormatTypeFromString(L"xml"), ArgumentException);
        VERIFY_NO_THROW(validation::ValidateFormatTypeFromString({L"json", L"table"}, L"formatArg"));
        VERIFY_THROWS(validation::ValidateFormatTypeFromString({L"JSON", L"TABLE", L"csv"}, L"formatArg"), ArgumentException);

        // Verify GPU device argument
        VERIFY_NO_THROW(validation::ValidateGpus({L"all"}, L"gpusArg"));
        VERIFY_THROWS(validation::ValidateGpus({L"none"}, L"gpusArg"), ArgumentException);
        VERIFY_THROWS(validation::ValidateGpus({L"0"}, L"gpusArg"), ArgumentException);
        VERIFY_THROWS(validation::ValidateGpus({L"gpu0"}, L"gpusArg"), ArgumentException);
        VERIFY_THROWS(validation::ValidateGpus({L""}, L"gpusArg"), ArgumentException);
    }

    // Test: Verify EnumVariantMap behavior with ArgTypes.
    TEST_METHOD(EnumVariantMap_AllDataTypes)
    {
        // ArgMap is an EnumVariantMap
        ArgMap argsContainer;

        // Verify basic add
        argsContainer.Add<ArgType::Help>(true);
        VERIFY_IS_TRUE(argsContainer.Contains(ArgType::Help));
        argsContainer.Add<ArgType::ContainerId>(std::wstring(L"test"));
        VERIFY_IS_TRUE(argsContainer.Contains(ArgType::ContainerId));
        argsContainer.Add<ArgType::ForwardArgs>(std::vector<std::wstring>{L"test1", L"test2"});
        VERIFY_IS_TRUE(argsContainer.Contains(ArgType::ForwardArgs));

        // Verify basic retrieval
        auto retrievedBool = argsContainer.Get<ArgType::Help>();
        VERIFY_ARE_EQUAL(retrievedBool, true);
        auto retrievedString = argsContainer.Get<ArgType::ContainerId>();
        VERIFY_ARE_EQUAL(retrievedString, std::wstring(L"test"));
        auto retrievedStringSet = argsContainer.Get<ArgType::ForwardArgs>();
        VERIFY_ARE_EQUAL(retrievedStringSet[0], std::wstring(L"test1"));
        VERIFY_ARE_EQUAL(retrievedStringSet[1], std::wstring(L"test2"));

        // Verify multimap functionality and Runtime Add
        argsContainer.Add(ArgType::Publish, std::wstring(L"test1"));
        argsContainer.Add(ArgType::Publish, std::wstring(L"test2"));
        argsContainer.Add(ArgType::Publish, std::wstring(L"test3"));
        VERIFY_ARE_EQUAL(argsContainer.Count(ArgType::Publish), 3);
        auto publishArgs = argsContainer.GetAll<ArgType::Publish>();
        VERIFY_ARE_EQUAL(publishArgs.size(), 3);
        VERIFY_ARE_EQUAL(publishArgs[0], std::wstring(L"test1"));
        VERIFY_ARE_EQUAL(publishArgs[1], std::wstring(L"test2"));
        VERIFY_ARE_EQUAL(publishArgs[2], std::wstring(L"test3"));

        // Verify Remove
        argsContainer.Remove(ArgType::Publish);
        VERIFY_ARE_EQUAL(argsContainer.Count(ArgType::Publish), 0);

        // Verify compile time add works like runtime add for multimap types.
        argsContainer.Add<ArgType::Publish>(L"test1");
        argsContainer.Add<ArgType::Publish>(L"test2");
        argsContainer.Add<ArgType::Publish>(L"test3");
        VERIFY_ARE_EQUAL(argsContainer.Count(ArgType::Publish), 3);
        publishArgs = argsContainer.GetAll<ArgType::Publish>();
        VERIFY_ARE_EQUAL(publishArgs.size(), 3);
        VERIFY_ARE_EQUAL(publishArgs[0], std::wstring(L"test1"));
        VERIFY_ARE_EQUAL(publishArgs[1], std::wstring(L"test2"));
        VERIFY_ARE_EQUAL(publishArgs[2], std::wstring(L"test3"));

        // Verify Keys
        auto allArgTypes = argsContainer.GetKeys();
        VERIFY_ARE_EQUAL(allArgTypes.size(), 4);
        VERIFY_IS_TRUE(std::find(allArgTypes.begin(), allArgTypes.end(), ArgType::Help) != allArgTypes.end());
        VERIFY_IS_TRUE(std::find(allArgTypes.begin(), allArgTypes.end(), ArgType::ContainerId) != allArgTypes.end());
        VERIFY_IS_TRUE(std::find(allArgTypes.begin(), allArgTypes.end(), ArgType::Publish) != allArgTypes.end());
        VERIFY_IS_TRUE(std::find(allArgTypes.begin(), allArgTypes.end(), ArgType::ForwardArgs) != allArgTypes.end());

        // Verify count
        VERIFY_ARE_EQUAL(argsContainer.Count(ArgType::Help), 1);
        VERIFY_ARE_EQUAL(argsContainer.Count(ArgType::ContainerId), 1);
        VERIFY_ARE_EQUAL(argsContainer.Count(ArgType::Publish), 3);
        VERIFY_ARE_EQUAL(argsContainer.Count(ArgType::ForwardArgs), 1);
        VERIFY_ARE_EQUAL(argsContainer.GetCount(), 6); // 1 Help + 1 ContainerId + 3 Publish + 1 ForwardArgs
        argsContainer.Remove(ArgType::Help);
        argsContainer.Remove(ArgType::ContainerId);
        argsContainer.Remove(ArgType::Publish);
        argsContainer.Remove(ArgType::ForwardArgs);
        VERIFY_ARE_EQUAL(argsContainer.GetCount(), 0);
    }
};
} // namespace WSLCCLIArgumentUnitTests