
#include <pthread.h>
#include <fcntl.h>
#include <eztrace-lib/eztrace.h>
#include "posix_io_otf2.h"

static int io_paradigm_ref_posix = 0;
static int comm_self_ref = 0;



pthread_rwlock_t open_files_lock;
struct ezt_file_handle* open_files = NULL;


static OTF2_IoAccessMode char_to_mode(const char*mode) {
  if(strcmp(mode, "r") == 0)
    return OTF2_IO_ACCESS_MODE_READ_ONLY;

  if(strcmp(mode, "w") == 0 || strcmp(mode, "a") == 0 )
    return OTF2_IO_ACCESS_MODE_WRITE_ONLY;

  if(strcmp(mode, "r+") == 0 ||
     strcmp(mode, "w+") == 0 ||
     strcmp(mode, "a+") == 0)
    return OTF2_IO_ACCESS_MODE_READ_WRITE;
  return OTF2_IO_ACCESS_MODE_READ_WRITE;
}

static OTF2_IoCreationFlag char_to_cflag(const char* mode) {
  OTF2_IoCreationFlag ret = OTF2_IO_CREATION_FLAG_NONE;

  if(strcmp(mode, "r+") == 0 ||
     strcmp(mode, "w+") == 0 ||
     strcmp(mode, "a+") == 0)
    ret |= OTF2_IO_CREATION_FLAG_CREATE;

  if(strcmp(mode, "w") == 0 ||
     strcmp(mode, "w+") == 0)
    ret |= OTF2_IO_CREATION_FLAG_TRUNCATE;

  if(strcmp(mode, "x") == 0)
    ret |= OTF2_IO_CREATION_FLAG_EXCLUSIVE;
  return ret;
}

static OTF2_IoStatusFlag char_to_sflag(const char* mode) {
  OTF2_IoStatusFlag ret = OTF2_IO_STATUS_FLAG_NONE;

  if(strcmp(mode, "e") == 0)
    ret |= OTF2_IO_STATUS_FLAG_CLOSE_ON_EXEC;

  if(strcmp(mode, "a") == 0 ||
     strcmp(mode, "a+") == 0)
    ret |= OTF2_IO_STATUS_FLAG_APPEND;

  return ret;
}


static OTF2_IoAccessMode int_to_mode(const int flags) {
  if(flags & O_RDONLY)
    return OTF2_IO_ACCESS_MODE_READ_ONLY;

  if(flags & O_WRONLY)
    return OTF2_IO_ACCESS_MODE_WRITE_ONLY;

  if(flags & O_RDWR)
    return OTF2_IO_ACCESS_MODE_READ_WRITE;
  return OTF2_IO_ACCESS_MODE_READ_WRITE;
}

static OTF2_IoCreationFlag int_to_cflag(const int flags) {
  OTF2_IoCreationFlag ret = OTF2_IO_CREATION_FLAG_NONE;
  if(flags & O_CREAT)
    ret |= OTF2_IO_CREATION_FLAG_CREATE;

  if(flags & O_TRUNC)
    ret |= OTF2_IO_CREATION_FLAG_TRUNCATE;

  if(flags & O_EXCL)
    ret |= OTF2_IO_CREATION_FLAG_EXCLUSIVE;
  return ret;
}

static OTF2_IoStatusFlag int_to_sflag(const int flags) {
  OTF2_IoStatusFlag ret = OTF2_IO_STATUS_FLAG_NONE;

  if(flags & O_CLOEXEC)
    ret |= OTF2_IO_STATUS_FLAG_CLOSE_ON_EXEC;

  if(flags & O_APPEND)
    ret |= OTF2_IO_STATUS_FLAG_APPEND;

  return ret;
}


struct ezt_file_handle* new_file_stream(const char* filename, FILE*stream);
struct ezt_file_handle* new_file_fd(const char* filename, int fd);

struct ezt_file_handle* get_file_handle_stream(FILE*stream) {

  pthread_rwlock_rdlock(&open_files_lock);
  struct ezt_file_handle* f = open_files;
  while(f) {
    if(f->stream == stream) {
      pthread_rwlock_unlock(&open_files_lock);
      return f;
    }
    f = f->next;
  }

  pthread_rwlock_unlock(&open_files_lock);
  eztrace_warn("Cannot find a handle that matches %p (fd: %d)\n",stream, fileno(stream));
  char buffer[128];
  snprintf(buffer, 128, "unknown_file_stream_%p", stream);
  return new_file_stream(buffer, stream);

}


