#include "h5lite.h"

/* Internal helper to dispatch H5Dread or H5Aread based on the object type. */
static herr_t h5_read_impl(hid_t loc_id, hid_t mem_type_id, void *buf, int is_dataset, hid_t mem_space_id, hid_t file_space_id) {
  if (is_dataset) { return H5Dread(loc_id, mem_type_id, mem_space_id, file_space_id, H5P_DEFAULT, buf); }
  else            { return H5Aread(loc_id, mem_type_id, buf); }
}


/* --- READERS FOR ATOMIC TYPES --- */

static SEXP read_numeric(hid_t loc_id, int is_dataset, hid_t file_type_id, H5T_class_t class_id, 
                         int ndims, hsize_t *dims, hsize_t total_elements, R_TYPE rtype, 
                         hid_t mem_space_id, hid_t file_space_id) {
  
  SEXP result = PROTECT(allocVector(REALSXP, (R_xlen_t)total_elements));
  
  herr_t status;
  char   as_bit64 = (rtype == R_TYPE_BIT64);
  if (as_bit64) { status = h5_read_impl(loc_id, H5T_NATIVE_INT64,  REAL(result), is_dataset, mem_space_id, file_space_id); }
  else          { status = h5_read_impl(loc_id, H5T_NATIVE_DOUBLE, REAL(result), is_dataset, mem_space_id, file_space_id); }
  if (status < 0) { UNPROTECT(1); return mkChar("Failed to read numeric data"); }
  
  UNPROTECT(1);
  result = PROTECT(coerce_to_rtype(result, rtype, file_type_id));
  
  if (ndims > 1) {
    SEXP   dup_result = PROTECT(duplicate(result));
    size_t el_size    = (TYPEOF(result) == REALSXP) ? sizeof(double) : sizeof(int);
    h5_transpose(get_R_data_ptr(dup_result), get_R_data_ptr(result), ndims, dims, el_size, 1);
    set_r_dimensions(result, ndims, dims);
    UNPROTECT(1); 
  }
  
  UNPROTECT(1);
  return result; 
}

static SEXP read_complex(
    hid_t loc_id, int is_dataset, int ndims, hsize_t *dims, hsize_t total_elements, 
    hid_t mem_space_id, hid_t file_space_id) {
  
  SEXP result = PROTECT(allocVector(CPLXSXP, (R_xlen_t)total_elements));
  hid_t mem_type_id = H5Tcomplex_create(H5T_NATIVE_DOUBLE);
  herr_t status     = h5_read_impl(loc_id, mem_type_id, COMPLEX(result), is_dataset, mem_space_id, file_space_id);
  H5Tclose(mem_type_id);
  
  if (status < 0) { UNPROTECT(1); return mkChar("Failed to read complex data"); }
  
  if (ndims > 1) {
    SEXP dup_result = PROTECT(duplicate(result));
    h5_transpose(COMPLEX(dup_result), COMPLEX(result), ndims, dims, sizeof(Rcomplex), 1);
    set_r_dimensions(result, ndims, dims);
    UNPROTECT(1); 
  }
  
  UNPROTECT(1);
  return result;
}

