#include <stdio.h>
#include <assert.h>

#include "pslib.h"
#include "pmFPA.h"
#include "psAdditionals.h"

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// File-static (private) functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// Move to the appropriate extension in FITS file for HDU
static bool hduMove(pmHDU *hdu,         // HDU with extname
                    psFits *fits        // FITS file in which to move
                   )
{
    // Deal with the PHU case
    if (strcasecmp(hdu->extname, "PHU") == 0 || hdu->phu) {
        if (! psFitsMoveExtNum(fits, 0, false)) {
            psError(PS_ERR_IO, false, "Unable to move to primary header!\n");
            return false;
        }
        hdu->phu = true;
        return true;
    }

    if (! psFitsMoveExtName(fits, hdu->extname)) {
        psError(PS_ERR_IO, false, "Unable to move to extension %s\n", hdu->extname);
        return false;
    }
    // Now, just in case for some reason the PHU has an extension name that we've moved to....
    if (psFitsGetExtNum(fits) == 0) {
        hdu->phu = true;
    } else {
        hdu->phu = false;
    }

    return true;
}

static void hduFree(pmHDU *hdu)
{
    psFree(hdu->extname);
    psFree(hdu->format);
    psFree(hdu->header);
    psFree(hdu->images);
    psFree(hdu->table);
    psFree(hdu->weights);
    psFree(hdu->masks);
}


//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Public functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

pmHDU *pmHDUAlloc(const char *extname   // Extension name
                 )
{
    pmHDU *hdu = psAlloc(sizeof(pmHDU));
    psMemSetDeallocator(hdu, (psFreeFunc)hduFree);

    if (! extname || strlen(extname) == 0) {
        hdu->phu = true;
        hdu->extname = psStringCopy("PHU");
    } else {
        if (strcasecmp(extname, "PHU") == 0) {
            hdu->phu = true;
        } else {
            hdu->phu = false;
        }
        hdu->extname = psStringCopy(extname);
    }
    hdu->format  = NULL;
    hdu->header  = NULL;
    hdu->images  = NULL;
    hdu->table   = NULL;
    hdu->weights = NULL;
    hdu->masks   = NULL;

    return hdu;
}

bool pmHDUReadHeader(pmHDU *hdu,        // HDU for which to read header
                     psFits *fits       // FITS file from which to read
                    )
{
    assert(hdu);
    assert(fits);

    // Move to the appropriate extension
    psTrace(__func__, 5, "Moving to extension %s...\n", hdu->extname);
    if (!hduMove(hdu, fits)) {
        return false;
    }

    if (!hdu->header) {
        psTrace(__func__, 5, "Reading the header...\n");
        hdu->header = psFitsReadHeader(NULL, fits);
        if (! hdu->header) {
            psError(PS_ERR_IO, false, "Unable to read header for extension %s\n", hdu->extname);
            return false;
        }
    }

    return true;
}

// Read the HDU
// XXX: Add a region specifier?
bool pmHDURead(pmHDU *hdu,              // HDU to read
               psFits *fits             // FITS file to read from
              )
{
    assert(hdu);
    assert(fits);

    // Read the header; includes the move
    if (!pmHDUReadHeader(hdu, fits)) {
        return false;
    }

    #ifdef FITS_TABLES
    // What type is it?
    if (psFitsIsImage(hdu->header)) {
        #endif
        if (hdu->images) {
            psLogMsg(__func__, PS_LOG_WARN, "HDU %s has already been read --- overwriting.\n", hdu->extname);
            psFree(hdu->images);        // Blow away anything existing
        }
        psTrace(__func__, 5, "Reading the pixels...\n");
        hdu->images = psFitsReadImageCube(fits, psRegionSet(0,0,0,0));
        if (! hdu->images) {
            psError(PS_ERR_IO, false, "Unable to read pixels for extension %s\n", hdu->extname);
            return false;
        }
        return true;
        #ifdef FITS_TABLES

    }
    if (psFitsIsTable(hdu->table)) {
        // Read the table
        if (hdu->table) {
            psFree(hdu->table);         // Blow away anything existing
        }
        hdu->table = psFitsReadTable(fits);
        return true;
    }

    psError(PS_ERR_UNKNOWN, true, "No idea what this HDU consists of!\n");
    return false;
    #endif
}

// Write the HDU
// XXX: Add a region specifier?
bool pmHDUWrite(pmHDU *hdu,             // HDU to write
                psFits *fits            // FITS file to write to
               )
{
    assert(hdu);
    assert(fits);

    psTrace(__func__, 7, "Writing HDU %s\n", hdu->extname);

    if (hdu->images && hdu->table && !hdu->header) {
        psError(PS_ERR_IO, true, "Both image and table data provided in HDU, but no header --- "
                "don't know which to write!\n");
        return false;
    }

    if (!hdu->images && !hdu->table && !hdu->header) {
        psLogMsg(__func__, PS_LOG_WARN, "Nothing to write for HDU %s\n", hdu->extname);
        return false;
    }

    // Preserve the extension name, if it's the PHU
    char *extname = hdu->extname;       // The name of the extension
    if (strcasecmp(extname, "PHU") == 0 && hdu->header) {
        bool mdok = true;               // Status of MD lookup
        extname = psMetadataLookupStr(&mdok, hdu->header, "EXTNAME");
        if (!mdok || !extname || strlen(extname) == 0) {
            extname = "";
        }
    }

    // Only a header
    if (!hdu->images && !hdu->table) {
        // Tell CFITSIO there's nothing there
        psMetadataItem *naxis = psMetadataLookup(hdu->header, "NAXIS");
        if (naxis) {
            naxis->data.S32 = 0;
        }

        if (!psFitsWriteHeader(hdu->header, fits)) {
            psError(PS_ERR_IO, false, "Unable to write header for extension %s\n", extname);
        }
    }

    #ifdef FITS_TABLES
    if (hdu->images && (!hdu->table || psFitsIsImage(hdu->header))) {
        psFitsWriteImageCube(fits, hdu->header, hdu->images, extname);
        return true;
    }

    if (hdu->table && (!hdu->images || psFitsIsTable(hdu->header))) {
        bool psFitsWriteTable(psFits* fits, const psMetadata *header, const psArray* table);
        return true;
    }

    psError(PS_ERR_IO, true, "No idea what this HDU consists of!\n");
    return false;
    #else

    if (hdu->images) {
        psTrace(__func__, 9, "Writing pixels for %s\n", hdu->extname);
        if (!psFitsWriteImageCube(fits, hdu->header, hdu->images, extname)) {
            psError(PS_ERR_IO, false, "Unable to write image to extension %s\n", hdu->extname);
            return false;
        }
    }
    return true;
    #endif
}