struct ezt_file_handle* get_file_handle_fd(int fd) {
  pthread_rwlock_rdlock(&open_files_lock);
  struct ezt_file_handle* f = open_files;
  while(f) {
    if(f->fd == fd) {
      pthread_rwlock_unlock(&open_files_lock);
      return f;
    }
    f = f->next;
  }
  pthread_rwlock_unlock(&open_files_lock);
  
  if(fd == 0)
    return new_file_fd("stdin", 0);
  if(fd == 1)
    return new_file_fd("stdout", 1);
  if(fd == 2)
    return new_file_fd("stderr", 1);
 
  eztrace_warn("Cannot find a handle that matches %d\n",fd);
  char buffer[128];
  snprintf(buffer, 128, "unknown_file_fd_%d", fd);
    
  return new_file_fd(buffer, fd);
}

/* eztrace initialization function (in
   eztrace-lib/eztrace_core.c). This is quite an internal function, so
   it is not declared in eztrace.h. We still need it here though.
 */
void _eztrace_init();

void init_otf2_posixio() {
  _eztrace_init();

  if( ezt_mpi_rank == 0) {
    
    static int next_io_paradigm_ref=0;
    io_paradigm_ref_posix=next_io_paradigm_ref++; /* TODO: atomic inc */
    int string_ref = ezt_otf2_register_string("posix io");

    /* In case MPI is use, all process assume they have the same
       io_paradigm_ref_posix
       Only MPI rank 0 is allowed to register this paradigm
    */
    OTF2_GlobalDefWriter_WriteIoParadigm(_ezt_trace.global_def_writer,
					 io_paradigm_ref_posix,
					 string_ref,
					 string_ref,
					 OTF2_IO_PARADIGM_CLASS_SERIAL,
					 OTF2_IO_PARADIGM_FLAG_OS,
					 0,
					 NULL,
					 NULL,
					 NULL);

    static int next_group_ref = 0;
    int group_ref = next_group_ref++; /* TODO: atomic inc */
    int group_string_ref = ezt_otf2_register_string("");
    uint64_t comm_locations[]={0};
    /* TODO: each MPI rank should register this group */
    OTF2_GlobalDefWriter_WriteGroup( _ezt_trace.global_def_writer,
				     group_ref,
				     group_string_ref,
				     OTF2_GROUP_TYPE_COMM_GROUP,
				     OTF2_PARADIGM_MPI,
				     OTF2_GROUP_FLAG_NONE,
				     1,
				     comm_locations );

    static int next_comm_ref = 0;

    /* TODO: each MPI rank should register this communicator */
    /* TODO: make mpi_comm_self be register by the MPI module ? */
    int comm_ref = next_comm_ref++; /* TODO: atomic inc */
    int comm_string_ref = ezt_otf2_register_string("MPI_Comm_Self");
    comm_self_ref = comm_ref;
    OTF2_GlobalDefWriter_WriteComm( _ezt_trace.global_def_writer,
				    comm_ref,
				    comm_string_ref,
				    group_ref,
				    OTF2_UNDEFINED_COMM
#if OTF2_MAJOR_VERSION >= 3
				    , OTF2_COMM_FLAG_NONE
#endif
				    );
  }
}

static int get_new_matching_id() {
  static _Atomic int next_id = 0;
  return next_id++;
}

struct ezt_file_handle* close_file(FILE*stream) {
  pthread_rwlock_wrlock(&open_files_lock);
  struct ezt_file_handle* f = open_files;
  struct ezt_file_handle* to_remove = NULL;

  if(f && f->stream == stream) {
    to_remove = f;
    open_files = f->next;
    to_remove->next = NULL;
    pthread_rwlock_unlock(&open_files_lock);
    return to_remove;
  }

  while(f && f->next) {
    if(f->next->stream == stream) {

      /* remove f->next from the list */
      /* TODO: make this thread-safe */
      to_remove = f->next;
      f->next = f->next->next;
      break;
    }
    f = f->next;
  }


  if(!to_remove) {
    eztrace_warn("Warning: when closing stream %p: could not find a matching file\n", stream);
  } else {
    to_remove->next = NULL;
  }

  pthread_rwlock_unlock(&open_files_lock);
  return to_remove;
}

struct ezt_file_handle* close_file_fd(int fd) {
  pthread_rwlock_wrlock(&open_files_lock);
  struct ezt_file_handle* f = open_files;
  struct ezt_file_handle* to_remove = NULL;

  if(f && f->fd == fd) {
    to_remove = f;
    open_files = f->next;
    to_remove->next = NULL;
    pthread_rwlock_unlock(&open_files_lock);
    return to_remove;
  }

