"""
ObjStorage 封装层
底层依赖 smartlibs 的 smart.utils.storage 抽象，提供文件列表、上传、下载能力。
需要安装 minio 可选依赖: pip install super-datascope[minio]
"""

import os
import logging
import tempfile
from pathlib import Path
from abc import abstractmethod

from smart.utils.tuple import tuple_fixed_len
from smart.utils.path import path_join
from smart.utils.loader import dyn_import
from smart.utils.storage.obj_factory import ObjStorageFactory
from smart.utils.storage.base import BaseStorageClient
from smart.utils.storage.obj_storage import ObjStorage
from smart.utils.storage.minio_storage import MinioClient
from smart.utils.storage.local_storage import LocalStorageClient

logger = logging.getLogger("super_datascope")


# ---------------------------------------------------------------------------
# Extended storage clients with list_files / fput_object
# ---------------------------------------------------------------------------

class ExtendBaseStorageClient(BaseStorageClient):
    @abstractmethod
    def list_files(self, prefix: str = '', recursive: bool = False):
        pass

    @abstractmethod
    def fput_object(self, local_file_path: str, object_name: str):
        pass


class ExtendLocalStorageClient(LocalStorageClient):
    def list_files(self, prefix: str = '', recursive: bool = False):
        if '/' in prefix:
            bucket_name, obj_prefix = prefix.split('/', 1)
        else:
            bucket_name, obj_prefix = prefix, ''

        dir_path = self.get_bucket_dir(bucket_name)
        base_path = Path(dir_path)
        target_path = base_path / obj_prefix if obj_prefix else base_path

        if not target_path.exists():
            return []

        files = []
        if recursive:
            for fp in target_path.rglob('*'):
                if fp.is_file():
                    files.append(str(fp.relative_to(base_path)))
        else:
            for fp in target_path.iterdir():
                rel = str(fp.relative_to(base_path))
                files.append(rel + '/' if fp.is_dir() else rel)
        return files

    def fput_object(self, local_file_path: str, object_name: str):
        import shutil
        if '/' in object_name:
            bucket_name, obj_path = object_name.split('/', 1)
        else:
            raise ValueError("Object path cannot be empty")

        bucket_dir = self.get_bucket_dir(bucket_name)
        target = Path(bucket_dir) / obj_path
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(local_file_path, target)


class ExtendMinioClient(MinioClient):
    def list_files(self, prefix: str = '', recursive: bool = False):
        if '/' in prefix:
            bucket_name, obj_prefix = prefix.split('/', 1)
        else:
            bucket_name, obj_prefix = prefix, ''

        files = []
        try:
            objects = self._minio.list_objects(bucket_name, prefix=obj_prefix, recursive=recursive)
            for obj in objects:
                name = obj.object_name
                if name.startswith(obj_prefix):
                    rel = name[len(obj_prefix):].lstrip('/')
                    if rel:
                        files.append(rel)
                else:
                    files.append(name)
        except Exception as e:
            logger.error(f"MinIO list_files error: {e}")
        return files

    def fput_object(self, local_file_path: str, object_name: str):
        if '/' in object_name:
            bucket_name, obj_path = object_name.split('/', 1)
        else:
            raise ValueError("Object path cannot be empty")
        self._minio.fput_object(bucket_name, obj_path, local_file_path)


# ---------------------------------------------------------------------------
# Extended ObjStorage / Factory
# ---------------------------------------------------------------------------

def _create_client(type_name: str, connection: dict, **kwargs) -> ExtendBaseStorageClient:
    if type_name == 'local':
        return ExtendLocalStorageClient(connection=connection, **kwargs)
    if type_name == 'minio':
        return ExtendMinioClient(connection=connection, **kwargs)
    client_cls = dyn_import(type_name)
    if issubclass(client_cls, BaseStorageClient):
        return client_cls(connection=connection, **kwargs)
    raise ValueError(f"Unsupported StorageClient type: {type_name}")


class ExtendObjStorage(ObjStorage):
    def list_files(self, prefix: str = '', recursive: bool = False):
        full_prefix = self._bucket_name or ''
        if self._bucket_root_path:
            full_prefix = f"{full_prefix}/{self._bucket_root_path.lstrip('/')}"
        if prefix:
            full_prefix = f"{full_prefix}/{prefix.lstrip('/')}"
        client: ExtendBaseStorageClient = self._client
        return client.list_files(full_prefix, recursive=recursive)

    def fput_object(self, local_file_path: str, object_name: str):
        full_name = self._bucket_name or ''
        if self._bucket_root_path:
            full_name = f"{full_name}/{self._bucket_root_path.lstrip('/')}"
        if object_name:
            full_name = f"{full_name}/{object_name.lstrip('/')}"
        client: ExtendBaseStorageClient = self._client
        client.fput_object(local_file_path, full_name)


class ExtendObjStorageFactory(ObjStorageFactory):
    def create_obj_store_by_config(self, config: dict, connection: dict = None,
                                   cache_dir: str = None) -> ExtendObjStorage:
        store_type = config.get('type') or ObjStorageFactory.Default_Store_Type
        conn_kwargs = connection or config.get('connection')
        bucket_name = config.get('bucket_name')
        bucket_root_path = config.get('bucket_root_path')
        cache_dir = cache_dir or config.get('cache_dir')
        assert conn_kwargs and bucket_name
        client = _create_client(type_name=store_type, connection=conn_kwargs)
        return ExtendObjStorage(
            client=client, cache_dir=cache_dir,
            bucket_name=bucket_name, bucket_root_path=bucket_root_path
        )


