import logging
from pathlib import Path
from typing import Any, Protocol

import requests

logger = logging.getLogger(__name__)


class SupportsGetObject(Protocol):
    def get_object(self, object_key: str, **kwargs: Any) -> requests.Response: ...


def download_object(
    client: SupportsGetObject,
    object_key: str,
    dest_path: str | Path,
    chunk_size: int = 1024 * 1024,
):
    with requests.Session() as session:
        return download_object_with_session(
            client,
            object_key,
            dest_path,
            chunk_size=chunk_size,
            session=session,
        )


def download_object_with_session(
    client: SupportsGetObject,
    object_key: str,
    dest_path: str | Path,
    chunk_size: int = 1024 * 1024,
    session: requests.Session | None = None,
) -> Path:
    dest_path = Path(dest_path)
    dest_path.parent.mkdir(parents=True, exist_ok=True)

    if not dest_path.exists():
        logger.info("Starting full download: %s -> %s", object_key, dest_path)
        response = client.get_object(object_key, session=session, stream=True)
        response.raise_for_status()
        _write_response_to_path(
            response,
            dest_path,
            mode="wb",
            chunk_size=chunk_size,
        )
        return dest_path

    local_size = dest_path.stat().st_size
    probe_response = client.get_object(
        object_key,
        session=session,
        headers={"Range": "bytes=0-0"},
        stream=True,
    )
    probe_response.raise_for_status()
    remote_size = _extract_total_size(probe_response)
    probe_response.close()

    if local_size == remote_size:
        logger.info("Local file already complete, skipping: %s", dest_path)
        return dest_path

    if local_size > remote_size:
        logger.info("Local file larger than remote, restarting: %s", dest_path)
        response = client.get_object(object_key, session=session, stream=True)
        response.raise_for_status()
        _write_response_to_path(
            response,
            dest_path,
            mode="wb",
            chunk_size=chunk_size,
        )
        return dest_path

    logger.info("Resuming download: %s -> %s", object_key, dest_path)
    response = client.get_object(
        object_key,
        session=session,
        headers={"Range": f"bytes={local_size}-"},
        stream=True,
    )
    response.raise_for_status()
    if response.status_code != 206:
        logger.warning(
            "Range request not honored for '%s'; restarting full download",
            object_key,
        )
        _write_response_to_path(
            response,
            dest_path,
            mode="wb",
            chunk_size=chunk_size,
        )
        return dest_path
    _write_response_to_path(
        response,
        dest_path,
        mode="ab",
        chunk_size=chunk_size,
    )
    return dest_path


def _extract_total_size(response: requests.Response) -> int:
    content_range = response.headers.get("Content-Range")
    if content_range and "/" in content_range:
        total_size = content_range.rsplit("/", 1)[-1]
        if total_size.isdigit():
            return int(total_size)

    content_length = response.headers.get("Content-Length")
    if content_length and content_length.isdigit():
        return int(content_length)

    raise RuntimeError("无法从响应头中解析远端文件总大小")


def _write_response_to_path(
    response: requests.Response,
    dest_path: Path,
    mode: str,
    chunk_size: int,
) -> None:
    try:
        with dest_path.open(mode) as fp:
            for chunk in response.iter_content(chunk_size=chunk_size):
                if chunk:
                    fp.write(chunk)
    finally:
        response.close()
