#include "h5lite.h"


/*
 * Writes an atomic R vector (numeric, character, etc.) to an already created HDF5
 * dataset or attribute. This function handles data transposition and NA values for strings.
 * It is a lower-level helper called by C_h5_write_dataset and write_atomic_attribute.
 */
SEXP write_atomic_dataset(hid_t obj_id, SEXP data, const char *dtype_str, int rank, hsize_t *h5_dims) {
  herr_t status = -1;
  int must_unprotect_data = 0;
  H5I_type_t obj_type = H5Iget_type(obj_id);

  if (obj_type != H5I_DATASET && obj_type != H5I_ATTR)
    return errmsg_1("Invalid object type provided to write_atomic_dataset for %s.", dtype_str); // # nocov

  /* --- Handle Character Data (Strings) --- */
  if (TYPEOF(data) == STRSXP) {

    hsize_t n           = (hsize_t)XLENGTH(data);
    size_t  fixed_width = get_fixed_byte_width(dtype_str);

    /* BRANCH A: Fixed-Length Strings (e.g. "ascii[10]", "utf8[50]") */
    if (fixed_width > 0) {
      
      /* 1. Allocate a flat buffer for the strings in R order (Col-Major) 
       * using calloc to ensure zero-padding for short strings/NA. */
      char *f_buffer = (char *)calloc(n, fixed_width);
      if (!f_buffer) return mkChar("Memory allocation failed for fixed-length string buffer.");

      /* 2. Fill Buffer */
      for (hsize_t i = 0; i < n; i++) {
        SEXP s = STRING_ELT(data, i);
        if (s != NA_STRING) {
          const char *utf8_s = Rf_translateCharUTF8(s);
          /* Copy up to fixed_width. 
           * If src < width, strncpy pads with nulls (safety).
           * If src > width, it truncates (standard behavior).
           * Note: We do not force a null-terminator at [width-1]. HDF5 allows fully filled strings. */
          strncpy(f_buffer + (i * fixed_width), utf8_s, fixed_width);
        }
        /* If NA_STRING, we leave it as 0 (from calloc) which is an empty string in HDF5 */
      }

      /* 3. Transpose from R (Col-Major) to C (Row-Major) */
      /* For fixed strings, the "element size" for transposition is the string width in bytes. */
      char *c_buffer = (char *)malloc(n * fixed_width);
      if (!c_buffer) { free(f_buffer); return mkChar("Memory allocation failed for string buffer"); }
      
      h5_transpose((void*)f_buffer, (void*)c_buffer, rank, h5_dims, fixed_width, 0);

      /* 4. Write */
      /* create_string_type handles fixed widths based on the slash syntax */
      hid_t mem_type_id = create_string_type(dtype_str);
      status = write_buffer_to_object(obj_id, mem_type_id, c_buffer);

      free(f_buffer); free(c_buffer); H5Tclose(mem_type_id);
    }
    
    /* BRANCH B: Variable-Length Strings (Default, e.g. "ascii", "utf8") */
    else {
      /* Create a buffer of C-style strings (pointers) */
      const char **f_buffer = (const char **)malloc(n * sizeof(const char *));
      if (!f_buffer) return mkChar("Memory allocation failed for string buffer.");
      
      for (hsize_t i = 0; i < n; i++) {
        SEXP s = STRING_ELT(data, i);
        f_buffer[i] = (s == NA_STRING) ? NULL : Rf_translateCharUTF8(s);
      }
      
      /* Transpose pointers */
      const char **c_buffer = (const char **)malloc(n * sizeof(const char *));
      if (!c_buffer) { free(f_buffer); return mkChar("Memory allocation failed for string buffer"); }
      
      h5_transpose((void*)f_buffer, (void*)c_buffer, rank, h5_dims, sizeof(char*), 0);

      hid_t mem_type_id = create_string_type(dtype_str);
      status = write_buffer_to_object(obj_id, mem_type_id, c_buffer);

      free(f_buffer); free(c_buffer); H5Tclose(mem_type_id);
    }
  }
  
  /* --- Handle Numeric, Logical, Opaque, Factor Data --- */
  else {
    hsize_t total_elements = 1;
    if (rank > 0 && h5_dims) {
      for(int i=0; i<rank; i++) total_elements *= h5_dims[i];
    }
    
    /* --- Special Handling for Integer/Logical to Double Promotion --- */
    if (TYPEOF(data) == INTSXP || TYPEOF(data) == LGLSXP) {
      if (strcmp(dtype_str, "float64") == 0 ||
          strcmp(dtype_str, "float32") == 0 ||
          strcmp(dtype_str, "float16") == 0) {
        
        data = PROTECT(coerceVector(data, REALSXP));
        must_unprotect_data = 1;
      }
    }

    /* Get a direct pointer to the R object's data. */
    void *r_data_ptr = get_R_data_ptr(data);
    if (!r_data_ptr) { // # nocov start
      if (must_unprotect_data) UNPROTECT(1);
      return errmsg_1("Failed to get data pointer for %s.", dtype_str);
    } // # nocov end

    hid_t mem_type_id = create_r_memory_type(data, dtype_str);
    if (mem_type_id < 0) { // # nocov start
      if (must_unprotect_data) UNPROTECT(1);
      return errmsg_1("Failed to get memory for %s.", dtype_str);
    } // # nocov end
    
    /* Allocate a C buffer and transpose the R data into it. */
    size_t el_size = H5Tget_size(mem_type_id);
    void *c_buffer = malloc(total_elements * el_size);
    if (!c_buffer) { // # nocov start
       H5Tclose(mem_type_id);
       if (must_unprotect_data) UNPROTECT(1);
       return errmsg_1("Memory allocation failed for %s.", dtype_str);
    } // # nocov end

    /* Transpose from R's column-major to C's row-major order. */
    h5_transpose(r_data_ptr, c_buffer, rank, h5_dims, el_size, 0);

    status = write_buffer_to_object(obj_id, mem_type_id, c_buffer);

    free(c_buffer);
    H5Tclose(mem_type_id);
    if (must_unprotect_data) UNPROTECT(1);
  }

  if (status < 0)
    return errmsg_1("Failed to write data to dataset for %s.", dtype_str); // # nocov

  return R_NilValue;
}


