#include "DebugService.h"
#include "DebugCore.h"
#include "ObjectRegistryDatabase.h"
#include "StandardParser.h"
#include "StreamString.h"
#include "BasicUDPSocket.h"
#include "BasicTCPSocket.h"
#include "RealTimeApplication.h"
#include "GlobalObjectsDatabase.h"
#include "MessageI.h"
#include <assert.h>
#include <stdio.h>

using namespace MARTe;

const char8 * const config_text = 
"+DebugService = {"
"    Class = DebugService "
"    ControlPort = 8080 "
"    UdpPort = 8081 "
"    StreamIP = \"127.0.0.1\" "
"}"
"+App = {"
"    Class = RealTimeApplication "
"    +Functions = {"
"        Class = ReferenceContainer "
"        +GAM1 = {"
"            Class = IOGAM "
"            InputSignals = {"
"                Counter = {"
"                    DataSource = Timer "
"                    Type = uint32 "
"                }"
"            }"
"            OutputSignals = {"
"                Counter = {"
"                    DataSource = DDB "
"                    Type = uint32 "
"                }"
"            }"
"        }"
"    }"
"    +Data = {"
"        Class = ReferenceContainer "
"        DefaultDataSource = DDB "
"        +Timer = {"
"            Class = LinuxTimer "
"            SleepTime = 100000 " // 100ms
"            Signals = {"
"                Counter = { Type = uint32 }"
"            }"
"        }"
"        +DDB = {"
"            Class = GAMDataSource "
"            Signals = { Counter = { Type = uint32 } }"
"        }"
"        +DAMS = { Class = TimingDataSource }"
"    }"
"    +States = {"
"        Class = ReferenceContainer "
"        +State1 = {"
"            Class = RealTimeState "
"            +Threads = {"
"                Class = ReferenceContainer "
"                +Thread1 = {"
"                    Class = RealTimeThread "
"                    Functions = {GAM1} "
"                }"
"            }"
"        }"
"    }"
"    +Scheduler = {"
"        Class = FastScheduler "
"        TimingDataSource = DAMS "
"    }"
"}";

void TestSchedulerControl() {
    printf("--- MARTe2 Scheduler Control Test ---\n");

    ConfigurationDatabase cdb;
    StreamString ss = config_text;
    ss.Seek(0);
    StandardParser parser(ss, cdb);
    assert(parser.Parse());
    assert(ObjectRegistryDatabase::Instance()->Initialise(cdb));

    ReferenceT<DebugService> service = ObjectRegistryDatabase::Instance()->Find("DebugService");
    assert(service.IsValid());
    
    ReferenceT<RealTimeApplication> app = ObjectRegistryDatabase::Instance()->Find("App");
    assert(app.IsValid());
    
    if (app->PrepareNextState("State1") != ErrorManagement::NoError) {
        printf("ERROR: Failed to prepare State1\n");
        return;
    }
    
    if (app->StartNextStateExecution() != ErrorManagement::NoError) {
        printf("ERROR: Failed to start execution\n");
        return;
    }

    printf("Application started. Waiting for cycles...\n");
    Sleep::MSec(1000); 

    // Enable Trace First
    {
        BasicTCPSocket tClient;
        if (tClient.Connect("127.0.0.1", 8080)) {
            const char* cmd = "TRACE Root.App.Data.Timer.Counter 1\n";
            uint32 s = StringHelper::Length(cmd);
            tClient.Write(cmd, s);
            tClient.Close();
        } else {
            printf("WARNING: Could not connect to DebugService to enable trace.\n");
        }
    }

    BasicUDPSocket listener;
    listener.Open();
    listener.Listen(8081);
    
    // Read current value
    uint32 valBeforePause = 0;
    char buffer[2048];
    uint32 size = 2048;
    TimeoutType timeout(500); 
    if (listener.Read(buffer, size, timeout)) {
        // [Header][ID][Size][Value]
        valBeforePause = *(uint32*)(&buffer[28]);
        printf("Value before/at pause: %u\n", valBeforePause);
    } else {
        printf("WARNING: No data received before pause.\n");
    }

    // Send PAUSE
    printf("Sending PAUSE command...\n");
    BasicTCPSocket client;
    if (client.Connect("127.0.0.1", 8080)) {
        const char* cmd = "PAUSE\n";
        uint32 s = StringHelper::Length(cmd);
        client.Write(cmd, s);
        client.Close();
    } else {
        printf("ERROR: Could not connect to DebugService to send PAUSE.\n");
    }

    Sleep::MSec(2000); // Wait 2 seconds

    // Read again - should be same or very close if paused
    uint32 valAfterWait = 0;
    size = 2048; // Reset size
    while(listener.Read(buffer, size, TimeoutType(10))) {
        valAfterWait = *(uint32*)(&buffer[28]);
        size = 2048;
    }
    
    printf("Value after 2s wait (drained): %u\n", valAfterWait);
    
    // Check if truly paused
    if (valAfterWait > valBeforePause + 5) {
        printf("FAILURE: Counter increased significantly while paused! (%u -> %u)\n", valBeforePause, valAfterWait);
    } else {
        printf("SUCCESS: Counter held steady (or close) during pause.\n");
    }

    // Resume
    printf("Sending RESUME command...\n");
    {
        BasicTCPSocket rClient;
        if (rClient.Connect("127.0.0.1", 8080)) {
            const char* cmd = "RESUME\n";
            uint32 s = StringHelper::Length(cmd);
            rClient.Write(cmd, s);
            rClient.Close();
        }
    }
    
    Sleep::MSec(1000);
    
    // Check if increasing
    uint32 valAfterResume = 0;
    size = 2048;
    if (listener.Read(buffer, size, timeout)) {
        valAfterResume = *(uint32*)(&buffer[28]);
        printf("Value after resume: %u\n", valAfterResume);
    }
    
    if (valAfterResume > valAfterWait) {
        printf("SUCCESS: Execution resumed.\n");
    } else {
        printf("FAILURE: Execution did not resume.\n");
    }

    app->StopCurrentStateExecution();
}

int main() {
    TestSchedulerControl();
    return 0;
}