  while(f && f->next) {
    if(f->next->fd == fd) {

      /* remove f->next from the list */
      /* TODO: make this thread-safe */
      to_remove = f->next;
      f->next = f->next->next;
      break;
    }
    f = f->next;
  }

  if(!to_remove) {
    eztrace_warn("Warning: when closing fd %d: could not find a matching file\n", fd);
  } else {
    to_remove->next = NULL;
  }
  pthread_rwlock_unlock(&open_files_lock);
  return to_remove;
}


struct ezt_file_handle* new_file(const char* filename) {
  /* todo: make sure the compiler is C11 compliant and knows _Atomic */
  static _Atomic int next_file_ref=0;
  static _Atomic int next_io_handle_ref = 0;
  int file_ref = next_file_ref++;
  int string_ref = 0;
  int io_handle_ref = next_io_handle_ref++;

  eztrace_log(dbg_lvl_debug, "New file: %s. io handle=%d\n", filename, io_handle_ref);
  if(EZTRACE_SAFE) {
    string_ref = ezt_otf2_register_string(filename);
    OTF2_GlobalDefWriter_WriteIoHandle( _ezt_trace.global_def_writer,
					io_handle_ref,
					string_ref,
					file_ref,
					io_paradigm_ref_posix,
					OTF2_IO_HANDLE_FLAG_NONE,
					comm_self_ref,
					OTF2_UNDEFINED_IO_HANDLE);
  }
  struct ezt_file_handle* f = malloc(sizeof(struct ezt_file_handle));
  f->io_handle_ref = io_handle_ref;
  f->stream=NULL;
  f->fd = -1;
  f->filename = strdup(filename);
  pthread_rwlock_wrlock(&open_files_lock);
  f->next = open_files;
  open_files = f;
  pthread_rwlock_unlock(&open_files_lock);

  return f;
}

struct ezt_file_handle* new_file_stream(const char* filename, FILE*stream) {
  struct ezt_file_handle* ret = new_file(filename);
  ret->stream = stream;
  return ret;
}

struct ezt_file_handle* new_file_fd(const char* filename, int fd) {
  struct ezt_file_handle* ret = new_file(filename);
  ret->fd = fd;
  return ret;
}

void otf2_fopen_file(const char* pathname, const char*mode, FILE*stream) {
  struct ezt_file_handle* handle = new_file_stream(pathname, stream);
  EZTRACE_PROTECT {
    EZTRACE_PROTECT_ON();

    if(EZTRACE_SAFE) {
      OTF2_EvtWriter_IoCreateHandle( evt_writer,
				     NULL,
				     ezt_get_timestamp(),
				     handle->io_handle_ref,
				     char_to_mode(mode),
				     char_to_cflag(mode),
				     char_to_sflag(mode));
      
    }
    EZTRACE_PROTECT_OFF();
  }
}

void otf2_fclose_file(FILE* stream) {
    struct ezt_file_handle* handle = close_file(stream);
    EZTRACE_PROTECT {
    EZTRACE_PROTECT_ON();

    if(EZTRACE_SAFE) {
      if(handle) {
	
	OTF2_EvtWriter_IoDestroyHandle(evt_writer,
				       NULL,
				       ezt_get_timestamp(),
				       handle->io_handle_ref);
	free(handle);
      }
    }
    EZTRACE_PROTECT_OFF();
  }
}

void otf2_open_file(const char*pathname, int flags, int fd) {
  struct ezt_file_handle* handle = new_file_fd(pathname, fd);
  EZTRACE_PROTECT {
    EZTRACE_PROTECT_ON();
    if(EZTRACE_SAFE) {
      OTF2_EvtWriter_IoCreateHandle( evt_writer,
				     NULL,
				     ezt_get_timestamp(),
				     handle->io_handle_ref,
				     int_to_mode(flags),
				     int_to_cflag(flags),
				     int_to_sflag(flags));
    }
    EZTRACE_PROTECT_OFF();
  }
}

void otf2_dup_fd(int oldfd, int new_fd) {
  struct ezt_file_handle* old_handle = get_file_handle_fd(oldfd);
  struct ezt_file_handle* new_handle = NULL;
  if(old_handle)
    new_handle = new_file_fd(old_handle->filename, new_fd);
  else
    new_handle = new_file_fd("unknown_filename", new_fd);
  EZTRACE_PROTECT {
    EZTRACE_PROTECT_ON();
    if(EZTRACE_SAFE) {


      /* TODO: we should copy the flags from oldfd (except that flags
	 are not stored by eztrace. */
      OTF2_EvtWriter_IoCreateHandle( evt_writer,
				     NULL,
				     ezt_get_timestamp(),
				     new_handle->io_handle_ref,
				     OTF2_IO_ACCESS_MODE_READ_WRITE,
				     OTF2_IO_CREATION_FLAG_NONE,
				     OTF2_IO_STATUS_FLAG_NONE);
    }
    EZTRACE_PROTECT_OFF();
  }
}

