import atexit
import gc
import logging
import multiprocessing
import os
from typing import Any, Optional, List

import tosnativeclient
from tosnativeclient import WriteStream, TosObject, TosBucket, ReadStream, TosInode
from .version import get_project_version

log = logging.getLogger(__name__)

import threading
import weakref
import traceback

_client_lock = threading.Lock()
_client_map = weakref.WeakSet()


def _before_fork():
    with _client_lock:
        clients = list(_client_map)

    if not clients or len(clients) == 0:
        return

    try:
        for client in clients:
            if client._inner_client is not None:
                if hasattr(client._inner_client, 'close') and callable(client._inner_client.close):
                    client._inner_client.close()
                client._inner_client = None
        _reset_client_map()
        gc.collect()
    except Exception as e:
        log.warning(f'failed to clean up native clients before fork, {str(e)}')
        traceback.print_exc()


def _after_fork_in_child():
    _reset_client_map()


def _reset_client_map():
    global _client_map
    with _client_lock:
        _client_map = weakref.WeakSet()


os.register_at_fork(before=_before_fork, after_in_child=_after_fork_in_child)


def tear_down():
    with _client_lock:
        clients = list(_client_map)

    if not clients or len(clients) == 0:
        return

    try:
        for client in clients:
            if client._inner_client is not None:
                if hasattr(client._inner_client, 'close') and callable(client._inner_client.close):
                    client._inner_client.close()
                client._inner_client = None
    except Exception as e:
        log.warning(f'failed to clean up native clients before atexit, {str(e)}')
        traceback.print_exc()


atexit.register(tear_down)


