#include <xavia/sensorFactory.h>
#include <xavia/iSensor.h>
#include <xavia/iCamera.h>
#include <xavia/iImage.h>
#include <xavia/error.h>
#include <xavia/logs.h>
#include <xavia/iPointCloud.h>
#include <xavia/pixelFormat.h>
#include <xavia/sensorState.h>

#include <iostream>
#include <fstream>
#include <chrono>
#include <thread>
#include <atomic>
#include <filesystem>

namespace xsdk = xavia::sdk;

std::atomic<int> pointcloudLineCount{0};
std::atomic<int> pointcloudFrameCount{0};
std::atomic<int> imageCount{0};
const std::string outputDir = "output";
const std::wstring outputDirW = L"output";
std::vector<std::unique_ptr<xsdk::IPointCloud>> pointClouds; // store point clouds for later processing

/* LIDAR AND CAMERA RECORDING EXAMPLE
 * 
 * This example demonstrates how to handle combine lidar and camera data.
 * Point clouds are saved to disk per frame in .xpc format, and camera images are saved as JPEGs.
 * The following steps are performed:
 * 1. Set up callbacks for point cloud and camera data.
 * 2. Build the sensor and get the camera interface.
 * 3. Start the sensor and camera.
 * 4. Process incoming point cloud lines and camera images in the callbacks and store them.
 * 5. Stop the sensor and camera after some time.
 * 
 * This example is naive in the sense that it assumes processing and disk speed are sufficient
 * to handle the storage in the callbacks. It uses the internal buffers to capture any variations.
 */

void SavePointCloud()
{
  if (pointClouds.empty())
    return;

  // use first line to get frame number for file naming.
  const auto& pointCloud = pointClouds.front();
  const auto& metadata = pointCloud->GetMetaData();

  std::ofstream outFile(outputDir + "/pointcloud_" + std::to_string(metadata.frameNumber) + ".xpc");
  if (!outFile)
  {
    std::wcerr << L"Failed to open file for writing point cloud data" << std::endl;
    return;
  }

  for(auto& pc : pointClouds)
  {
    auto bytes = pc->Serialize();
    outFile.write(reinterpret_cast<const char*>(bytes.data()), bytes.size());
  }
}

void pointCloudCallback(std::unique_ptr<xsdk::IPointCloud> pointCloud)
{
  // step 4a. Process incoming point cloud line
  ++pointcloudLineCount;
  auto md = pointCloud->GetMetaData();
  if(md.lineNumber == 0)
  {
    ++pointcloudFrameCount;
    std::wcout << L"PointCloud # " << pointcloudFrameCount << L": " << md.frameNumber <<std::endl;

    SavePointCloud();
    pointClouds.clear(); // clear stored point clouds after saving

  }
  pointClouds.push_back(std::move(pointCloud)); // aggregate the point cloud lines
}

void imageCallback(std::unique_ptr<xsdk::IImage> image)
{
  // step 4b. Process incoming camera image
  ++imageCount;
  const auto& metadata = image->GetMetaData();
  std::wcout << L"Image # " << imageCount << L": " << metadata.frameNumber << std::endl;

  image->SaveAsJpeg(outputDirW + L"/image_" + std::to_wstring(metadata.frameNumber) + L".jpg", 90, true);
}

void errorCallback(const xsdk::Exception& ex)
{
  std::wcerr << L"Error: " << ex.what() << std::endl;
}

int main(int argc, char** argv)
{
  try
  {
    std::wcout << L"=== Lidar and Camera Recording Example ===" << std::endl;

    // create output dir if not already exists
    std::filesystem::create_directory(outputDir);

    // logging is not mandatory, but we need to convert to std::wcout.
    xsdk::SetLogCallBack([](const xsdk::Severity& sev, const std::wstring& msg)
                       { std::wcout << xsdk::GetLogTypeString(sev) << L": " << msg << std::endl; });
    xsdk::SetLogSeverity(xsdk::Severity::XM_INFO);
    
    xsdk::SensorFactory factory;
    factory.SetErrorCallback(errorCallback); // optional
    
    // step 1. Set up callbacks
    factory.SetPointCloudCallback(pointCloudCallback);
    factory.SetImageCallback(imageCallback);
    
    // other configuration
    factory.SetImageOutputFormat(xsdk::PixelFormat::RGB8);

    // step 2. Build the sensor and get the camera
    auto sensor = factory.Build();
    auto camera = sensor->GetCamera();
    
    // always check if the camera is available before trying to use it.
    // The camera may not be available due to various reasons.
    if (!camera)
    {
      std::wcerr << L"Camera not available" << std::endl;
      return EXIT_FAILURE;
    }
    std::wcout << L"Connected to Camera ID: " << camera->GetCameraIdW() << std::endl;

    // step 3. Start the sensors
    if(sensor->GetCurrentState()!= xsdk::SensorState::RUNNING)
    {
      sensor->Start();
      camera->Start();
    }

    std::this_thread::sleep_for(std::chrono::seconds(5));

    // step 5. Stop the sensors
    sensor->Stop();
    camera->Stop();

    std::wcout << L"Total point cloud lines received: " << pointcloudLineCount << std::endl;
    std::wcout << L"Total point cloud frames received: " << pointcloudFrameCount << std::endl;
    std::wcout << L"Total images received: " << imageCount << std::endl;
  }
  catch (const std::exception& ex)
  {
    std::wcerr << L"Exception: " << ex.what() << std::endl;
    return EXIT_FAILURE;
  }

  std::wcout << L"=== Lidar and Camera Recording Example Completed ===" << std::endl;
  return EXIT_SUCCESS;
}
