import os
import random
import string
import time
from typing import Iterable, Union

IS_WINDOWS = os.name == 'nt'
# Max attempt if another process is reading/writing, effective only on Windows
WINDOWS_MAX_ATTEMPT = 5
# Base time to wait between retries (seconds)
WINDOWS_RETRY_DELAY = 0.05


def random_id():
    """
    Returns:
        str: Random ID, like "sTD2kF"
    """
    # 6 random letter (62^6 combinations) would be enough
    return ''.join(random.sample(string.ascii_letters + string.digits, 6))


def is_tmp_file(file: str) -> bool:
    """
    Check if a filename is tmp file
    """
    # Check suffix first to reduce regex calls
    if not file.endswith('.tmp'):
        return False
    # Check temp file format
    dot = file[-11:-10]
    if not dot:
        return False
    rid = file[-10:-4]
    return rid.isalnum()


def to_tmp_file(file: str) -> str:
    """
    Convert a filename or directory name to tmp
    filename -> filename.sTD2kF.tmp
    """
    suffix = random_id()
    return f'{file}.{suffix}.tmp'


def to_nontmp_file(file: str) -> str:
    """
    Convert a tmp filename or directory name to original file
    filename.sTD2kF.tmp -> filename
    """
    if is_tmp_file(file):
        return file[:-11]
    else:
        return file


def windows_attempt_delay(attempt: int) -> float:
    """
    Exponential Backoff if file is in use on Windows

    Args:
        attempt: Current attempt, starting from 0

    Returns:
        float: Seconds to wait
    """
    return 2 ** attempt * WINDOWS_RETRY_DELAY


def replace_tmp(tmp: str, file: str):
    """
    Replace temp file to file

    Raises:
        PermissionError: (Windows only) If another process is still reading the file and all retries failed
        FileNotFoundError: If tmp file gets deleted unexpectedly
    """
    if IS_WINDOWS:
        # PermissionError on Windows if another process is reading
        last_error = None
        for attempt in range(WINDOWS_MAX_ATTEMPT):
            try:
                # Atomic operation
                os.replace(tmp, file)
                # success
                return
            except PermissionError as e:
                last_error = e
                delay = windows_attempt_delay(attempt)
                time.sleep(delay)
                continue
            except FileNotFoundError:
                # tmp file gets deleted unexpectedly
                raise
            except Exception as e:
                last_error = e
                break
    else:
        # Linux and Mac allow existing reading
        try:
            # Atomic operation
            os.replace(tmp, file)
            # success
            return
        except FileNotFoundError:
            raise
        except Exception as e:
            last_error = e

    # Clean up tmp file on failure
    try:
        os.unlink(tmp)
    except FileNotFoundError:
        # tmp file already get deleted
        pass
    except:
        pass
    if last_error is not None:
        raise last_error from None


def atomic_replace(replace_from: str, replace_to: str):
    """
    Replace file or directory

    Raises:
        PermissionError: (Windows only) If another process is still reading the file and all retries failed
        FileNotFoundError:
    """
    if IS_WINDOWS:
        # PermissionError on Windows if another process is reading
        last_error = None
        for attempt in range(WINDOWS_MAX_ATTEMPT):
            try:
                # Atomic operation
                os.replace(replace_from, replace_to)
                # success
                return
            except PermissionError as e:
                last_error = e
                delay = windows_attempt_delay(attempt)
                time.sleep(delay)
                continue
            except FileNotFoundError:
                raise
            except Exception as e:
                last_error = e
                break
        if last_error is not None:
            raise last_error from None
    else:
        # Linux and Mac
        os.replace(replace_from, replace_to)