class TosClient(object):
    def __init__(self, region: str, endpoint: str = '', ak='', sk='',
                 part_size: int = 8 * 1024 * 1024, max_retry_count: int = 3, shared_prefetch_tasks: int = 32,
                 shared_upload_part_tasks: int = 32, shared_upload_part_copy_tasks: int = 32, enable_crc: bool = True,
                 connection_timeout: int = 10000, request_timeout: int = 120000, security_token: Optional[str] = None):
        self._region = region
        self._endpoint = endpoint
        self._ak = ak
        self._sk = sk
        self._part_size = part_size
        self._max_retry_count = max_retry_count
        self._shared_prefetch_tasks = shared_prefetch_tasks
        self._shared_upload_part_tasks = shared_upload_part_tasks
        self._shared_upload_part_copy_tasks = shared_upload_part_copy_tasks
        self._enable_crc = enable_crc
        self._connection_timeout = connection_timeout
        self._request_timeout = request_timeout
        self._user_agent_conf = ('tosfsspec', get_project_version())
        self._security_token = security_token
        self._inner_client = None
        self._client_pid = None

    def is_fork(self) -> bool:
        return multiprocessing.get_start_method() == 'fork'

    @property
    def _client(self) -> Any:
        if self._client_pid is None or self._client_pid != os.getpid() or self._inner_client is None:
            with _client_lock:
                if self._client_pid is None or self._client_pid != os.getpid() or self._inner_client is None:
                    if self._inner_client is not None:
                        if hasattr(self._inner_client, 'close') and callable(self._inner_client.close):
                            self._inner_client.close()
                    self._inner_client = tosnativeclient.TosClient(self._region, self._endpoint, self._ak,
                                                                   self._sk,
                                                                   self._part_size,
                                                                   self._max_retry_count,
                                                                   shared_prefetch_tasks=self._shared_prefetch_tasks,
                                                                   shared_upload_part_tasks=self._shared_upload_part_tasks,
                                                                   shared_upload_part_copy_tasks=self._shared_upload_part_copy_tasks,
                                                                   enable_crc=self._enable_crc,
                                                                   connection_timeout=self._connection_timeout,
                                                                   request_timeout=self._request_timeout,
                                                                   user_agent_conf=self._user_agent_conf,
                                                                   dns_cache_async_refresh=False,
                                                                   security_token=self._security_token)
                    self._client_pid = os.getpid()
                    _client_map.add(self)

        assert self._inner_client is not None
        return self._inner_client

    def put_object(self, bucket: str, key: str, storage_class: Optional[str] = None,
                   write_part_size: Optional[int] = None, content_type: Optional[str] = None) -> WriteStream:
        return self._client.put_object(bucket, key, storage_class=storage_class, write_part_size=write_part_size,
                                       content_type=content_type)

    def append_object(self, bucket: str, key: str, etag: str, size: int, crc64: int,
                      storage_class: Optional[str] = None, write_part_size: Optional[int] = None,
                      content_type: Optional[str] = None) -> WriteStream:
        return self._client.append_object(bucket, key, etag, size, crc64, storage_class=storage_class,
                                          write_part_size=write_part_size, content_type=content_type)

    def get_object(self, bucket: str, key: str, etag: Optional[str] = None,
                   size: Optional[int] = None, crc64: Optional[int] = None, read_part_size: Optional[int] = None,
                   preload: bool = False, with_sliding_window: bool = True) -> ReadStream:

        if (size is None or etag is None or size == -1) and not preload:
            tos_object = self.head_object(bucket, key)
            etag = tos_object.etag
            size = tos_object.size
            crc64 = tos_object.crc64
        return self._client.get_object(bucket, key, etag=etag, size=size, crc64=crc64, read_part_size=read_part_size,
                                       with_sliding_window=with_sliding_window, preload=preload)

    def copy_object(self, bucket: str, key: str, src_bucket: str, src_key: str, copy_metadata: bool = False,
                    recursive: bool = False) -> None:
        self._client.copy(bucket, key, src_bucket=src_bucket, src_key=src_key, copy_metadata=copy_metadata,
                          recursive=recursive)

    def list_buckets(self, project_name: Optional[str] = None, bucket_type: Optional[str] = None) -> List[TosBucket]:
        return self._client.list_buckets(project_name, bucket_type)

    def list_objects(self, bucket: str, prefix: str, max_keys: int = 1000,
                     delimiter: Optional[str] = None,
                     continuation_token: Optional[str] = None,
                     start_after: Optional[str] = None,
                     delimiter_recursive: bool = False) -> tosnativeclient.ListStream:

        delimiter = delimiter if delimiter is not None else ''
        continuation_token = continuation_token if continuation_token is not None else ''
        start_after = start_after if start_after is not None else ''
        return self._client.list_objects(bucket, prefix, max_keys=max_keys, delimiter=delimiter,
                                         continuation_token=continuation_token, start_after=start_after,
                                         delimiter_recursive=delimiter_recursive)

    def head_bucket(self, bucket: str) -> TosBucket:
        return self._client.head_bucket(bucket)

    def create_bucket(self, bucket: str, bucket_type: Optional[str] = None) -> None:
        return self._client.create_bucket(bucket, bucket_type=bucket_type)

    def delete_bucket(self, bucket: str) -> None:
        return self._client.delete_bucket(bucket)

    def get_bucket_type(self, bucket: str) -> str:
        return self._client.get_bucket_type(bucket)

    def head_object(self, bucket: str, key: str) -> TosObject:
        return self._client.head_object(bucket, key)

    def delete_object(self, bucket: str, key: str) -> None:
        self._client.delete_object(bucket, key)

    def delete_folder(self, bucket: str, key: str, recursive: bool = False) -> None:
        self._client.delete_folder(bucket, folder=key, recursive=recursive)

    def walk(self, bucket: str, prefix: str, max_keys: int = 1000,
             maxdepth: Optional[int] = None) -> tosnativeclient.WalkStream:

        return self._client.walk(bucket, prefix, max_keys=max_keys, max_depth=maxdepth)

    def lookup(self, bucket: str, key: str) -> Optional[TosInode]:
        return self._client.lookup(bucket, key)

    def mkdir(self, bucket: str, key: str, create_parents: bool = False, check_exists: bool = False) -> None:
        try:
            return self._client.mkdir(bucket, key, create_parents=create_parents, check_exists=check_exists,
                                      content_type='application/x-directory')
        except tosnativeclient.TosException as ex:
            # CHECK_EXISTS_ERROR
            if len(ex.args) > 0 and ex.args[0].ec == '-1':
                raise FileExistsError(ex.args[0].message)
            raise

    def rmdir(self, bucket: str, key: str, recursive: bool = False) -> None:
        self._client.delete_folder(bucket, folder=key, recursive=recursive)

    def rm(self, bucket: str, key: str, recursive: bool = False) -> None:
        self._client.delete(bucket, key=key, recursive=recursive)

    def copy(self, bucket: str, key: str, src_key: str, recursive: bool = False,
             copy_metadata: bool = False, src_bucket: Optional[str] = None, maxdepth: Optional[int] = None,
             with_src_parent: bool = False) -> None:
        try:
            return self._client.copy(bucket, key, src_key=src_key, recursive=recursive,
                                     copy_metadata=copy_metadata, src_bucket=src_bucket, max_depth=maxdepth,
                                     with_src_parent=with_src_parent)
        except tosnativeclient.TosException as ex:
            # COPY_FOLDER_IS_NOT_ALLOWED
            if not recursive and len(ex.args) > 0 and ex.args[0].ec == '-2':
                # do nothing
                return None
            raise

    def rename(self, bucket: str, key: str, src_key: str, recursive: bool = False,
               copy_metadata: bool = False, src_bucket: Optional[str] = None) -> None:
        return self._client.rename(bucket, key, src_key=src_key, recursive=recursive,
                                   copy_metadata=copy_metadata, src_bucket=src_bucket)

    def touch(self, bucket: str, key: str, content_type: Optional[str] = None) -> str:
        return self._client.touch(bucket, key, content_type=content_type)

    def close(self):
        self._client.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        if exc_type is not None:
            try:
                log.info(f'exception occurred before closing tos client: {exc_type.__name__}: {exc_val}')
            except Exception:
                pass
            finally:
                self.close()
        else:
            self.close()
