import os
from pathlib import Path

from minio import Minio

from ldnode import logger
from ldnode.loader.schemas.enumerations import SyncStatus
from ldnode.loader.schemas.models import DataSource, SyncResult


class S3QuotationDataSource(DataSource):
    """
    S3 行情后台预拉取数据源。

    由 LoaderWorker 注册为 LoadRule（每 3s，仅交易时段），替代原本地行情计算逻辑。

    使用 MinIO 客户端全量 GET 下载整个文件，
    绕过 polars scan_parquet 的 HTTP Range Request（RustFS alpha 对 >1MB 文件的
    range read 存在 bug: read_file_stream offset+length 计算越界）。
    output_lazyframe() 继续读本地文件，选股请求无 S3 网络延迟。
    """

    def __init__(self, s3_uri: str, storage_options: dict, local_data_dir: Path):
        self.local_data_dir = local_data_dir
        opts = storage_options or {}
        # s3://bucket/key → bucket, key
        stripped = s3_uri.replace("s3://", "")
        slash = stripped.index("/")
        self._bucket = stripped[:slash]
        self._object_key = stripped[slash + 1:]
        # 构建 MinIO 客户端
        endpoint = opts.get("endpoint_url", "")
        secure = not endpoint.startswith("http://")
        # MinIO 客户端 endpoint 不带 scheme
        host = endpoint.replace("http://", "").replace("https://", "")
        self._client = Minio(
            host,
            access_key=opts.get("aws_access_key_id", ""),
            secret_key=opts.get("aws_secret_access_key", ""),
            secure=secure,
            region=opts.get("aws_region", "us-east-1"),
        )

    def load(self, key: str) -> list[SyncResult]:
        """全量 GET 下载 S3 对象并写入本地文件"""
        local_path = self.local_data_dir / key
        try:
            resp = self._client.get_object(self._bucket, self._object_key)
            try:
                data = resp.read()
            finally:
                resp.close()
                resp.release_conn()
            # 原子写：先写临时文件，再 rename 替换，避免读写竞态
            tmp_path = local_path.with_suffix(".tmp")
            tmp_path.write_bytes(data)
            os.rename(tmp_path, local_path)
            logger.loader.debug(f"S3 行情已同步: {key}  size={len(data)}")
            return [SyncResult(key, SyncStatus.SUCCESS, f"size={len(data)}")]
        except Exception as e:
            logger.loader.warning(f"S3 行情同步失败: {key}  error={e}")
            return [SyncResult(key, SyncStatus.FAILED, str(e))]