/*
 * Orchestrates the creation and writing of an atomic (non-data.frame) attribute.
 * It creates the dataspace, file type, and attribute, then calls write_atomic_dataset.
 */
static SEXP write_atomic_attribute(hid_t file_id, hid_t obj_id, const char *attr_name, SEXP data, SEXP dtype, SEXP dims) {
  
  const char *dtype_str = CHAR(STRING_ELT(dtype, 0));
  int rank = 0;
  hsize_t *h5_dims = NULL;
  
  /* 1. Create the dataspace for the attribute. */
  hid_t space_id = create_dataspace(dims, data, &rank, &h5_dims);
  if (space_id < 0) return errmsg_1("Failed to create dataspace for attribute '%s'. Check dims.", attr_name);
  
  /* 2. Determine the HDF5 file data type. */
  hid_t file_type_id = create_h5_file_type(data, dtype_str); // This also handles special types like 'factor'
  if (file_type_id < 0) { // # nocov start
    H5Sclose(space_id);
    return errmsg_1("Failed to get file type for attribute '%s'", attr_name);
  } // # nocov end
  
  /* 3. Create Attribute Creation Property List (ACPL). 
   * We must set the character encoding to UTF-8 so the attribute NAME is stored correctly. 
   */
  hid_t acpl_id = H5Pcreate(H5P_ATTRIBUTE_CREATE);
  H5Pset_char_encoding(acpl_id, H5T_CSET_UTF8);

  /* 4. Create the attribute on the specified object. */
  hid_t attr_id = H5Acreate2(obj_id, attr_name, file_type_id, space_id, acpl_id, H5P_DEFAULT);
  
  H5Pclose(acpl_id); /* Always close the property list */

  if (attr_id < 0) { // # nocov start
    H5Sclose(space_id); H5Tclose(file_type_id);
    return errmsg_1("Failed to create attribute '%s'", attr_name);
  } // # nocov end
  
  /* 5. Write the data to the newly created attribute. */
  SEXP errmsg = write_atomic_dataset(attr_id, data, dtype_str, rank, h5_dims);
  
  H5Aclose(attr_id); H5Tclose(file_type_id); H5Sclose(space_id);
  
  if (TYPEOF(errmsg) == CHARSXP)
    return errmsg_2("Failed to write data to attribute '%s'\n%s", attr_name, CHAR(errmsg)); // # nocov
  
  return R_NilValue;
}

