#include "imageloader.h"
#include <algorithm>
#include <cctype>
#include <stdexcept>
#include <fitsio.h>
#include <jpeglib.h>
#include <cstring>
#include <memory>

std::shared_ptr<ImageBuffer8bit> ImageLoader::loadImage(const std::string &filename)
{
    std::string ext = getFileExtension(filename);

    if (ext == "jpg" || ext == "jpeg")
    {
        return loadJpeg(filename);
    }
    else if (ext == "fit" || ext == "fits")
    {
        return loadFits(filename);
    }
    else
    {
        throw std::runtime_error("Unsupported file format: " + ext);
    }
}

void ImageLoader::saveImage(const std::string &filename, const std::shared_ptr<ImageBuffer8bit> &image)
{
    std::string ext = getFileExtension(filename);

    if (ext == "jpg" || ext == "jpeg")
    {
        saveJpeg(filename, image);
    }
    else if (ext == "fit" || ext == "fits")
    {
        saveFits(filename, image);
    }
    else
    {
        throw std::runtime_error("Unsupported file format: " + ext);
    }
}

std::shared_ptr<ImageBuffer8bit> ImageLoader::loadJpeg(const std::string &filename)
{
    FILE *file = fopen(filename.c_str(), "rb");
    if (!file)
    {
        throw std::runtime_error("Could not open file: " + filename);
    }

    // Initialize JPEG decompression
    jpeg_decompress_struct cinfo;
    jpeg_error_mgr jerr;

    cinfo.err = jpeg_std_error(&jerr);
    jpeg_create_decompress(&cinfo);
    jpeg_stdio_src(&cinfo, file);

    // Read JPEG header
    if (jpeg_read_header(&cinfo, TRUE) != JPEG_HEADER_OK)
    {
        fclose(file);
        jpeg_destroy_decompress(&cinfo);
        throw std::runtime_error("Invalid JPEG file: " + filename);
    }

    // Start decompression
    jpeg_start_decompress(&cinfo);

    // Create image buffer
    int width = cinfo.output_width;
    int height = cinfo.output_height;
    int channels = cinfo.output_components;

    auto buffer = std::make_shared<ImageBuffer8bit>(width, height, channels);

    // Read scanlines
    JSAMPARRAY scanlines = (*cinfo.mem->alloc_sarray)
                           ((j_common_ptr)&cinfo, JPOOL_IMAGE, width * channels, 1);

    while (cinfo.output_scanline < cinfo.output_height)
    {
        jpeg_read_scanlines(&cinfo, scanlines, 1);

        // Copy scanline to buffer
        int y = cinfo.output_scanline - 1;
        for (int x = 0; x < width; ++x)
        {
            for (int c = 0; c < channels; ++c)
            {
                buffer->setPixel(x, y, scanlines[0][x * channels + c], c);
            }
        }
    }

    // Finish decompression
    jpeg_finish_decompress(&cinfo);
    jpeg_destroy_decompress(&cinfo);
    fclose(file);

    return buffer;
}

std::shared_ptr<ImageBuffer8bit> ImageLoader::loadFits(const std::string &filename)
{
    fitsfile *fptr;
    int status = 0;

    // Open FITS file
    if (fits_open_file(&fptr, filename.c_str(), READONLY, &status))
    {
        throw std::runtime_error("Could not open FITS file: " + filename);
    }

    // Get image dimensions
    int bitpix, naxis;
    long naxes[2];

    if (fits_get_img_param(fptr, 2, &bitpix, &naxis, naxes, &status))
    {
        fits_close_file(fptr, &status);
        throw std::runtime_error("Could not read FITS image parameters");
    }

    if (naxis != 2)
    {
        fits_close_file(fptr, &status);
        throw std::runtime_error("Only 2D FITS images are supported");
    }

    int width = naxes[0];
    int height = naxes[1];

    // Create image buffer (grayscale)
    auto buffer = std::make_shared<ImageBuffer8bit>(width, height, 1);

    // Read image data
    long fpixel[2] = {1, 1};  // FITS uses 1-based indexing

    // Allocate temporary buffer for FITS data
    std::vector<uint8_t> tempBuffer(width * height);

    if (fits_read_pix(fptr, TBYTE, fpixel, width * height, nullptr, tempBuffer.data(), nullptr, &status))
    {
        fits_close_file(fptr, &status);
        throw std::runtime_error("Could not read FITS image data");
    }

    // Copy data to image buffer
    for (int y = 0; y < height; ++y)
    {
        for (int x = 0; x < width; ++x)
        {
            buffer->setPixel(x, y, tempBuffer[y * width + x]);
        }
    }

    // Close FITS file
    fits_close_file(fptr, &status);

    return buffer;
}