def file_write(file: str, data: Union[str, bytes]):
    """
    Write data into file, auto create directory
    Auto determines write mode based on the type of data.
    """
    if isinstance(data, str):
        mode = 'w'
        encoding = 'utf-8'
        newline = ''
    elif isinstance(data, bytes):
        mode = 'wb'
        encoding = None
        newline = None
        # Create memoryview as Pathlib do
        data = memoryview(data)
    else:
        typename = str(type(data))
        if typename == "<class 'numpy.ndarray'>":
            mode = 'wb'
            encoding = None
            newline = None
        else:
            mode = 'w'
            encoding = 'utf-8'
            newline = ''

    try:
        # Write temp file
        with open(file, mode=mode, encoding=encoding, newline=newline) as f:
            f.write(data)
            # Ensure data flush to disk
            f.flush()
            os.fsync(f.fileno())
    except FileNotFoundError:
        # Create parent directory
        directory = os.path.dirname(file)
        if directory:
            os.makedirs(directory, exist_ok=True)
        # Write again
        with open(file, mode=mode, encoding=encoding, newline=newline) as f:
            f.write(data)
            # Ensure data flush to disk
            f.flush()
            os.fsync(f.fileno())


def file_write_stream(file: str, data_generator):
    """
    Only creates a file if the generator yields at least one data chunk.
    Auto determines write mode based on the type of first chunk.

    Args:
        file: Target file path
        data_generator: An iterable that yields data chunks (str or bytes)
    """
    # Convert generator to iterator to ensure we can peek at first chunk
    data_iter = iter(data_generator)

    # Try to get the first chunk
    try:
        first_chunk = next(data_iter)
    except StopIteration:
        # Generator is empty, no file will be created
        return

    # Determine mode, encoding and newline from first chunk
    if isinstance(first_chunk, str):
        mode = 'w'
        encoding = 'utf-8'
        newline = ''
    elif isinstance(first_chunk, bytes):
        mode = 'wb'
        encoding = None
        newline = None
    else:
        # Default to text mode for other types
        mode = 'w'
        encoding = 'utf-8'
        newline = ''

    try:
        # Write temp file
        with open(file, mode=mode, encoding=encoding, newline=newline) as f:
            f.write(first_chunk)
            for chunk in data_iter:
                f.write(chunk)
            # Ensure data flush to disk
            f.flush()
            os.fsync(f.fileno())
    except FileNotFoundError:
        # Create parent directory
        directory = os.path.dirname(file)
        if directory:
            os.makedirs(directory, exist_ok=True)
        # Write again
        with open(file, mode=mode, encoding=encoding, newline=newline) as f:
            f.write(first_chunk)
            for chunk in data_iter:
                f.write(chunk)
            # Ensure data flush to disk
            f.flush()
            os.fsync(f.fileno())


def atomic_write(
        file: str,
        data: Union[str, bytes],
):
    """
    Atomic file write with minimal IO operation
    and handles cases where file might be read by another process.

    os.replace() is an atomic operation among all OS,
    we write to temp file then do os.replace()

    Args:
        file:
        data:
    """
    temp = to_tmp_file(file)
    file_write(temp, data)
    replace_tmp(temp, file)


def atomic_write_stream(
        file: str,
        data_generator,
):
    """
    Atomic file write with streaming data support.
    Handles cases where file might be read by another process.

    os.replace() is an atomic operation among all OS,
    we write to temp file then do os.replace()

    Args:
        file: Target file path
        data_generator: An iterable that yields data chunks (str or bytes)
    """
    temp = to_tmp_file(file)
    file_write_stream(temp, data_generator)
    replace_tmp(temp, file)


def file_read_text(
        file: str,
        encoding: str = 'utf-8',
        errors: str = 'strict'
) -> str:
    """
    Args:
        file:
        encoding:
        errors: 'strict', 'ignore', 'replace' and any other errors mode in open()
    """
    try:
        with open(file, mode='r', encoding=encoding, errors=errors) as f:
            return f.read()
    except FileNotFoundError:
        return ''


