import os
import bz2
import pickle
import tempfile
from time import time

class _PickleSerializationMixin(object):
  """
  Mixin for pickle serialization functionalities that are attached to `ratio1.Logger`.

  This mixin cannot be instantiated because it is built just to provide some additional
  functionalities for `ratio1.Logger`

  In this mixin we can use any attribute/method of the Logger.
  """

  def __init__(self):
    super(_PickleSerializationMixin, self).__init__()
    return

  def _save_compressed_pickle(self, full_filename, myobj, locking=False):
    """
    save object to file using pickle

    @param full_filename: name of destination file
    @param myobj: object to save (has to be pickleable)
    """
    with self.managed_lock_resource(full_filename, condition=locking):
      try:
        fhandle = bz2.BZ2File(full_filename, 'wb')
        pickle.dump(myobj, fhandle, protocol=pickle.HIGHEST_PROTOCOL)
        fhandle.close()
      except:
        self.P('ERROR: File ' + full_filename + ' cannot be written!')
        return False
    # endwith conditional lock
    return True


  def _load_compressed_pickle(self, full_filename, locking=False):
    """
    Load from filename using pickle

    @param full_filename: name of file to load from
    """
    with self.managed_lock_resource(full_filename, condition=locking):
      try:
        fhandle = bz2.BZ2File(full_filename, 'rb')
        myobj = pickle.load(fhandle)
        fhandle.close()
      except:
        self.P('ERROR: File ' + full_filename + ' cannot be read!')
        return None
    # endwith conditional lock

    return myobj


  def _fsync_dir(self, dirpath: str):
    """Best-effort directory fsync for durable rename; silently ignore if unsupported."""
    try:
      # Linux: O_DIRECTORY available; on other OSes this may fail -> we ignore.
      dir_fd = os.open(dirpath, getattr(os, "O_DIRECTORY", 0))
      try:
        os.fsync(dir_fd)
      finally:
        os.close(dir_fd)
    except Exception:
      # Not all platforms/filesystems support directory fsync; ignore if it fails.
      pass
  # enddef

  def _fsync_file(self, filepath: str):
    """Best-effort file fsync for durability; silently ignore if unsupported."""
    # Ensure file contents hit disk even if helper didn't fsync
    try:
      rd_fd = os.open(filepath, os.O_RDONLY)
      try:
        os.fsync(rd_fd)
      finally:
        os.close(rd_fd)
    except Exception:
      # Best effort; continue to replace
      pass
  # enddef

  def save_pickle(
      self, data, fn, folder=None,
      use_prefix=False, verbose=True,
      compressed=False,
      subfolder_path=None,
      locking=True,
  ):
    """
    compressed: True if compression is required OR you can just add '.pklz' to `fn`
    """

    def P(s):
      if verbose:
        self.P(s)
      return
    # enddef

    lfld = self.get_target_folder(target=folder)

    if lfld is None:
      P("Assuming `fn` param ({}) is a full path".format(fn))
      datafile = fn
    else:
      if use_prefix:
        fn = self.file_prefix + '_' + fn
      datafolder = lfld
      if subfolder_path is not None:
        datafolder = os.path.join(datafolder, subfolder_path.lstrip('/'))
        os.makedirs(datafolder, exist_ok=True)
      datafile = os.path.join(datafolder, fn)

    os.makedirs(os.path.split(datafile)[0], exist_ok=True)
    target_dir = os.path.dirname(datafile)

    tm_start = time()
    tm_elapsed = None
    err_msg = None

    # Create a temp file in the SAME directory so os.replace is atomic
    tmp_fd, tmp_path = tempfile.mkstemp(prefix=os.path.basename(fn) + ".", suffix=".tmp", dir=target_dir)
    os.close(tmp_fd)  # we'll reopen with Python I/O or let helper write to it

    try:
      if compressed or '.pklz' in fn:
        if not compressed:
          P("Saving pickle with compression=True forced due to extension")
        else:
          P("Saving pickle with compression...")

        ok = self._save_compressed_pickle(tmp_path, myobj=data, locking=locking)
        if ok:
          # Ensure data is flushed to disk before rename
          self._fsync_file(tmp_path)
          os.replace(tmp_path, datafile)  # atomic move
          self._fsync_dir(target_dir)
          tm_elapsed = time() - tm_start
          P("  Compressed pickle {} saved in {} folder in {:.1f}s".format(fn, folder, tm_elapsed))
        else:
          P("  FAILED compressed pickle save!")
      else:
        P("Saving uncompressed pikle (lock:{}) : {} ".format(locking, datafile))
        with self.managed_lock_resource(datafile, condition=locking):
          try:
            with open(tmp_path, 'wb') as fhandle:
              pickle.dump(data, fhandle, protocol=pickle.HIGHEST_PROTOCOL)
              fhandle.flush()
              os.fsync(fhandle.fileno())  # ensure data is written to disk
            # Atomic replace + dir fsync
            os.replace(tmp_path, datafile)
            self._fsync_dir(target_dir)
            tm_elapsed = time() - tm_start
          except Exception as e:
            err_msg = e
        if tm_elapsed is not None:
          if verbose:
            P("  Saved pickle '{}' in '{}' folder in {:.1f}s".format(fn, folder, tm_elapsed))
        else:
          # maybe show this only if verbose?
          P(f"  FAILED pickle save! Error: {err_msg}")
      # endif compressed or not
    except Exception as e:
      err_msg = e
      try:
        if os.path.exists(datafile):
          os.remove(tmp_path)
      except:
        pass
      P(f"  FAILED pickle save! Error: {err_msg}")
    return datafile


  def save_pickle_to_data(self, data, fn, compressed=False, verbose=True, 
                          subfolder_path=None, locking=True):
    """
    compressed: True if compression is required OR you can just add '.pklz' to `fn`
    """
    return self.save_pickle(
      data, fn,
      folder='data',
      compressed=compressed,
      subfolder_path=subfolder_path,
      verbose=verbose,
      locking=locking,
    )


  def save_pickle_to_models(self, data, fn, compressed=False, verbose=True, 
                            subfolder_path=None, locking=True):
    """
    compressed: True if compression is required OR you can just add '.pklz' to `fn`
    """
    return self.save_pickle(
      data, fn,
      folder='models',
      compressed=compressed,
      subfolder_path=subfolder_path,
      verbose=verbose,
      locking=locking,
    )


  def save_pickle_to_output(self, data, fn, compressed=False, verbose=True, 
                            subfolder_path=None, locking=True):
    """
    compressed: True if compression is required OR you can just add '.pklz' to `fn`
    """
    return self.save_pickle(
      data, fn,
      folder='output',
      compressed=compressed,
      subfolder_path=subfolder_path,
      verbose=verbose,
      locking=locking,
    )


  def load_pickle_from_models(self, fn, decompress=False, verbose=True, 
                              subfolder_path=None, locking=True):
    """
     decompressed : True if the file was saved with `compressed=True` or you can just use '.pklz'
    """
    return self.load_pickle(
      fn,
      folder='models',
      decompress=decompress,
      verbose=verbose,
      subfolder_path=subfolder_path,
      locking=locking,
    )


  def load_pickle_from_data(self, fn, decompress=False, verbose=True, 
                            subfolder_path=None, locking=True):
    """
     decompressed : True if the file was saved with `compressed=True` or you can just use '.pklz'
    """
    return self.load_pickle(
      fn,
      folder='data',
      decompress=decompress,
      verbose=verbose,
      subfolder_path=subfolder_path,
      locking=locking,
    )


  def load_pickle_from_output(self, fn, decompress=False, verbose=True, 
                              subfolder_path=None, locking=True):
    """
     decompressed : True if the file was saved with `compressed=True` or you can just use '.pklz'
    """
    return self.load_pickle(
      fn,
      folder='output',
      decompress=decompress,
      verbose=verbose,
      subfolder_path=subfolder_path,
      locking=locking,
    )


  def load_pickle(self, fn, folder=None, decompress=False, verbose=True,
                  subfolder_path=None, locking=True):
    """
     load_from: 'data', 'output', 'models'
     decompressed : True if the file was saved with `compressed=True` or you can just use '.pklz'
    """
    if verbose:
      P = self.P
    else:
      P = lambda x, color=None: x

    lfld = self.get_target_folder(target=folder)

    if lfld is None:
      P("Loading pickle ... Assuming `fn` param ({}) is a full path".format(fn))
      datafile = fn
    else:
      datafolder = lfld
      if subfolder_path is not None:
        datafolder = os.path.join(datafolder, subfolder_path.lstrip('/'))
      datafile = os.path.join(datafolder, fn)
      P("Loading pickle (locked:{}) from {}".format(locking, datafile))
    #endif full path or not
    data = None
    exc = None
    if os.path.isfile(datafile):
      if decompress or '.pklz' in datafile:
        if not decompress:
          P("Loading pickle with decompress=True forced due to extension")
        else:
          P("Loading pickle with decompression...")
        data = self._load_compressed_pickle(datafile)
      else:
        with self.managed_lock_resource(datafile, condition=locking):
          try:
            with open(datafile, "rb") as f:
              data = pickle.load(f)
          except Exception as _exc:
            data = None
            exc = _exc
        # endwith conditional lock
      #endif decompress or not
      if data is None:
        P("  {} load failed with error {}".format(datafile, exc), color='r')
      else:
        P("  Loaded: {}".format(datafile))
      #endif data is None
    else:
      P("  File not found! Pickle load failed.", color='r')
    return data
  
  
  def update_pickle_from_data(self, 
                              fn, 
                              update_callback, 
                              decompress=False, 
                              verbose=False, 
                              subfolder_path=None,
                              force_update=False):
    assert update_callback is not None, "update_callback must be defined!"
    datafile = self.get_file_path(
      fn=fn,
      folder='data',
      subfolder_path=subfolder_path,
      force=True
      )
    if datafile is None:
      self.P("update_pickle_from_data failed due to missing {}".format(datafile), color='error')
      return False

    with self.managed_lock_resource(datafile):
      result = None
      try:
        data = self.load_pickle_from_data(
          fn=fn,
          decompress=decompress,
          verbose=verbose,
          subfolder_path=subfolder_path,
          locking=False,
          )
        
        if data is not None or force_update:
          data = update_callback(data)
          
          self.save_pickle_to_data(
            data=data, 
            fn=fn,
            compressed=decompress,
            verbose=verbose,
            subfolder_path=subfolder_path,
            locking=False,
            )
          result = True
      except Exception as e:
        self.P("update_pickle_from_data failed: {}".format(e), color='error')
        result = False
    # endwith lock
    return result
      