SEXP read_character(
    hid_t loc_id, int is_dataset, hid_t file_type_id, hid_t space_id, int ndims, 
    hsize_t *dims, hsize_t total_elements, hid_t mem_space_id, hid_t file_space_id) {
  
  SEXP result = R_NilValue;
  
  /* 1. Detect the character set of the file dataset */
  H5T_cset_t file_cset = H5Tget_cset(file_type_id);
  
  if (H5Tis_variable_str(file_type_id)) {
    char **c_buffer = (char **)malloc(total_elements * sizeof(char *));
    char **f_buffer = (char **)malloc(total_elements * sizeof(char *));
    if (!c_buffer || !f_buffer) { // # nocov start
      free(c_buffer); free(f_buffer);
      return mkChar("Memory allocation failed for variable-length strings");
    } // # nocov end
    
    hid_t mem_type = H5Tcopy(H5T_C_S1);
    H5Tset_size(mem_type, H5T_VARIABLE); 
    H5Tset_cset(mem_type, file_cset);
    
    if (h5_read_impl(loc_id, mem_type, c_buffer, is_dataset, mem_space_id, file_space_id) < 0) {
      result = PROTECT(mkChar("Failed to read variable-length strings")); // # nocov
    } else {
      result = PROTECT(allocVector(STRSXP, (R_xlen_t)total_elements));
      h5_transpose(c_buffer, f_buffer, ndims, dims, sizeof(char*), 1);
      for (hsize_t i = 0; i < total_elements; i++) {
        if (f_buffer[i]) { SET_STRING_ELT(result, i, mkCharCE(f_buffer[i], CE_UTF8)); }
        else             { SET_STRING_ELT(result, i, NA_STRING); } // # nocov
      }
    }
    /* --- CLEANUP (Hyperslab-safe) --- */
    hid_t reclaim_space = (mem_space_id != H5S_ALL) ? mem_space_id : space_id;
    if (is_dataset) { H5Dvlen_reclaim(mem_type, reclaim_space, H5P_DEFAULT, c_buffer); }
    else            { H5Treclaim(mem_type, reclaim_space, H5P_DEFAULT, c_buffer);      }
    free(c_buffer); free(f_buffer); H5Tclose(mem_type);
  }
  else {
    size_t type_size = H5Tget_size(file_type_id);
    char *c_buffer = (char *)malloc(total_elements * type_size);
    char *f_buffer = (char *)malloc(total_elements * type_size);
    if (!c_buffer || !f_buffer) { // # nocov start
      free(c_buffer); free(f_buffer);
      return mkChar("Memory allocation failed for fixed-length strings");
    } // # nocov end
    
    hid_t mem_type = H5Tcopy(H5T_C_S1);
    H5Tset_size(mem_type, type_size);
    H5Tset_cset(mem_type, file_cset);
    
    if (h5_read_impl(loc_id, mem_type, c_buffer, is_dataset, mem_space_id, file_space_id) < 0) {
      result = PROTECT(mkChar("Failed to read fixed-length strings")); // # nocov
    }else {
      result = PROTECT(allocVector(STRSXP, (R_xlen_t)total_elements));
      h5_transpose(c_buffer, f_buffer, ndims, dims, type_size, 1);
      char *single_str = (char *)R_alloc(type_size + 1, sizeof(char));
      for (hsize_t i = 0; i < total_elements; i++) {
        memcpy(single_str, f_buffer + (i * type_size), type_size);
        single_str[type_size] = '\0';
        SET_STRING_ELT(result, i, mkCharCE(single_str, CE_UTF8));
      }
    }
    free(c_buffer); free(f_buffer); H5Tclose(mem_type);
  }
  
  if (TYPEOF(result) != CHARSXP)
    set_r_dimensions(result, ndims, dims);
  
  UNPROTECT(1);
  return result;
}

static SEXP read_raw(
    hid_t loc_id, int is_dataset, hid_t file_type_id, int ndims, 
    hsize_t *dims, hsize_t total_elements, hid_t mem_space_id, hid_t file_space_id) {
  
  size_t type_size = H5Tget_size(file_type_id);
  if (type_size != 1) return mkChar("h5lite only supports reading 1-byte opaque types as raw vectors");
  
  hid_t  mem_type = H5Tcreate(H5T_OPAQUE, type_size);
  SEXP   result   = PROTECT(allocVector(RAWSXP, (R_xlen_t)total_elements));
  herr_t status   = h5_read_impl(loc_id, mem_type, RAW(result), is_dataset, mem_space_id, file_space_id);
  
  if (status < 0) {
    result = PROTECT(mkChar("Failed to read raw data")); // # nocov
  }
  else if (ndims > 1) {
    SEXP dup_result = PROTECT(duplicate(result));
    h5_transpose(RAW(dup_result), RAW(result), ndims, dims, type_size, 1);
    set_r_dimensions(result, ndims, dims);
    UNPROTECT(1); 
  }
  
  H5Tclose(mem_type);
  UNPROTECT(1);
  return result;
}