# ---------------------------------------------------------------------------
# Config-based factory (reads ~/.super-datascope/config.yaml or env vars)
# ---------------------------------------------------------------------------

def _load_storage_config() -> dict:
    """Load storage config from ~/.super-datascope/config.yaml"""
    config_path = Path.home() / '.super-datascope' / 'config.yaml'
    if config_path.exists():
        try:
            from smart.utils.yaml import yaml_load
            cfg = yaml_load(str(config_path))
            return cfg.get('storage', {})
        except Exception as e:
            logger.warning(f"Failed to load config: {e}")
    return {}


def _load_storage_config_from_env() -> dict | None:
    """Load storage config from environment variables"""
    stype = os.environ.get('DATASCOPE_STORAGE_TYPE')
    if not stype:
        return None
    return {
        'env_storage': {
            'type': stype,
            'connection': {
                'endpoint': os.environ.get('DATASCOPE_STORAGE_ENDPOINT', ''),
                'access_key': os.environ.get('DATASCOPE_STORAGE_ACCESS_KEY', ''),
                'secret_key': os.environ.get('DATASCOPE_STORAGE_SECRET_KEY', ''),
                'secure': os.environ.get('DATASCOPE_STORAGE_SECURE', 'false').lower() == 'true',
            },
            'bucket_name': os.environ.get('DATASCOPE_STORAGE_BUCKET', ''),
            'bucket_root_path': os.environ.get('DATASCOPE_STORAGE_ROOT_PATH', ''),
        }
    }


def get_storage_configs() -> dict:
    """Get all configured storage backends (merged from file + env)"""
    configs = _load_storage_config()
    env_cfg = _load_storage_config_from_env()
    if env_cfg:
        configs.update(env_cfg)
    return configs


def create_storage(name: str) -> ExtendObjStorage | None:
    """Create an ExtendObjStorage instance by config name"""
    configs = get_storage_configs()
    cfg = configs.get(name)
    if not cfg:
        logger.warning(f"Storage config not found: {name}")
        return None
    factory = ExtendObjStorageFactory()
    return factory.create_obj_store_by_config(cfg)


# ---------------------------------------------------------------------------
# Convenience singleton
# ---------------------------------------------------------------------------

class _ObjStorages:
    cache_root_dir = './logs'

    def get_temp_dir(self, store_name: str = None) -> str:
        if store_name:
            temp_base = os.path.join(tempfile.gettempdir(), 'obj_storage', store_name)
        else:
            temp_base = os.path.join(tempfile.gettempdir(), 'obj_storage')
        os.makedirs(temp_base, exist_ok=True)
        return temp_base

    def get_file_path(self, obj_file_path: str, use_cache: bool = True,
                      local_file_path: str = None, **kwargs) -> str:
        """Download obj:// path to local and return local path"""
        if not obj_file_path or not obj_file_path.startswith('obj://'):
            return obj_file_path
        store_name, obj_path = tuple_fixed_len(obj_file_path[6:].split('/', 1), fix_len=2)
        assert obj_path
        storage = create_storage(store_name)
        if storage is None:
            raise ValueError(f"Storage not configured: {store_name}")
        if not storage._cache_dir:
            storage._cache_dir = path_join(self.cache_root_dir, store_name)
        return storage.fget(obj_path, use_cache=use_cache,
                            local_file_path=local_file_path, **kwargs)

    def get_file_list(self, obj_path: str, file_types: list = None,
                      max_depth: int = 1) -> list[str]:
        """List files under an obj:// path"""
        if not obj_path or not obj_path.startswith('obj://'):
            return []
        store_name, obj_prefix = tuple_fixed_len(obj_path[6:].split('/', 1), fix_len=2)
        storage = create_storage(store_name)
        if storage is None:
            return []

        file_list = []
        if max_depth == -1:
            for fp in storage.list_files(obj_prefix, recursive=True):
                fp_str = fp if isinstance(fp, str) else fp.get('key', '')
                if file_types:
                    ext = fp_str.rsplit('.', 1)[-1].lower() if '.' in fp_str else ''
                    if ext not in file_types:
                        continue
                if obj_prefix and not fp_str.startswith(obj_prefix):
                    full = f"obj://{store_name}/{obj_prefix.lstrip('/')}/{fp_str.lstrip('/')}"
                else:
                    full = f"obj://{store_name}/{fp_str.lstrip('/')}"
                file_list.append(full)
        else:
            prefixes = [obj_prefix] if obj_prefix else ['']
            for depth in range(max_depth):
                next_prefixes = []
                for prefix in prefixes:
                    try:
                        for fp in storage.list_files(prefix, recursive=False):
                            fp_str = fp if isinstance(fp, str) else fp.get('key', '')
                            if fp_str.endswith('/'):
                                if depth + 1 < max_depth:
                                    next_prefixes.append(fp_str)
                            else:
                                if file_types:
                                    ext = fp_str.rsplit('.', 1)[-1].lower() if '.' in fp_str else ''
                                    if ext not in file_types:
                                        continue
                                if prefix and not fp_str.startswith(prefix):
                                    full = f"obj://{store_name}/{prefix.lstrip('/')}/{fp_str.lstrip('/')}"
                                else:
                                    full = f"obj://{store_name}/{fp_str.lstrip('/')}"
                                file_list.append(full)
                    except Exception:
                        pass
                prefixes = next_prefixes
        return file_list


ObjStorages = _ObjStorages()