def file_read_text_stream(
        file: str,
        encoding: str = 'utf-8',
        errors: str = 'strict',
        chunk_size: int = 8192
) -> Iterable[str]:
    """
    Args:
        file:
        encoding:
        errors: 'strict', 'ignore', 'replace' and any other errors mode in open()
        chunk_size:
    """
    try:
        with open(file, mode='r', encoding=encoding, errors=errors) as f:
            while 1:
                chunk = f.read(chunk_size)
                if not chunk:
                    return
                yield chunk
    except FileNotFoundError:
        return


def file_read_bytes(file: str) -> bytes:
    """
    Args:
        file:
    """
    try:
        # No python-side buffering when reading the entire file to speedup reading
        # https://github.com/python/cpython/pull/122111
        with open(file, mode='rb', buffering=0) as f:
            return f.read()
    except FileNotFoundError:
        return b''


def file_read_bytes_stream(file: str, chunk_size: int = 8192) -> Iterable[bytes]:
    """
    Args:
        file:
        chunk_size:
    """
    try:
        with open(file, mode='rb') as f:
            while 1:
                chunk = f.read(chunk_size)
                if not chunk:
                    return
                yield chunk
    except FileNotFoundError:
        return


def atomic_read_text(
        file: str,
        encoding: str = 'utf-8',
        errors: str = 'strict'
) -> str:
    """
    Atomic file read with minimal IO operation

    Args:
        file:
        encoding:
        errors: 'strict', 'ignore', 'replace' and any other errors mode in open()
    """
    if IS_WINDOWS:
        # PermissionError on Windows if another process is replacing
        last_error = None
        for attempt in range(WINDOWS_MAX_ATTEMPT):
            try:
                return file_read_text(file, encoding=encoding, errors=errors)
            except PermissionError as e:
                last_error = e
                delay = windows_attempt_delay(attempt)
                time.sleep(delay)
                continue
        if last_error is not None:
            raise last_error from None
    else:
        # Linux and Mac allow reading while replacing
        return file_read_text(file, encoding=encoding, errors=errors)


def atomic_read_text_stream(
        file: str,
        encoding: str = 'utf-8',
        errors: str = 'strict',
        chunk_size: int = 8192
) -> Iterable[str]:
    """
    Args:
        file:
        encoding:
        errors: 'strict', 'ignore', 'replace' and any other errors mode in open()
        chunk_size:
    """
    if IS_WINDOWS:
        # PermissionError on Windows if another process is replacing
        last_error = None
        for attempt in range(WINDOWS_MAX_ATTEMPT):
            try:
                yield from file_read_text_stream(file, encoding=encoding, errors=errors, chunk_size=chunk_size)
                return
            except PermissionError as e:
                last_error = e
                delay = windows_attempt_delay(attempt)
                time.sleep(delay)
                continue
        if last_error is not None:
            raise last_error from None
    else:
        # Linux and Mac allow reading while replacing
        yield from file_read_text_stream(file, encoding=encoding, errors=errors, chunk_size=chunk_size)
        return


def atomic_read_bytes(file: str) -> bytes:
    """
    Atomic file read with minimal IO operation
    """
    if IS_WINDOWS:
        # PermissionError on Windows if another process is replacing
        last_error = None
        for attempt in range(WINDOWS_MAX_ATTEMPT):
            try:
                return file_read_bytes(file)
            except PermissionError as e:
                last_error = e
                delay = windows_attempt_delay(attempt)
                time.sleep(delay)
                continue
        if last_error is not None:
            raise last_error from None
    else:
        # Linux and Mac allow reading while replacing
        return file_read_bytes(file)


def atomic_read_bytes_stream(file: str, chunk_size: int = 8192) -> Iterable[bytes]:
    """
    Args:
        file:
        chunk_size:
    """
    if IS_WINDOWS:
        # PermissionError on Windows if another process is replacing
        last_error = None
        for attempt in range(WINDOWS_MAX_ATTEMPT):
            try:
                yield from file_read_bytes_stream(file, chunk_size=chunk_size)
                return
            except PermissionError as e:
                last_error = e
                delay = windows_attempt_delay(attempt)
                time.sleep(delay)
                continue
        if last_error is not None:
            raise last_error from None
    else:
        # Linux and Mac allow reading while replacing
        yield from file_read_bytes_stream(file, chunk_size=chunk_size)
        return


