#ifndef STACKER_SATURATION_ADJUSTMENT_TPP
#define STACKER_SATURATION_ADJUSTMENT_TPP

#include "../io/timer.h"
#include <opencv2/opencv.hpp>

template <typename T>
std::shared_ptr<ImageBuffer<T>> Stacker<T>::applySaturationAdjustment(const std::shared_ptr<ImageBuffer<T>> &frame,
                             float percentage)
{
    TIME_FUNCTION("applySaturationAdjustment");

    // Create a copy of the frame
    auto result = frame->clone();

    // Only applicable to color images
    if (result->getChannels() < 3)
    {
        return result;
    }

    // For astronomical images, we'll use HSV color space for more natural results
    // and apply selective saturation based on luminance

    T maxValue = std::is_same<T, uint8_t>::value ? 255 : 65535;
    float factor = percentage / 100.0f;

    // Convert to OpenCV Mat for HSV conversion
    cv::Mat imageMat;
    if (std::is_same<T, uint8_t>::value)
    {
        imageMat = cv::Mat(result->getHeight(), result->getWidth(), CV_8UC3, result->getData());
    }
    else
    {
        // For 16-bit images, we need to convert to 8-bit for HSV conversion
        // OpenCV doesn't support 16-bit HSV conversion directly
        cv::Mat tempMat(result->getHeight(), result->getWidth(), CV_16UC3, result->getData());
        tempMat.convertTo(imageMat, CV_8UC3, 255.0 / 65535.0);
    }

    // Create a deep copy to avoid memory issues
    imageMat = imageMat.clone();

    // Convert to HSV color space
    cv::Mat hsvMat;
    cv::cvtColor(imageMat, hsvMat, cv::COLOR_BGR2HSV);

    // Split into channels
    std::vector<cv::Mat> hsvChannels;
    cv::split(hsvMat, hsvChannels);

    // Saturation is the second channel (index 1)
    cv::Mat &saturationChannel = hsvChannels[1];
    cv::Mat &valueChannel = hsvChannels[2]; // Value (brightness) channel

    // Apply saturation adjustment with protection for bright areas
    for (int y = 0; y < saturationChannel.rows; ++y)
    {
        for (int x = 0; x < saturationChannel.cols; ++x)
        {
            // Get current saturation and value
            uchar &saturation = saturationChannel.at<uchar>(y, x);
            uchar value = valueChannel.at<uchar>(y, x);

            // Calculate adjustment factor based on brightness
            // Reduce saturation for very bright pixels (stars)
            float brightnessFactor = 1.0f;
            if (value > 200) // Very bright pixels (stars)
            {
                // Reduce saturation enhancement for bright areas
                brightnessFactor = std::max(0.2f, 1.0f - (value - 200) / 55.0f);
            }

            // Enhance saturation more in mid-tones (nebulosity)
            float midtoneFactor = 1.0f;
            if (value > 50 && value < 180)
            {
                // Boost saturation in midtones (peak at around 120)
                midtoneFactor = 1.0f + 0.3f * (1.0f - std::abs(value - 120) / 70.0f);
            }

            // Apply combined adjustment
            float adjustmentFactor = factor * brightnessFactor * midtoneFactor;

            // Calculate new saturation
            float newSaturation = saturation * adjustmentFactor;

            // Clamp to valid range
            saturation = static_cast<uchar>(std::max(0.0f, std::min(255.0f, newSaturation)));
        }
    }

    // Merge channels back
    cv::merge(hsvChannels, hsvMat);

    // Convert back to BGR
    cv::cvtColor(hsvMat, imageMat, cv::COLOR_HSV2BGR);

    // If original was 16-bit, convert back to 16-bit
    if (!std::is_same<T, uint8_t>::value)
    {
        cv::Mat outputMat;
        imageMat.convertTo(outputMat, CV_16UC3, 65535.0 / 255.0);

        // Copy data back to ImageBuffer
        std::memcpy(result->getData(), outputMat.data,
                    result->getWidth() * result->getHeight() * result->getChannels() * sizeof(T));
    }
    else
    {
        // Copy data back to ImageBuffer
        std::memcpy(result->getData(), imageMat.data,
                    result->getWidth() * result->getHeight() * result->getChannels() * sizeof(T));
    }

    return result;
}

#endif // STACKER_SATURATION_ADJUSTMENT_TPP
