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

using namespace MARTe;

const char8 * const scheduler_config_text = 
"DebugService = {"
"    Class = DebugService "
"    ControlPort = 8098 "
"    UdpPort = 8099 "
"    StreamIP = \"127.0.0.1\" "
"}"
"App = {"
"    Class = RealTimeApplication "
"    +Functions = {"
"        Class = ReferenceContainer "
"        +GAM1 = {"
"            Class = IOGAM "
"            InputSignals = {"
"                Counter = { DataSource = Timer Type = uint32 Frequency = 1000 }"
"                Time = { DataSource = Timer Type = uint32 }"
"            }"
"            OutputSignals = {"
"                Counter = { DataSource = DDB Type = uint32 }"
"                Time = { DataSource = DDB Type = uint32 }"
"            }"
"        }"
"    }"
"    +Data = {"
"        Class = ReferenceContainer "
"        DefaultDataSource = DDB "
"        +Timer = { Class = LinuxTimer SleepTime = 1000 Signals = { Counter = { Type = uint32 } Time = { Type = uint32 } } }"
"        +DDB = { Class = GAMDataSource Signals = { Counter = { Type = uint32 } Time = { Type = uint32 } } }"
"        +DAMS = { Class = TimingDataSource }"
"    }"
"    +States = {"
"        Class = ReferenceContainer "
"        +State1 = { Class = RealTimeState +Threads = { Class = ReferenceContainer +Thread1 = { Class = RealTimeThread Functions = {GAM1} } } }"
"    }"
"    +Scheduler = { Class = GAMScheduler TimingDataSource = DAMS }"
"}";

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

  ObjectRegistryDatabase::Instance()->Purge();

  ConfigurationDatabase cdb;
  StreamString ss = scheduler_config_text;
  ss.Seek(0);
  StandardParser parser(ss, cdb);
  if (!parser.Parse()) {
      printf("ERROR: Failed to parse configuration\n");
      return;
  }
  
  cdb.MoveToRoot();
  uint32 n = cdb.GetNumberOfChildren();
  for (uint32 i=0; i<n; i++) {
      const char8* name = cdb.GetChildName(i);
      ConfigurationDatabase child;
      cdb.MoveRelative(name);
      cdb.Copy(child);
      cdb.MoveToAncestor(1u);
      
      StreamString className;
      child.Read("Class", className);
      
      Reference ref(className.Buffer(), GlobalObjectsDatabase::Instance()->GetStandardHeap());
      if (!ref.IsValid()) {
          printf("ERROR: Could not create object %s of class %s\n", name, className.Buffer());
          continue;
      }
      ref->SetName(name);
      if (!ref->Initialise(child)) {
          printf("ERROR: Failed to initialise object %s\n", name);
          continue;
      }
      ObjectRegistryDatabase::Instance()->Insert(ref);
  }

  ReferenceT<DebugService> service =
      ObjectRegistryDatabase::Instance()->Find("DebugService");
  if (!service.IsValid()) {
      printf("ERROR: DebugService not found in registry\n");
      return;
  }
  service->SetFullConfig(cdb);

  ReferenceT<RealTimeApplication> app =
      ObjectRegistryDatabase::Instance()->Find("App");
  if (!app.IsValid()) {
      printf("ERROR: App not found in registry\n");
      return;
  }

  if (!app->ConfigureApplication()) {
      printf("ERROR: ConfigureApplication failed.\n");
      return;
  }

  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(2000);

  // Enable Trace First - with retry logic
  {
    bool connected = false;
    for (int retry=0; retry<10 && !connected; retry++) {
        BasicTCPSocket tClient;
        if (tClient.Open()) {
            if (tClient.Connect("127.0.0.1", 8098)) {
                connected = true;
                const char *cmd = "TRACE App.Data.Timer.Counter 1\n";
                uint32 s = StringHelper::Length(cmd);
                tClient.Write(cmd, s);
                tClient.Close();
            } else {
                printf("[SchedulerTest] Connect failed (retry %d)\n", retry);
                Sleep::MSec(500);
            }
        } else {
            printf("[SchedulerTest] Open failed (retry %d)\n", retry);
            Sleep::MSec(500);
        }
    }

    if (!connected) {
      printf("WARNING: Could not connect to DebugService to enable trace.\n");
    }
  }

  BasicUDPSocket listener;
  listener.Open();
  listener.Listen(8099);

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

  // Send PAUSE
  printf("Sending PAUSE command...\n");
  {
      bool connected = false;
      for (int retry=0; retry<10 && !connected; retry++) {
          BasicTCPSocket client;
          if (client.Open()) {
              if (client.Connect("127.0.0.1", 8098)) {
                  connected = true;
                  const char *cmd = "PAUSE\n";
                  uint32 s = StringHelper::Length(cmd);
                  client.Write(cmd, s);
                  client.Close();
              } else {
                  Sleep::MSec(200);
              }
          } else {
              Sleep::MSec(200);
          }
      }

      if (!connected) {
        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(100))) {
    valAfterWait = *(uint32 *)(&buffer[sizeof(TraceHeader) + 16]);
    size = 2048;
  }

  printf("Value after 2s wait (drained): %u\n", valAfterWait);

  // Check if truly paused
  if (valAfterWait > valBeforePause + 10) {
    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");
  {
    bool connected = false;
    for (int retry=0; retry<10 && !connected; retry++) {
        BasicTCPSocket rClient;
        if (rClient.Open()) {
            if (rClient.Connect("127.0.0.1", 8098)) {
                connected = true;
                const char *cmd = "RESUME\n";
                uint32 s = StringHelper::Length(cmd);
                rClient.Write(cmd, s);
                rClient.Close();
            } else {
                Sleep::MSec(200);
            }
        } else {
            Sleep::MSec(200);
        }
    }
  }

  Sleep::MSec(1000);

  // Check if increasing
  uint32 valAfterResume = 0;
  size = 2048;
  if (listener.Read(buffer, size, timeout)) {
    valAfterResume = *(uint32 *)(&buffer[sizeof(TraceHeader) + 16]);
    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();
}