#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   multipart.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Multipart upload support for large files.
"""

import base64
import os
from collections.abc import Iterator
from concurrent.futures import Future, ThreadPoolExecutor, as_completed
from pathlib import Path
from queue import Empty, Queue
from threading import Lock
from typing import TYPE_CHECKING

import msgspec
from vi.api.resources.datasets.assets import responses
from vi.api.resources.datasets.assets.consts import (
    ASSET_FILE_EXTENSION_TO_MIME_TYPE_MAP,
    MULTIPART_UPLOAD_WORKERS,
)
from vi.api.resources.datasets.assets.links import MultipartUploadLinkParser
from vi.api.resources.datasets.assets.types import (
    FilePart,
    MultipartProgress,
    MultipartSession,
    MultipartSessionSpec,
    PartSignedUrl,
)
from vi.api.resources.datasets.assets.utils.helper import calculate_parts, read_part
from vi.logging import get_logger

if TYPE_CHECKING:
    from rich.progress import TaskID
    from vi.api.resources.datasets.assets.uploader import (
        AssetUploader,
        AssetUploadProgressTracker,
    )
    from vi.utils.graceful_exit import GracefulExit
    from vi.utils.progress import ViProgress

logger = get_logger("assets.multipart")


class ProgressTrackingContent:
    """Wrapper around bytes content that tracks upload progress and checks for interruption."""

    def __init__(
        self,
        content: bytes,
        progress: "ViProgress | None",
        task: "TaskID | None",
        handler: "GracefulExit | None",
        progress_tracker: "AssetUploadProgressTracker | None" = None,
        chunk_size: int = 8192,
    ):
        """Initialize progress tracking content wrapper.

        Args:
            content: The bytes content to upload.
            progress: Optional ViProgress instance for updating UI.
            task: Optional TaskID for the progress bar.
            handler: Optional GracefulExit handler for checking interruption signals.
            progress_tracker: Optional asset-level progress tracker to update bytes.
            chunk_size: Size of chunks to yield (default 8KB for responsive interruption).

        """
        self._content = content
        self._progress = progress
        self._task = task
        self._handler = handler
        self._progress_tracker = progress_tracker
        self._chunk_size = chunk_size
        self._position = 0

    def __iter__(self) -> Iterator[bytes]:
        """Iterate over content chunks while updating progress and checking for interruption.

        Raises:
            InterruptedError: If upload is cancelled via handler.exit_now.

        """
        while self._position < len(self._content):
            # Check for interruption before yielding each chunk
            # This allows the upload to be interrupted at chunk boundaries
            if self._handler and self._handler.exit_now:
                raise InterruptedError("Upload cancelled by user")

            chunk = self._content[self._position : self._position + self._chunk_size]
            self._position += len(chunk)

            # Update progress only if progress tracking is enabled
            if self._progress and self._task is not None:
                self._progress.update(self._task, advance=len(chunk))

            # Update asset-level progress tracker with bytes uploaded (real-time per chunk)
            if self._progress_tracker:
                self._progress_tracker.update_bytes(len(chunk))

            yield chunk


class MultipartProgressTracker:
    """Thread-safe progress tracker wrapper for multipart uploads with byte-level feedback."""

    def __init__(
        self,
        progress: "ViProgress",
        task: "TaskID",
        state: MultipartProgress,
        file_size: int,
    ):
        """Initialize the progress tracker.

        Args:
            progress: ViProgress instance for updating UI.
            task: TaskID for the progress bar.
            state: Initial MultipartProgress state.
            file_size: Total file size in bytes for percentage calculation.

        """
        self._progress = progress
        self._task = task
        self._state = state
        self._file_size = file_size
        self._bytes_uploaded = 0
        self._lock = Lock()

        # Update the task total to reflect the number of parts
        self._progress.update(
            self._task,
            total=state.total_parts,
            description=self._format_description(0, ""),
        )

    def update(self, success: bool = True, part_size: int = 0) -> None:
        """Update progress after part completion.

        Args:
            success: Whether the part upload succeeded.
            part_size: Size of the completed part in bytes.

        """
        with self._lock:
            self._state = MultipartProgress(
                total_parts=self._state.total_parts,
                filename=self._state.filename,
                completed=self._state.completed + 1,
            )
            if success:
                self._bytes_uploaded += part_size

            status = "uploaded" if success else "failed"
            self._progress.update(
                self._task,
                description=self._format_description(self._state.completed, status),
                advance=1,
            )

    def _format_description(self, completed: int, status: str = "") -> str:
        """Format the progress description with clear part and byte information.

        Args:
            completed: Number of completed parts.
            status: Optional status message (e.g., "uploaded", "failed").

        Returns:
            Formatted description string.

        """
        # Calculate percentage and format sizes
        percentage = (
            (self._bytes_uploaded / self._file_size * 100) if self._file_size > 0 else 0
        )

        # Format byte sizes based on file size
        if self._file_size >= 1024**3:  # >= 1 GB
            uploaded_size = self._bytes_uploaded / (1024**3)
            total_size = self._file_size / (1024**3)
            size_unit = "GB"
        elif self._file_size >= 1024**2:  # >= 1 MB
            uploaded_size = self._bytes_uploaded / (1024**2)
            total_size = self._file_size / (1024**2)
            size_unit = "MB"
        else:  # < 1 MB
            uploaded_size = self._bytes_uploaded / 1024
            total_size = self._file_size / 1024
            size_unit = "KB"

        # Build description with clear part information
        if completed < self._state.total_parts:
            next_part = completed + 1
            part_info = f"Part {next_part}/{self._state.total_parts}"
        else:
            part_info = f"{completed}/{self._state.total_parts} parts"

        status_suffix = f" {status}" if status else ""

        return (
            f"{self._state.filename} - {part_info} "
            f"({percentage:.1f}% - {uploaded_size:.2f}{size_unit}/{total_size:.2f}{size_unit}{status_suffix})"
        )


class MultipartUploader:
    """Orchestrates multipart upload for a single large file."""

    def __init__(
        self,
        uploader: "AssetUploader",
        file_path: str,
        filename: str,
        max_workers: int = MULTIPART_UPLOAD_WORKERS,
    ):
        """Initialize the multipart uploader.

        Args:
            uploader: Parent AssetUploader instance.
            file_path: Path to the file to upload.
            filename: Name of the file.
            max_workers: Maximum number of parallel part uploads.

        """
        self._uploader = uploader
        self._file_path = file_path
        self._filename = filename
        self._max_workers = max_workers
        self._file_size = os.path.getsize(file_path)

        link_parser = self._uploader._link_parser
        self._multipart_link_parser = MultipartUploadLinkParser(
            organization_id=link_parser._organization_id,
            dataset_id=link_parser._dataset_id,
            asset_source_class=link_parser._asset_source_class,
            asset_source_provider=link_parser._asset_source_provider,
            asset_source_id=link_parser._asset_source_id,
        )

    def upload(
        self,
        ingestion_id: str,
        handler: "GracefulExit",
        progress_tracker: "AssetUploadProgressTracker | None" = None,
    ) -> bool:
        """Execute multipart upload with live progress bars per part.

        Multiple parts can show live progress bars simultaneously, each on its own line.
        All parts share a single ViProgress instance, and each part adds its own task.

        Args:
            ingestion_id: The ingestion ID.
            handler: GracefulExit handler for cancellation.
            progress_tracker: Optional progress tracker to update bytes uploaded in main progress bar.

        Returns:
            True if all parts uploaded successfully, False otherwise.

        """
        # Check for cancellation before starting
        if handler.exit_now:
            logger.info(
                f"Multipart upload cancelled before starting for '{self._filename}'"
            )
            return False

        multipart_upload_id = None
        try:
            # Create multipart session
            session = self._create_session(ingestion_id)
            multipart_upload_id = session.multipart_upload_id

            parts = calculate_parts(self._file_size, session.status.part_count)

            # Upload parts in parallel with round-robin distribution
            uploaded_parts: list[responses.UploadedPart] = []
            success = False
            cancelled = False

            try:
                uploaded_parts = self._upload_parts_round_robin(
                    multipart_upload_id,
                    parts,
                    handler,
                    progress_tracker,
                )
                # Check if upload was cancelled
                if handler.exit_now and len(uploaded_parts) < len(parts):
                    cancelled = True
                    logger.info(
                        f"Multipart upload cancelled for '{self._filename}' "
                        f"({len(uploaded_parts)}/{len(parts)} parts completed)"
                    )
                else:
                    success = len(uploaded_parts) == len(parts)
            except Exception as e:
                logger.error(
                    f"Multipart upload failed for '{self._filename}': {e}",
                    extra={
                        "file": self._filename,
                        "file_size": self._file_size,
                        "multipart_upload_id": multipart_upload_id,
                        "total_parts": len(parts),
                        "uploaded_parts": len(uploaded_parts),
                        "error": str(e),
                        "error_type": type(e).__name__,
                    },
                )
                success = False

            # Check for cancellation one more time before completing
            # This handles the case where cancellation occurred after parts uploaded
            if handler.exit_now:
                logger.info(
                    f"Multipart upload cancelled before completion for '{self._filename}' "
                    f"({len(uploaded_parts)}/{len(parts)} parts completed)"
                )
                self._abort_session(multipart_upload_id)
                return False

            # Abort session on failure or cancellation to prevent resource leaks
            if not success or cancelled:
                self._abort_session(multipart_upload_id)
                return False

            # Complete multipart upload with part ETags
            self._complete_session(multipart_upload_id, uploaded_parts)
            return True

        except (KeyboardInterrupt, Exception) as e:
            # Ensure we abort the multipart session on any exception to prevent resource leaks on the server
            if multipart_upload_id is not None:
                logger.info(
                    f"Aborting multipart session {multipart_upload_id} due to {type(e).__name__}"
                )
                self._abort_session(multipart_upload_id)
            # Re-raise the exception to propagate cancellation signal
            raise

    def _create_session(self, ingestion_id: str) -> responses.MultipartUpload:
        """Create multipart session via API.

        Args:
            ingestion_id: The ingestion ID to associate with this upload.

        Returns:
            MultipartUpload with multipart_upload_id and part_count.

        """
        # Get content type from filename extension
        file_ext = Path(self._filename).suffix.lower()
        content_type = ASSET_FILE_EXTENSION_TO_MIME_TYPE_MAP.get(file_ext)

        # Construct request with spec containing ingestionId, filename, and size
        session_spec = MultipartSession(
            spec=MultipartSessionSpec(
                ingestion_id=ingestion_id,
                filename=self._filename,
                content_type=content_type,
                size_bytes=self._file_size,
            )
        )

        response = self._uploader._requester.post(
            self._multipart_link_parser(),
            json_data=msgspec.to_builtins(session_spec, str_keys=True),
            response_type=responses.MultipartUpload,
        )
        return response

    def _upload_parts_round_robin(
        self,
        multipart_upload_id: str,
        parts: list[FilePart],
        handler: "GracefulExit",
        progress_tracker: "AssetUploadProgressTracker | None" = None,
    ) -> list[responses.UploadedPart]:
        """Upload all parts using round-robin distribution across workers.

        Distributes parts evenly across workers for predictable load balancing.
        Each worker processes its assigned parts sequentially, but workers
        run in parallel.

        Args:
            multipart_upload_id: Multipart upload ID.
            parts: List of file parts to upload.
            handler: GracefulExit handler for cancellation.
            progress_tracker: Optional progress tracker to update bytes uploaded.

        Returns:
            List of successfully uploaded parts with ETags.
            Empty list if any part fails or upload is cancelled.

        """
        total_parts = len(parts)
        uploaded_parts: list[responses.UploadedPart] = []
        uploaded_parts_lock = Lock()
        failed = False
        failed_lock = Lock()

        # Distribute parts to workers using round-robin
        worker_queues: list[Queue[FilePart]] = [
            Queue() for _ in range(self._max_workers)
        ]
        for i, part in enumerate(parts):
            worker_idx = i % self._max_workers
            worker_queues[worker_idx].put(part)

        def worker_task(worker_queue: Queue[FilePart]) -> None:
            """Process all parts assigned to this worker."""
            nonlocal failed

            while not worker_queue.empty():
                # Check for cancellation or failure
                with failed_lock:
                    if failed or handler.exit_now:
                        return

                try:
                    part = worker_queue.get_nowait()
                except Empty:
                    return

                try:
                    uploaded_part = self._upload_single_part_with_etag(
                        multipart_upload_id,
                        part,
                        total_parts,
                        handler,
                        progress_tracker,
                    )

                    # Check for cancellation after part upload completes
                    if handler.exit_now:
                        logger.info(
                            f"Multipart upload cancelled after part {part.part_number} "
                            f"completed for '{self._filename}'"
                        )
                        return

                    with uploaded_parts_lock:
                        uploaded_parts.append(uploaded_part)
                except Exception as e:
                    logger.error(
                        f"Failed to upload part {part.part_number}/{total_parts} "
                        f"of file '{self._filename}': {e}",
                        extra={
                            "file": self._filename,
                            "file_size": self._file_size,
                            "part_number": part.part_number,
                            "total_parts": total_parts,
                            "part_size": part.size,
                            "part_byte_range": f"{part.start_byte}-{part.end_byte}",
                            "multipart_upload_id": multipart_upload_id,
                            "error": str(e),
                            "error_type": type(e).__name__,
                        },
                    )
                    with failed_lock:
                        failed = True
                    return

        # Execute workers in parallel
        with ThreadPoolExecutor(max_workers=self._max_workers) as executor:
            futures: dict[Future[None], int] = {}
            for worker_idx, worker_queue in enumerate(worker_queues):
                if not worker_queue.empty():
                    future = executor.submit(worker_task, worker_queue)
                    futures[future] = worker_idx

            # Wait for completion with cancellation checks
            # When cancellation is detected, we still wait for all workers to finish
            # so they can properly abort their uploads and clean up resources
            cancellation_detected = False
            remaining_futures = set(futures.keys())

            while remaining_futures:
                # Check for cancellation
                if handler.exit_now and not cancellation_detected:
                    cancellation_detected = True
                    logger.info(
                        f"Cancellation detected, waiting for {len(remaining_futures)} "
                        f"workers to finish cleanup for '{self._filename}'"
                    )
                    # Cancel futures that haven't started yet
                    for f in remaining_futures:
                        f.cancel()

                try:
                    # Wait for next completion with short timeout to check handler.exit_now frequently
                    for future in as_completed(remaining_futures, timeout=0.1):
                        remaining_futures.discard(future)
                        try:
                            future.result()
                        except Exception as e:
                            logger.error(f"Worker failed: {e}")
                            with failed_lock:
                                failed = True
                        break  # Process one at a time to check handler.exit_now
                except TimeoutError:
                    # No futures completed, loop again to check handler.exit_now
                    continue

        # Return empty list if any failure occurred
        with failed_lock:
            if failed:
                return []

        # Sort by part number for completion request
        uploaded_parts.sort(key=lambda p: p.part_number)
        return uploaded_parts

    def _upload_single_part_with_etag(
        self,
        multipart_upload_id: str,
        part: FilePart,
        total_parts: int,
        handler: "GracefulExit",
        progress_tracker: "AssetUploadProgressTracker | None" = None,
    ) -> responses.UploadedPart:
        """Upload a single part with retry logic and ETag capture.

        Args:
            multipart_upload_id: Multipart upload ID.
            part: The file part to upload.
            total_parts: Total number of parts.
            handler: GracefulExit handler for checking interruption signals.
            progress_tracker: Optional progress tracker to update bytes uploaded.

        Returns:
            UploadedPart with part number, ETag, and optional checksum.

        Raises:
            Exception: If upload fails after retries.

        """
        url_info = self._get_part_url(multipart_upload_id, part.part_number)

        etag_result: str | None = None
        headers_result: dict[str, str] = {}
        body_result: str | None = None

        def do_upload() -> None:
            nonlocal etag_result, headers_result, body_result
            etag, headers, body = self._do_upload_part_with_verification(
                part, url_info, None, None, handler, progress_tracker
            )
            etag_result = etag
            headers_result = headers
            body_result = body

        self._uploader._retry_executor.execute(
            operation=do_upload,
            operation_name=f"Upload part {part.part_number}/{total_parts}",
            max_retries=3,
            context_info={
                "file": self._filename,
                "part": part.part_number,
                "size": part.size,
            },
        )

        # Register part completion with response headers and body
        self._register_part_completion(
            multipart_upload_id=multipart_upload_id,
            part_number=part.part_number,
            response_headers=headers_result,
            response_body=body_result,
        )

        return responses.UploadedPart(
            part_number=part.part_number,
            etag=etag_result or "",
        )

    def _get_part_url(
        self, multipart_upload_id: str, part_number: int
    ) -> PartSignedUrl:
        """Get signed URL for uploading a part.

        Args:
            multipart_upload_id: Multipart upload ID.
            part_number: Part number to get URL for.

        Returns:
            PartSignedUrl with URL and headers.

        """
        request_body = {}

        return self._uploader._requester.post(
            f"{self._multipart_link_parser(multipart_upload_id)}/parts/{part_number}",
            json_data=request_body,
            response_type=PartSignedUrl,
        )

    def _register_part_completion(
        self,
        multipart_upload_id: str,
        part_number: int,
        response_headers: dict[str, str],
        response_body: str | None = None,
    ) -> None:
        """Register that a part upload has completed.

        This notifies the API that the part was successfully uploaded
        to the signed URL.

        Args:
            multipart_upload_id: Multipart upload ID.
            part_number: Part number that was uploaded.
            response_headers: HTTP response headers from the signed URL upload.
            response_body: Optional response body from the signed URL upload.

        """
        request_body = {
            "responseHeaders": dict(response_headers),
        }
        if response_body:
            request_body["responseBody"] = response_body

        self._uploader._requester.post(
            f"{self._multipart_link_parser(multipart_upload_id)}/parts/{part_number}/completion",
            json_data=request_body,
        )

    def _do_upload_part_with_verification(
        self,
        part: FilePart,
        url_info: PartSignedUrl,
        progress: "ViProgress | None" = None,
        part_task: "TaskID | None" = None,
        handler: "GracefulExit | None" = None,
        progress_tracker: "AssetUploadProgressTracker | None" = None,
    ) -> tuple[str, dict[str, str], str | None]:
        """Upload a single part's data with integrity verification.

        Args:
            part: The file part to upload.
            url_info: Signed URL information for upload.
            progress: Optional ViProgress for tracking upload progress.
            part_task: Optional TaskID for the progress bar.
            handler: Optional GracefulExit handler for checking interruption signals.
            progress_tracker: Optional asset-level progress tracker to update bytes.

        Returns:
            Tuple of (etag, response_headers, response_body).
            - etag: ETag from the response header
            - response_headers: Headers requested by backend (case-insensitive lookup)
            - response_body: Base64-encoded response body if requested by backend

        Raises:
            HTTPError: If the upload request fails.
            InterruptedError: If upload is cancelled via handler.exit_now.

        """
        data = read_part(self._file_path, part)

        # Build headers with optional checksum
        headers = dict(url_info.headers or {})

        # Wrap content with progress tracking and interruption checking
        # Use ProgressTrackingContent if either progress tracking or interruption is needed
        if handler or (progress and part_task is not None) or progress_tracker:
            content = ProgressTrackingContent(
                data, progress, part_task, handler, progress_tracker
            )
        else:
            content = data

        # Use the HTTP method specified by the backend
        http_client = self._uploader._get_upload_http_client()
        response = http_client.request(
            method=url_info.method,
            url=url_info.url,
            content=content,
            headers=headers,
        )
        response.raise_for_status()

        # Capture ETag from response (required for S3-style multipart completion)
        etag = response.headers.get("ETag", "").strip('"')

        # Collect response headers as requested by backend
        # Use case-insensitive lookup and return empty string for missing headers
        response_headers: dict[str, str] = {}
        for header_name in url_info.response_headers:
            # headers.get() is case-insensitive in httpx
            value = response.headers.get(header_name)
            response_headers[header_name] = value if value is not None else ""

        # Collect response body if requested by backend
        response_body: str | None = None
        if url_info.response_body:
            body_bytes = response.content
            response_body = base64.b64encode(body_bytes).decode("utf-8")

        return etag, response_headers, response_body

    def _abort_session(self, multipart_upload_id: str) -> None:
        """Abort multipart upload session to clean up server-side resources.

        This prevents resource leaks when upload fails or is cancelled.
        Failures are logged but not raised to avoid masking the original error.

        Args:
            multipart_upload_id: Multipart upload ID to abort.

        """
        try:
            self._uploader._requester.delete(
                self._multipart_link_parser(multipart_upload_id),
            )
            logger.info(
                f"Aborted multipart session {multipart_upload_id} for '{self._filename}'",
            )
        except Exception as e:
            # Log but don't raise - we don't want to mask the original error
            logger.warning(
                f"Failed to abort multipart session {multipart_upload_id}: {e}"
            )

    def _complete_session(
        self, multipart_upload_id: str, uploaded_parts: list[responses.UploadedPart]
    ) -> None:
        """Mark multipart upload as complete with part ETags.

        Args:
            multipart_upload_id: Multipart upload ID to complete.
            uploaded_parts: List of uploaded parts with ETags for verification.

        """
        # Build parts payload for S3-style multipart completion
        parts_payload = [
            {
                "partNumber": part.part_number,
                "etag": part.etag,
                **({"crc32c": part.crc32c} if part.crc32c else {}),
            }
            for part in uploaded_parts
        ]

        self._uploader._requester.post(
            f"{self._multipart_link_parser(multipart_upload_id)}/completion",
            json_data={"parts": parts_payload} if parts_payload else None,
        )
