from __future__ import annotations

import json
import logging
from pathlib import Path

from botocore.config import Config as BotoConfig


logger = logging.getLogger(__name__)


class FinalExportR2Client:
    def __init__(self, config):
        self.config = config
        self.mock_mode = config.mock_mode
        if not self.mock_mode:
            import boto3

            self.s3 = boto3.client(
                "s3",
                endpoint_url=config.base.r2_endpoint_url,
                aws_access_key_id=config.base.r2_access_key_id,
                aws_secret_access_key=config.base.r2_secret_access_key,
                region_name="auto",
                config=BotoConfig(max_pool_connections=32),
            )

    def upload_file(self, local_path: Path, bucket: str, key: str) -> int:
        if self.mock_mode:
            logger.info("[MOCK] upload %s -> s3://%s/%s", local_path, bucket, key)
            return local_path.stat().st_size if local_path.exists() else 0
        logger.info("Uploading %s -> s3://%s/%s", local_path, bucket, key)
        self.s3.upload_file(str(local_path), bucket, key)
        return local_path.stat().st_size

    def download_file(self, bucket: str, key: str, local_path: Path):
        if self.mock_mode:
            raise RuntimeError("Mock mode does not support R2 download")
        local_path.parent.mkdir(parents=True, exist_ok=True)
        logger.info("Downloading s3://%s/%s -> %s", bucket, key, local_path)
        self.s3.download_file(bucket, key, str(local_path))

    def head_size(self, bucket: str, key: str) -> int:
        if self.mock_mode:
            return 0
        meta = self.s3.head_object(Bucket=bucket, Key=key)
        return int(meta.get("ContentLength", 0))

    def list_keys(self, bucket: str, prefix: str) -> list[str]:
        if self.mock_mode:
            return []
        paginator = self.s3.get_paginator("list_objects_v2")
        keys: list[str] = []
        for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
            for item in page.get("Contents", []):
                key = item.get("Key")
                if key:
                    keys.append(str(key))
        return keys

    def download_json(self, bucket: str, key: str) -> dict:
        if self.mock_mode:
            raise RuntimeError("Mock mode does not support R2 download")
        obj = self.s3.get_object(Bucket=bucket, Key=key)
        body = obj["Body"].read().decode("utf-8")
        return json.loads(body)