# storage/init_storage.py
"""初始化 RustFS 对象存储"""
import asyncio
import sys
from pathlib import Path

# 添加项目根目录和 src 目录到 Python 路径
project_root = Path(__file__).parent.parent.parent
src_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(src_root))

from storage.rustfs_storage import rustfs_manager
from config.settings import settings
from utils.logger import get_logger

logger = get_logger(__name__)


async def init_rustfs_storage():
    """初始化 RustFS 对象存储"""
    try:
        # 连接到 RustFS (S3v4 API)
        await rustfs_manager.connect(
            endpoint=settings.RUSTFS_ENDPOINT,
            access_key=settings.RUSTFS_ACCESS_KEY,
            secret_key=settings.RUSTFS_SECRET_KEY,
            timeout=settings.RUSTFS_TIMEOUT
        )

        logger.info("RustFS 对象存储初始化完成")
        return True

    except Exception as e:
        logger.error(f"RustFS 对象存储初始化失败: {e}")
        return False


async def test_storage():
    """测试 RustFS 对象存储功能"""
    try:
        import json

        # 测试上传 JSON
        test_data = {"test": True, "timestamp": "2024-01-01", "storage": "rustfs"}
        result = await rustfs_manager.upload_json_data(
            file_type='stp_files',
            json_data=test_data,
            file_hash='test-hash'
        )

        logger.info(f"RustFS 测试上传成功: {result['object_key']}")

        # 测试下载
        downloaded_bytes = await rustfs_manager.download_file(
            file_type='stp_files',
            object_key=result['object_key']
        )
        downloaded_data = json.loads(downloaded_bytes.decode('utf-8'))
        logger.info(f"RustFS 测试下载成功: {downloaded_data}")

        # 测试预签名 URL
        url = await rustfs_manager.generate_presigned_url(
            file_type='stp_files',
            object_key=result['object_key'],
            expires=3600
        )
        logger.info(f"RustFS 预签名URL: {url}")

        # 清理测试文件
        await rustfs_manager.delete_file(
            file_type='stp_files',
            object_key=result['object_key']
        )
        logger.info("RustFS 测试文件已清理")

        return True

    except Exception as e:
        logger.error(f"RustFS 存储测试失败: {e}")
        return False


if __name__ == "__main__":
    async def main():
        try:
            print("=== 初始化 RustFS 对象存储 ===")
            # 初始化存储
            init_result = await init_rustfs_storage()
            if init_result:
                print("[OK] RustFS 连接成功")

            print("\n=== 测试 RustFS 功能 ===")
            # 运行测试
            test_result = await test_storage()
            if test_result:
                print("[OK] RustFS 测试全部通过")
            else:
                print("[FAIL] RustFS 测试失败")

        finally:
            # 关闭连接
            await rustfs_manager.close()
            print("\n=== 连接已关闭 ===")

    # 运行主函数
    asyncio.run(main())