def file_remove(file: str):
    """
    Remove a file non-atomic
    """
    try:
        os.unlink(file)
    except FileNotFoundError:
        # If file not exist, just no need to remove
        pass


def atomic_remove(file: str):
    """
    Atomic file remove

    Args:
        file:
    """
    if IS_WINDOWS:
        # PermissionError on Windows if another process is replacing
        last_error = None
        for attempt in range(WINDOWS_MAX_ATTEMPT):
            try:
                return file_remove(file)
            except PermissionError as e:
                last_error = e
                delay = windows_attempt_delay(attempt)
                time.sleep(delay)
                continue
        if last_error is not None:
            raise last_error from None
    else:
        # Linux and Mac allow deleting while another process is reading
        # The directory entry is removed but the storage allocated to the file is not made available
        # until the original file is no longer in use.
        return file_remove(file)


def folder_rmtree(folder, may_symlinks=True):
    """
    Recursively remove a folder and its content

    Args:
        folder:
        may_symlinks: Default to True
            False if you already know it's not a symlink

    Returns:
        bool: If success
    """
    try:
        # If it's a symlinks, unlink it
        if may_symlinks and os.path.islink(folder):
            file_remove(folder)
            return True
        # Iter folder
        with os.scandir(folder) as entries:
            for entry in entries:
                if entry.is_dir(follow_symlinks=False):
                    folder_rmtree(entry.path, may_symlinks=False)
                else:
                    # File or symlink
                    # Just remove the symlink, not what it points to
                    try:
                        file_remove(entry.path)
                    except PermissionError:
                        # Another process is reading/writing
                        pass

    except FileNotFoundError:
        # directory to clean up does not exist, no need to clean up
        return True
    except NotADirectoryError:
        file_remove(folder)
        return True

    # Remove empty folder
    # May raise OSError if it's still not empty
    try:
        os.rmdir(folder)
        return True
    except FileNotFoundError:
        return True
    except NotADirectoryError:
        file_remove(folder)
        return True
    except OSError:
        return False


def atomic_rmtree(folder: str):
    """
    Atomic folder rmtree
    Rename folder as temp folder and remove it,
    folder can be removed by atomic_failure_cleanup at next startup if remove gets interrupted
    """
    temp = to_tmp_file(folder)
    try:
        atomic_replace(folder, temp)
    except FileNotFoundError:
        # Folder not exist, no need to rmtree
        return
    folder_rmtree(temp)


def atomic_failure_cleanup(folder: str, recursive: bool = False):
    """
    Cleanup remaining temp file under given path.
    In most cases there should be no remaining temp files unless write process get interrupted.

    This method should only be called at startup
    to avoid deleting temp files that another process is writing.
    """
    try:
        with os.scandir(folder) as entries:
            for entry in entries:
                if is_tmp_file(entry.name):
                    try:
                        # Delete temp file or directory
                        if entry.is_dir(follow_symlinks=False):
                            folder_rmtree(entry.path, may_symlinks=False)
                        else:
                            file_remove(entry.path)
                    except PermissionError:
                        # Another process is reading/writing
                        pass
                    except:
                        pass
                else:
                    if recursive:
                        try:
                            if entry.is_dir(follow_symlinks=False):
                                # Normal directory
                                atomic_failure_cleanup(entry.path, recursive=True)
                        except:
                            pass

    except FileNotFoundError:
        # directory to clean up does not exist, no need to clean up
        pass
    except NotADirectoryError:
        file_remove(folder)
    except:
        # Ignore all failures, it doesn't matter if tmp files still exist
        pass