void ImageLoader::saveJpeg(const std::string &filename, const std::shared_ptr<ImageBuffer8bit> &image)
{
    FILE *file = fopen(filename.c_str(), "wb");
    if (!file)
    {
        throw std::runtime_error("Could not open file for writing: " + filename);
    }

    // Initialize JPEG compression
    jpeg_compress_struct cinfo;
    jpeg_error_mgr jerr;

    cinfo.err = jpeg_std_error(&jerr);
    jpeg_create_compress(&cinfo);
    jpeg_stdio_dest(&cinfo, file);

    // Set image parameters
    cinfo.image_width = image->getWidth();
    cinfo.image_height = image->getHeight();
    cinfo.input_components = image->getChannels();
    cinfo.in_color_space = (image->getChannels() == 3) ? JCS_RGB : JCS_GRAYSCALE;

    // Set default compression parameters
    jpeg_set_defaults(&cinfo);
    jpeg_set_quality(&cinfo, 95, TRUE);

    // Start compression
    jpeg_start_compress(&cinfo, TRUE);

    // Allocate scanline buffer
    int row_stride = image->getWidth() * image->getChannels();
    JSAMPARRAY scanlines = (*cinfo.mem->alloc_sarray)
                           ((j_common_ptr)&cinfo, JPOOL_IMAGE, row_stride, 1);

    // Write scanlines
    while (cinfo.next_scanline < cinfo.image_height)
    {
        int y = cinfo.next_scanline;

        // Copy buffer to scanline
        for (int x = 0; x < image->getWidth(); ++x)
        {
            for (int c = 0; c < image->getChannels(); ++c)
            {
                scanlines[0][x * image->getChannels() + c] = image->getPixel(x, y, c);
            }
        }

        jpeg_write_scanlines(&cinfo, scanlines, 1);
    }

    // Finish compression
    jpeg_finish_compress(&cinfo);
    jpeg_destroy_compress(&cinfo);
    fclose(file);
}

void ImageLoader::saveFits(const std::string &filename, const std::shared_ptr<ImageBuffer8bit> &image)
{
    fitsfile *fptr;
    int status = 0;

    // Create FITS file
    if (fits_create_file(&fptr, filename.c_str(), &status))
    {
        throw std::runtime_error("Could not create FITS file: " + filename);
    }

    // Set image dimensions
    long naxes[2] = {image->getWidth(), image->getHeight()};

    // Create image
    if (fits_create_img(fptr, BYTE_IMG, 2, naxes, &status))
    {
        fits_close_file(fptr, &status);
        throw std::runtime_error("Could not create FITS image");
    }

    // Allocate temporary buffer for FITS data
    std::vector<uint8_t> tempBuffer(image->getWidth() * image->getHeight());

    // Copy data to temporary buffer
    for (int y = 0; y < image->getHeight(); ++y)
    {
        for (int x = 0; x < image->getWidth(); ++x)
        {
            tempBuffer[y * image->getWidth() + x] = image->getPixel(x, y, 0);
        }
    }

    // Write image data
    long fpixel[2] = {1, 1};  // FITS uses 1-based indexing

    if (fits_write_pix(fptr, TBYTE, fpixel, image->getWidth() * image->getHeight(), tempBuffer.data(), &status))
    {
        fits_close_file(fptr, &status);
        throw std::runtime_error("Could not write FITS image data");
    }

    // Close FITS file
    fits_close_file(fptr, &status);
}

std::string ImageLoader::getFileExtension(const std::string &filename)
{
    size_t pos = filename.find_last_of('.');
    if (pos == std::string::npos)
    {
        return "";
    }

    std::string ext = filename.substr(pos + 1);
    std::transform(ext.begin(), ext.end(), ext.begin(),
                   [](unsigned char c)
    {
        return std::tolower(c);
    });

    return ext;
}
