#include <Renderer.h>
// #include <common.hpp>
#include <chrono>
#include <filesystem>

void Renderer::initRenderer(int width, int height, int maxSamples, int samplesPerFrame) {
    this->width = width;
    this->height = height;
    this->maxSamples = maxSamples;
    this->samplesPerFrame = samplesPerFrame;

    int device_count;
    cudaGetDeviceCount(&device_count);
    if (device_count > 0) {
        cudaSetDevice(0);
    } else {
        throw std::runtime_error("No Cuda Device Found");
    }

    CUDA::free(deviceData);
    CUDA::malloc(deviceData, this->width * this->height);
}

Renderer::~Renderer() {
    CUDA::free(deviceData);
}

void Renderer::drawGui(float deltaTime, float frametime) {
    frameTimes[frameCounter] = deltaTime * 1000.f;
    frameCounter = (frameCounter + 1) % sizeAvgFrames;
    float timeMs = 0;
    for (int i = 0; i < sizeAvgFrames; ++i) {
        timeMs += frameTimes[i] / static_cast<float>(sizeAvgFrames);
    }

    float sampleTime = frametime / static_cast<float>(samplesPerFrame);

    const std::string frameRateString = "Frametime Window (ms): " + std::to_string(timeMs);
    const std::string FPSString = "FPS: " + std::to_string(1000.f / timeMs);
    const std::string frametimeString = "Frametime Cuda (ms): " + std::to_string(frametime);
    const std::string sampletimeString = "Frametime per Sample (ms): " + std::to_string(sampleTime);
    const std::string samplesString = "Num Samples: " + std::to_string(currSamples);

    ImVec2 center = ImGui::GetMainViewport()->GetCenter() + ImVec2(width / 2.f - 310.f, -height / 2.f + 10.f);
    ImGui::SetNextWindowPos(center);
    ImGui::SetNextWindowSize(ImVec2(300, 100));
    ImGui::PushStyleColor(ImGuiCol_WindowBg, ImVec4(0.8f, 0.8f, 0.8f, 0.8f));
    ImGui::PushStyleColor(ImGuiCol_Text, ImVec4(0.0f, 0.0f, 0.0f, 1.0f));
    ImGui::Begin("Performance", 0,
                 ImGuiWindowFlags_NoMove | ImGuiWindowFlags_NoResize | ImGuiWindowFlags_NoTitleBar |
                         ImGuiWindowFlags_NoCollapse);
    ImGui::Text("%s\n", FPSString.c_str());
    ImGui::Text("%s\n", frametimeString.c_str());
    ImGui::Text("%s\n", samplesString.c_str());
    ImGui::PushItemWidth(-120);
    ImGui::InputInt("Samples per Frame", &samplesPerFrame, 1);
    ImGui::PopItemWidth();
    ImGui::End();
    ImGui::PopStyleColor(2);
}

void Renderer::saveImage(std::string filepath) const {
    float4* hostData = new float4[width * height * 4];
    CUDA::memcopyToHost(hostData, deviceData, width * height);
    auto* imageData = new uchar3[width * height];
    for (int i = 0; i < width * height; ++i) {
        float4 pixel = hostData[i];
        uchar3 color;
        color.x = (unsigned char) (clamp(pixel.x, 0.f, 1.f) * 255.f);
        color.y = (unsigned char) (clamp(pixel.y, 0.f, 1.f) * 255.f);
        color.z = (unsigned char) (clamp(pixel.z, 0.f, 1.f) * 255.f);
        imageData[i] = color;
    }
    stbi_flip_vertically_on_write(true);

    if (stbi_write_png(filepath.c_str(), width, height, 3, imageData, width * sizeof(uchar3)) != 1) {
        std::cerr << "Couldn't write PNG File" << std::endl;
    }
}

void Renderer::saveImage() const {
    std::time_t result = std::time(nullptr);
    auto time = std::localtime(&result);
    std::filesystem::create_directories("../images/");
    const std::string filepath = "../images/Test_Image_" + std::to_string(time->tm_mday) + "_" +
                                 std::to_string(time->tm_mon + 1) + "_" + std::to_string(time->tm_hour) + "-" +
                                 std::to_string(time->tm_min) + "-" + std::to_string(time->tm_sec) + ".png";

    saveImage(filepath);
}