void otf2_close_file(int fd) {
  struct ezt_file_handle* handle = close_file_fd(fd);
  EZTRACE_PROTECT {
    EZTRACE_PROTECT_ON();
    if(EZTRACE_SAFE) {
      if(handle) {
	OTF2_EvtWriter_IoDestroyHandle(evt_writer,
				       NULL,
				       ezt_get_timestamp(),
				       handle->io_handle_ref);
	free(handle);
      }
      
    }
    EZTRACE_PROTECT_OFF();
  }
}

static int otf2_begin_operation(struct ezt_file_handle *handle,
				OTF2_IoOperationMode mode,
				size_t count) {
  int matching_id = -1;

  EZTRACE_PROTECT {
    EZTRACE_PROTECT_ON();
    if(EZTRACE_SAFE) {
      if(handle && handle->io_handle_ref >= 0) {
      	matching_id = get_new_matching_id();

	OTF2_EvtWriter_IoOperationBegin(evt_writer,
					NULL,
					ezt_get_timestamp(),
					handle->io_handle_ref,
					mode,
					OTF2_IO_OPERATION_FLAG_NONE,
					count,
					matching_id );
      }
    }
    EZTRACE_PROTECT_OFF();
  }
  return matching_id;
}

void otf2_begin_fd_operation(int fd,
			     OTF2_IoOperationMode mode,
			     size_t count,
			     struct context*context) {
  context->handle = get_file_handle_fd(fd);
  if(context->handle)
    context->matching_id = otf2_begin_operation(context->handle, mode, count);
  else {
    if(mode==OTF2_IO_OPERATION_MODE_READ)
      eztrace_warn("read unknown fd %d\n", fd);
    else
      eztrace_warn("write unknown fd %d\n", fd);
  }
}

void otf2_begin_stream_operation(FILE* stream,
				 OTF2_IoOperationMode mode,
				 size_t count,
				 struct context*context) {
  context->handle = get_file_handle_stream(stream);
  if(context->handle)
    context->matching_id = otf2_begin_operation(context->handle, mode, count);
  else {
    if(mode==OTF2_IO_OPERATION_MODE_READ)
      eztrace_warn("read unknown stream %p\n", stream);
    else
      eztrace_warn("write unknown stream %p\n", stream);
  }
}


void otf2_end_operation(struct context* context, size_t nbytes) {
  EZTRACE_PROTECT {
    EZTRACE_PROTECT_ON();
    if(EZTRACE_SAFE) {
      if(context->handle && context->handle->io_handle_ref >= 0) {

	OTF2_EvtWriter_IoOperationComplete(evt_writer,
					   NULL,
					   ezt_get_timestamp(),
					   context->handle->io_handle_ref,
					   nbytes,
					   context->matching_id);
      }
    }
    EZTRACE_PROTECT_OFF();
  }
}

void otf2_end_fd_operation(struct context* context, size_t nbytes) {
  if(context->handle)
    otf2_end_operation(context, nbytes);
}

void otf2_end_stream_operation(struct context* context, size_t nbytes) {
  if(context->handle)
    otf2_end_operation(context, nbytes);
}



void otf2_seek_operation(struct ezt_file_handle *handle, off_t offset, int whence, off_t result) {
  EZTRACE_PROTECT {
    EZTRACE_PROTECT_ON();
    if(EZTRACE_SAFE) {
      if(handle && handle->io_handle_ref >= 0) {
      	OTF2_EvtWriter_IoSeek(evt_writer,
			      NULL,
			      ezt_get_timestamp(),
			      handle->io_handle_ref,
			      offset,
			      whence,
			      result);
      }
    }
    EZTRACE_PROTECT_OFF();
  }
}

void otf2_fd_seek_operation(int fd, off_t offset, int whence, off_t result) {
  struct ezt_file_handle *handle = get_file_handle_fd(fd);
  otf2_seek_operation(handle, offset, whence, result);
}

void otf2_stream_seek_operation(FILE* stream, off_t offset, int whence, off_t result) {
  struct ezt_file_handle *handle = get_file_handle_stream(stream);
  otf2_seek_operation(handle, offset, whence, result);
}