static SEXP read_factor(
    hid_t loc_id, int is_dataset, hid_t file_type_id, int ndims, 
    hsize_t *dims, hsize_t total_elements, hid_t mem_space_id, hid_t file_space_id) {
  
  int n_members = H5Tget_nmembers(file_type_id);
  if (n_members <= 0) return mkChar("enum type has no members");
  
  /* 1. Create a memory Enum type based on integers */
  hid_t mem_type_id = H5Tcreate(H5T_ENUM, sizeof(int));
  if (mem_type_id < 0) return mkChar("Failed to create memory enum type");

  /* 2. Build the R levels vector and populate the memory Enum */
  /* We map the file's names to R's 1-based indices (1, 2, 3...) */
  SEXP levels = PROTECT(allocVector(STRSXP, n_members));
  
  for (int i = 0; i < n_members; i++) {
    char *name = H5Tget_member_name(file_type_id, i);
    int r_index = i + 1; /* R uses 1-based indexing */
    
    if (name) {
      /* Tell HDF5: "When you see this name, read it as integer 'r_index'" */
      H5Tenum_insert(mem_type_id, name, &r_index);
      SET_STRING_ELT(levels, i, mkCharCE(name, CE_UTF8));
      H5free_memory(name);
    } else {
      SET_STRING_ELT(levels, i, NA_STRING); // # nocov
    }
  }

  /* 3. Read the data using the custom memory Enum type */
  /* HDF5 will automatically match names from the file to our 1-based integers */
  SEXP result = PROTECT(allocVector(INTSXP, (R_xlen_t)total_elements));
  herr_t status = h5_read_impl(loc_id, mem_type_id, INTEGER(result), is_dataset, mem_space_id, file_space_id);
  
  H5Tclose(mem_type_id);

  if (status < 0) { // # nocov start
    UNPROTECT(2); 
    return mkChar("Failed to read enum values"); 
  } // # nocov end
  
  /* 4. Handle Transposition (if necessary) */
  if (ndims > 1) {
    SEXP dup_result = PROTECT(duplicate(result));
    h5_transpose(INTEGER(dup_result), INTEGER(result), ndims, dims, sizeof(int), 1);
    set_r_dimensions(result, ndims, dims);
    UNPROTECT(1); 
  }
  
  /* 5. Set Class and Levels */
  setAttrib(result, R_LevelsSymbol, levels);
  
  SEXP class_attr = PROTECT(allocVector(STRSXP, 1));
  SET_STRING_ELT(class_attr, 0, mkChar("factor"));
  setAttrib(result, R_ClassSymbol, class_attr);
  UNPROTECT(1); // class_attr
  
  UNPROTECT(2); // levels, result
  return result;
}


/* --- DATASET READ ENTRY POINT --- */