/*
 * Creates a dataset with a null dataspace.
 */
static SEXP write_null_dataset(hid_t file_id, const char *dname) {
  hid_t space_id = H5Screate(H5S_NULL);
  
  /* Create Link Creation Property List (LCPL) */
  hid_t lcpl_id = H5Pcreate(H5P_LINK_CREATE);
  H5Pset_create_intermediate_group(lcpl_id, 1);
  H5Pset_char_encoding(lcpl_id, H5T_CSET_UTF8); /* Explicitly set UTF-8 for the dataset name */
  
  herr_t status = handle_overwrite(file_id, dname);
  if (status < 0) { // # nocov start
    H5Sclose(space_id); H5Pclose(lcpl_id);
    return errmsg_1("Failed to overwrite existing dataset: %s", dname);
  } // # nocov end
  
  hid_t dset_id = H5Dcreate2(file_id, dname, H5T_STD_I32LE, space_id, lcpl_id, H5P_DEFAULT, H5P_DEFAULT);

  H5Pclose(lcpl_id); H5Sclose(space_id);
  if (dset_id < 0) return errmsg_1("Failed to create null dataset: %s", dname);

  H5Dclose(dset_id);
  return R_NilValue;
}

/*
 * Creates an attribute with a null dataspace.
 */
static SEXP write_null_attribute(hid_t file_id, hid_t obj_id, const char *attr_name) {
  hid_t space_id = H5Screate(H5S_NULL);
  
  /* Create Attribute Creation Property List (ACPL) */
  hid_t acpl_id = H5Pcreate(H5P_ATTRIBUTE_CREATE);
  H5Pset_char_encoding(acpl_id, H5T_CSET_UTF8); /* Explicitly set UTF-8 for the attribute name */
  
  hid_t attr_id = H5Acreate2(obj_id, attr_name, H5T_STD_I32LE, space_id, acpl_id, H5P_DEFAULT);
  
  H5Pclose(acpl_id);
  H5Sclose(space_id);
  if (attr_id < 0) return errmsg_1("Failed to create null attribute '%s'", attr_name);
  
  H5Aclose(attr_id);
  return R_NilValue;
}


