"""High-level upload helpers for Concentriq."""

from __future__ import annotations

import os
from concurrent.futures import ThreadPoolExecutor
from typing import IO, TYPE_CHECKING, Any, Union

import httpx

if TYPE_CHECKING:
    from concentriq.client import ConcentriqClient

# Default multipart part size: 15 MB
DEFAULT_PART_SIZE = 15 * 1024 * 1024


def upload_image(
    client: ConcentriqClient,
    file_path_or_obj: Union[str, os.PathLike[str], IO[bytes]],
    image_set_id: int,
    file_name: str | None = None,
    file_size: int | None = None,
    part_size: int = DEFAULT_PART_SIZE,
    max_concurrency: int = 4,
) -> dict[str, Any]:
    """Create an image record and upload the file in one call.

    For small files (< part_size), uses a single presigned PUT.
    For large files, uses S3 multipart upload with presigned URLs per part.

    Args:
        client: ConcentriqClient instance
        file_path_or_obj: File path string or file-like object (opened in binary mode)
        image_set_id: Repository ID to upload into
        file_name: Filename (required if file_path_or_obj is a file object)
        file_size: File size in bytes (required if file_path_or_obj is a file object)
        part_size: Part size for multipart upload in bytes (default: 15 MB)
        max_concurrency: Maximum number of concurrent part uploads (default: 4)

    Returns:
        Created image data dict (contains at minimum {"id": ...})

    Raises:
        ValueError: If file_name or file_size not provided for file objects
        AuthenticationError: If not authenticated
        ValidationError: If image data is invalid

    Example:
        # Upload from file path
        image = upload_image(client, "/path/to/slide.tif", image_set_id=100)

        # Upload from file object
        with open("slide.tif", "rb") as f:
            image = upload_image(
                client, f, image_set_id=100,
                file_name="slide.tif", file_size=os.path.getsize("slide.tif")
            )
    """
    # Resolve file path vs file object
    if isinstance(file_path_or_obj, (str, os.PathLike)):
        path = str(file_path_or_obj)
        if file_name is None:
            file_name = os.path.basename(path)
        if file_size is None:
            file_size = os.path.getsize(path)
        file_obj: IO[bytes] = open(path, "rb")  # noqa: SIM115
        should_close = True
    else:
        file_obj = file_path_or_obj
        should_close = False
        if file_name is None:
            raise ValueError("file_name is required when passing a file object")
        if file_size is None:
            raise ValueError("file_size is required when passing a file object")

    try:
        return _do_upload(client, file_obj, file_name, file_size, image_set_id, part_size, max_concurrency)
    finally:
        if should_close:
            file_obj.close()


def _do_upload(
    client: ConcentriqClient,
    file_obj: IO[bytes],
    file_name: str,
    file_size: int,
    image_set_id: int,
    part_size: int,
    max_concurrency: int = 4,
) -> dict[str, Any]:
    """Internal: create image record and upload file data."""
    # Step 1: Create image record
    image = client.v1.images.create({
        "name": file_name,
        "size": file_size,
        "imageSetId": image_set_id,
    })
    image_id = image["id"]

    # Step 2: Get storageKey from the created image
    image_data = client.v1.images.get(image_id)
    storage_key = image_data["storageKey"]

    # Step 3: Upload the file
    if file_size <= part_size:
        _single_part_upload(client, file_obj, file_size, image_id, storage_key)
    else:
        _multipart_upload(client, file_obj, file_size, image_id, storage_key, part_size, max_concurrency)

    return image


def _single_part_upload(
    client: ConcentriqClient,
    file_obj: IO[bytes],
    file_size: int,
    image_id: int,
    storage_key: str,
) -> None:
    """Upload file using a single presigned PUT."""
    url = client.v3.files.get_upload_url(
        resource_type="Image",
        resource_id=image_id,
        storage_key=storage_key,
    )

    data = file_obj.read()
    httpx.put(url, content=data, timeout=300.0)


def _upload_one_part(
    client: ConcentriqClient,
    chunk: bytes,
    image_id: int,
    storage_key: str,
    upload_id: str,
    part_number: int,
) -> dict[str, Any]:
    """Upload a single part and return its PartNumber + ETag."""
    url = client.v3.files.get_part_upload_url(
        resource_type="Image",
        resource_id=image_id,
        storage_key=storage_key,
        upload_id=upload_id,
        part_number=part_number,
    )
    response = httpx.put(url, content=chunk, timeout=300.0)
    etag = response.headers.get("ETag", "")
    return {"PartNumber": part_number, "ETag": etag}


def _multipart_upload(
    client: ConcentriqClient,
    file_obj: IO[bytes],
    file_size: int,
    image_id: int,
    storage_key: str,
    part_size: int,
    max_concurrency: int = 4,
) -> None:
    """Upload file using S3 multipart upload with presigned URLs.

    Parts are uploaded concurrently using a thread pool.
    """
    # Initiate multipart upload
    upload_info = client.v3.files.create_multipart_upload(
        resource_type="Image",
        resource_id=image_id,
        storage_key=storage_key,
    )
    upload_id = upload_info["UploadId"]

    try:
        # Read all chunks upfront so they can be dispatched to threads
        chunks: list[tuple[int, bytes]] = []
        part_number = 1
        bytes_remaining = file_size

        while bytes_remaining > 0:
            chunk_size = min(part_size, bytes_remaining)
            chunk = file_obj.read(chunk_size)
            chunks.append((part_number, chunk))
            part_number += 1
            bytes_remaining -= chunk_size

        # Upload parts concurrently
        with ThreadPoolExecutor(max_workers=max_concurrency) as executor:
            futures = [
                executor.submit(
                    _upload_one_part,
                    client, chunk, image_id, storage_key, upload_id, pn,
                )
                for pn, chunk in chunks
            ]
            parts = [f.result() for f in futures]

        # Sort by part number to ensure correct order
        parts.sort(key=lambda p: p["PartNumber"])

        # Complete the multipart upload
        client.v3.files.complete_multipart_upload(
            resource_type="Image",
            resource_id=image_id,
            storage_key=storage_key,
            upload_id=upload_id,
            parts=parts,
        )
    except Exception:
        # Abort on failure
        client.v3.files.abort_multipart_upload(
            resource_type="Image",
            resource_id=image_id,
            storage_key=storage_key,
            upload_id=upload_id,
        )
        raise