SEXP C_h5_read_dataset(SEXP filename, SEXP dataset_name, SEXP rmap, SEXP element_name, SEXP start, SEXP count) {
  
  const char *fname   = Rf_translateCharUTF8(STRING_ELT(filename,     0));
  const char *dname   = Rf_translateCharUTF8(STRING_ELT(dataset_name, 0));
  const char *el_name = Rf_translateCharUTF8(STRING_ELT(element_name, 0));
  
  hid_t file_id = H5Fopen(fname, H5F_ACC_RDONLY, H5P_DEFAULT);
  if (file_id < 0) error("Failed to open file: %s", fname);
  
  hid_t dset_id = H5Dopen2(file_id, dname, H5P_DEFAULT);
  if (dset_id < 0) { H5Fclose(file_id); error("Failed to open dataset: %s", dname); }
  
  hid_t       file_type_id = H5Dget_type(dset_id);
  H5T_class_t class_id     = H5Tget_class(file_type_id);
  hid_t       space_id     = H5Dget_space(dset_id);
  R_TYPE      rtype        = rtype_from_map(file_type_id, rmap, el_name);
  
  if (H5Sget_simple_extent_type(space_id) == H5S_NULL || rtype == R_TYPE_NULL) {
    H5Sclose(space_id); H5Tclose(file_type_id); H5Dclose(dset_id); H5Fclose(file_id);
    return R_NilValue;
  }
  
  int ndims = H5Sget_simple_extent_ndims(space_id);
  hsize_t total_elements = 1;
  hsize_t *dims      = NULL;
  hsize_t *mem_dims  = NULL;
  hsize_t *offset    = NULL;
  hid_t mem_space_id = H5S_ALL;
  int has_hyperslab = 0;
  
  if (ndims > 0) {
    dims     = (hsize_t *)R_alloc(ndims, sizeof(hsize_t));
    mem_dims = (hsize_t *)R_alloc(ndims, sizeof(hsize_t));
    offset   = (hsize_t *)R_alloc(ndims, sizeof(hsize_t));
    H5Sget_simple_extent_dims(space_id, dims, NULL);
    
    for (int i = 0; i < ndims; i++) {
      offset[i] = 0;
      mem_dims[i] = dims[i];
    }
    
    /* 1. Parse Hyperslab Constraints (Guaranteed by R to both be NULL, or both be valid REALSXP) */
    if (start != R_NilValue && count != R_NilValue) {
      has_hyperslab = 1;
      int start_len = Rf_length(start);
      
      hsize_t count_val = (hsize_t) REAL(count)[0];
      
      /* Generate Generalized Mapping (C is 0-indexed) */
      /* Pattern for ndims >= 3: (ndims-1), (ndims-2), ..., 2, 0, 1 */
      int *dim_map = (int *)R_alloc(ndims, sizeof(int));
      
      if (ndims >= 3) {
        for (int i = 0; i < ndims - 2; i++) {
          dim_map[i] = ndims - 1 - i;
        }
        dim_map[ndims - 2] = 0; /* Rows */
      dim_map[ndims - 1] = 1; /* Cols */
      } else {
        /* Fallback for 1D and 2D: 0, 1 */
        for (int i = 0; i < ndims; i++) {
          dim_map[i] = i;
        }
      }
      
      for(int i = 0; i < start_len && i < ndims; i++) {
        int target_dim = dim_map[i]; /* Look up the targeted HDF5 dimension */
        
        hsize_t start_val = (hsize_t) REAL(start)[i];
        
        offset[target_dim] = start_val - 1; /* Adjust 1-based R index to 0-based HDF5 index */
        mem_dims[target_dim] = 1;           /* Default to 1 element */
      }
      
      /* Apply the single count value to the LAST dimension targeted by start */
      if (start_len > 0 && start_len <= ndims) {
        int target_dim = dim_map[start_len - 1];
        mem_dims[target_dim] = count_val;
      }
    }
    
    /* 2. Apply Selection */
    if (has_hyperslab) {
      H5Sselect_hyperslab(space_id, H5S_SELECT_SET, offset, NULL, mem_dims, NULL);
      mem_space_id = H5Screate_simple(ndims, mem_dims, NULL);
      
      /* Recalculate total elements based on the newly selected hyperslab block */
      total_elements = 1;
      for (int i = 0; i < ndims; i++) total_elements *= mem_dims[i];
    } else {
      for (int i = 0; i < ndims; i++) total_elements *= dims[i];
    }
  }
  
  SEXP result;
  hsize_t *final_dims = has_hyperslab ? mem_dims : dims;
  
  /* Dispatch to specific readers, passing the memory and file dataspace IDs */
  if (class_id == H5T_INTEGER || class_id == H5T_FLOAT) {
    result = read_numeric(dset_id, 1, file_type_id, class_id, ndims, final_dims, total_elements, rtype, mem_space_id, space_id);
  }
  else if (class_id == H5T_COMPLEX) {
    result = read_complex(dset_id, 1, ndims, final_dims, total_elements, mem_space_id, space_id);
  }
  else if (class_id == H5T_STRING) {
    result = read_character(dset_id, 1, file_type_id, space_id, ndims, final_dims, total_elements, mem_space_id, space_id);
  }
  else if (class_id == H5T_OPAQUE) {
    result = read_raw(dset_id, 1, file_type_id, ndims, final_dims, total_elements, mem_space_id, space_id);
  }
  else if (class_id == H5T_ENUM) {
    result = read_factor(dset_id, 1, file_type_id, ndims, final_dims, total_elements, mem_space_id, space_id);
  }
  else if (class_id == H5T_COMPOUND) {
    result = read_data_frame(dset_id, 1, file_type_id, space_id, rmap, mem_space_id, space_id, 
                             has_hyperslab ? offset : NULL, has_hyperslab ? mem_dims : NULL);
  }
  else {
    result = mkChar("Unsupported HDF5 type"); // # nocov
  }
  
  /* For Atomic types, check for attached Dimension Scales and restore names/dimnames */
  if (class_id != H5T_COMPOUND && TYPEOF(result) != CHARSXP) {
    read_r_dimscales(dset_id, ndims, result, has_hyperslab ? offset : NULL, has_hyperslab ? mem_dims : NULL);
  }
  
  /* Clean up */
  if (mem_space_id != H5S_ALL) H5Sclose(mem_space_id);
  H5Tclose(file_type_id); H5Sclose(space_id); H5Dclose(dset_id); H5Fclose(file_id);
  
  if (TYPEOF(result) == CHARSXP)
    error("Error reading dataset '%s'\n%s", dname, CHAR(result)); // # nocov
  
  return result;
}