/* --- WRITER: DATASET --- */
SEXP C_h5_write_dataset(SEXP filename, SEXP dset_name, SEXP data, SEXP dtype, SEXP dims, SEXP compress) {
  const char *fname = Rf_translateCharUTF8(STRING_ELT(filename, 0));
  const char *dname = Rf_translateCharUTF8(STRING_ELT(dset_name, 0));
  
  hid_t file_id = open_or_create_file(fname);

  /* --- Overwrite Logic --- */
  herr_t status = handle_overwrite(file_id, dname);
  if (status < 0) { // # nocov start
    H5Fclose(file_id);
    error("Failed to overwrite existing dataset: %s", dname);
  } // # nocov end

  SEXP errmsg = R_NilValue;
  
  /* --- Dispatch for Compound Types (Data Frames) --- */
  if (TYPEOF(data) == VECSXP) {
    /* Now creates Dimension Scale for row.names inside this function.
       write_dataframe creates its own property lists, so we don't need to do it here. */
    errmsg = write_dataframe(file_id, file_id, dname, data, dtype, compress, 0);
    H5Fclose(file_id);
  }

  /* --- Atomic Dataset Logic --- */
  else {

    const char *dtype_str = CHAR(STRING_ELT(dtype, 0));
    int rank = 0;
    hsize_t *h5_dims = NULL; /* Will be created by R_alloc() - no need to free() */
    
    if (strcmp(dtype_str, "null") == 0) {
      errmsg = write_null_dataset(file_id, dname);
      H5Fclose(file_id);
    }
    else {
  
      hid_t space_id     = create_dataspace(dims, data, &rank, &h5_dims);
      hid_t file_type_id = create_h5_file_type(data, dtype_str);
      
      /* Create Link Creation Property List (LCPL)
       * 1. create_intermediate_group = 1: Auto-create parents (like mkdir -p).
       * 2. char_encoding = UTF8: Store the dataset name as UTF-8. 
       */
      hid_t lcpl_id = H5Pcreate(H5P_LINK_CREATE);
      H5Pset_create_intermediate_group(lcpl_id, 1);
      H5Pset_char_encoding(lcpl_id, H5T_CSET_UTF8);
      
      /* Create Dataset Creation Property List (DCPL) for compression */
      hid_t dcpl_id = H5Pcreate(H5P_DATASET_CREATE);
      
      /* Only apply chunking/compression if requested and applicable */
      if (rank > 0 && XLENGTH(data) > 0) {
        size_t type_size = H5Tget_size(file_type_id);
        hsize_t *chunk_dims = (hsize_t *) R_alloc(rank, sizeof(hsize_t));
        calculate_chunk_dims(rank, h5_dims, type_size, compress, chunk_dims);
        
        apply_compression(dcpl_id, file_type_id, rank, chunk_dims, compress);
      }
      
      /* Create the dataset passing both LCPL (name encoding) and DCPL (compression) */
      hid_t dset_id = H5Dcreate2(file_id, dname, file_type_id, space_id, lcpl_id, dcpl_id, H5P_DEFAULT);
      H5Pclose(lcpl_id);
      H5Pclose(dcpl_id);
      
      if (dset_id < 0) {
        errmsg = errmsg_1("Failed to create dataset for '%s'", dname); // # nocov
      } else {
        /* Write the main data */
        errmsg = write_atomic_dataset(dset_id, data, dtype_str, rank, h5_dims);
        
        /* --- Write Dimension Scales if present --- */
        if (errmsg == R_NilValue) {
          write_r_dimscales(file_id, dset_id, dname, data);
        }
        
        H5Dclose(dset_id);
      }
      
      H5Tclose(file_type_id); H5Sclose(space_id); H5Fclose(file_id);
    }
  }

  if (TYPEOF(errmsg) == CHARSXP) error("%s", CHAR(errmsg));
  
  return R_NilValue;
}


/* --- WRITER: ATTRIBUTE --- */
SEXP C_h5_write_attribute(SEXP filename, SEXP obj_name, SEXP attr_name, SEXP data, SEXP dtype, SEXP dims) {
  const char *fname = Rf_translateCharUTF8(STRING_ELT(filename, 0));
  const char *oname = Rf_translateCharUTF8(STRING_ELT(obj_name, 0));
  const char *aname = Rf_translateCharUTF8(STRING_ELT(attr_name, 0));
  
  hid_t file_id = H5Fopen(fname, H5F_ACC_RDWR, H5P_DEFAULT);
  if (file_id < 0) error("File must exist to write attributes: %s", fname);
  
  hid_t obj_id = H5Oopen(file_id, oname, H5P_DEFAULT);
  if (obj_id < 0) { H5Fclose(file_id); error("Failed to open object: %s", oname); }

  SEXP errmsg = R_NilValue;
  
  /* --- Overwrite Logic --- */
  herr_t status = handle_attribute_overwrite(file_id, obj_id, aname);

  if (status < 0) { /* Attribute exists and could not be overwritten */
    errmsg = errmsg_1("Failed to overwrite existing attribute '%s'", aname); // # nocov
  }
  else if (TYPEOF(data) == VECSXP) { /* a data.frame */
    /* Attribute writing mode: is_attribute = 1. Scales will NOT be written.
       write_dataframe handles internal ACPL creation for UTF-8 names. */
    errmsg = write_dataframe(file_id, obj_id, aname, data, dtype, R_NilValue, 1);
  }
  else { /* an atomic type or NULL */
    const char *dtype_str_check = CHAR(STRING_ELT(dtype, 0));
    if (strcmp(dtype_str_check, "null") == 0) {
      errmsg = write_null_attribute(file_id, obj_id, aname);
    } else {
      errmsg = write_atomic_attribute(file_id, obj_id, aname, data, dtype, dims);
    }
  }
  
  H5Oclose(obj_id);
  H5Fclose(file_id);

  if (TYPEOF(errmsg) == CHARSXP) error("%s", CHAR(errmsg));

  return R_NilValue;
}