/* --- ATTRIBUTE READ ENTRY POINT --- */

SEXP C_h5_read_attribute(SEXP filename, SEXP obj_name, SEXP attr_name, SEXP rmap) {
  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_RDONLY, H5P_DEFAULT);
  if (file_id < 0) error("Failed to open file: %s", fname);
  
  hid_t attr_id = H5Aopen_by_name(file_id, oname, aname, H5P_DEFAULT, H5P_DEFAULT);
  if (attr_id < 0) { H5Fclose(file_id); error("Failed to open attribute: %s", aname); }
  
  hid_t       file_type_id = H5Aget_type(attr_id);
  H5T_class_t class_id     = H5Tget_class(file_type_id);
  hid_t       space_id     = H5Aget_space(attr_id);
  R_TYPE      rtype        = rtype_from_map(file_type_id, rmap, aname);
  
  if (H5Sget_simple_extent_type(space_id) == H5S_NULL || rtype == R_TYPE_NULL) { // # nocov start
    H5Sclose(space_id); H5Tclose(file_type_id); H5Aclose(attr_id); H5Fclose(file_id);
    return R_NilValue;
  } // # nocov end
  
  int ndims = H5Sget_simple_extent_ndims(space_id);
  hsize_t total_elements = 1;
  hsize_t *dims = NULL;
  
  if (ndims > 0) {
    dims = (hsize_t *)R_alloc(ndims, sizeof(hsize_t));
    H5Sget_simple_extent_dims(space_id, dims, NULL);
    for (int i = 0; i < ndims; i++) total_elements *= dims[i];
  }
  
  SEXP result;
  
  if (class_id == H5T_INTEGER || class_id == H5T_FLOAT) {
    result = read_numeric(attr_id, 0, file_type_id, class_id, ndims, dims, total_elements, rtype, H5S_ALL, H5S_ALL);
  }
  else if (class_id == H5T_COMPLEX) {
    result = read_complex(attr_id, 0, ndims, dims, total_elements, H5S_ALL, H5S_ALL);
  }
  else if (class_id == H5T_STRING) {
    result = read_character(attr_id, 0, file_type_id, space_id, ndims, dims, total_elements, H5S_ALL, H5S_ALL);
  }
  else if (class_id == H5T_OPAQUE) {
    result = read_raw(attr_id, 0, file_type_id, ndims, dims, total_elements, H5S_ALL, H5S_ALL);
  }
  else if (class_id == H5T_ENUM) {
    result = read_factor(attr_id, 0, file_type_id, ndims, dims, total_elements, H5S_ALL, H5S_ALL);
  }
  else if (class_id == H5T_COMPOUND) {
    result = read_data_frame(attr_id, 0, file_type_id, space_id, rmap, H5S_ALL, H5S_ALL, NULL, NULL);
  }
  else {
    result = mkChar("Unsupported HDF5 type"); // # nocov
  }
  
  H5Tclose(file_type_id); H5Sclose(space_id); H5Aclose(attr_id); H5Fclose(file_id);
  
  if (TYPEOF(result) == CHARSXP)
    error("Error reading attribute '%s': %s", aname, CHAR(result)); // # nocov
  
  return result;
}
