"""
Chunked File Uploader

A robust Python wrapper for uploading files in chunks to backend services
supporting multipart uploads (like S3-compatible storage).

This module provides a clean interface for uploading both local files and URLs
with progress tracking, retry mechanisms, and comprehensive error handling.
"""

import logging
import math
import mimetypes
import os
import tempfile
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Callable, Dict, List, Optional, Tuple, Union
from urllib.parse import parse_qs, urlparse

import requests

# Configure logging
logger = logging.getLogger(__name__)

# Constants
CHUNK_SIZE = 1024 * 8  # 8KB for downloading
DEFAULT_TIMEOUT = 30
DEFAULT_MAX_RETRIES = 3
DEFAULT_RETRY_DELAY = 1


class ChunkUploaderError(Exception):
    """Base exception for chunk uploader errors."""
    pass


class AuthenticationError(ChunkUploaderError):
    """Raised when authentication fails."""
    pass


class UploadInitiationError(ChunkUploaderError):
    """Raised when upload initiation fails."""
    pass


class ChunkUploadError(ChunkUploaderError):
    """Raised when chunk upload fails."""
    pass


class UploadCompletionError(ChunkUploaderError):
    """Raised when upload completion fails."""
    pass


class FileDownloadError(ChunkUploaderError):
    """Raised when file download from URL fails."""
    pass


class ValidationError(ChunkUploaderError):
    """Raised when input validation fails."""
    pass


class ChunkedUploader:
    """
    A robust chunked file uploader for uploading local files or file URLs
    to backend services supporting multipart uploads.
    
    This class handles the complete upload lifecycle including:
    - Upload initiation
    - Chunked file upload with progress tracking
    - Upload completion
    - Automatic retry mechanisms
    - Comprehensive error handling and logging
    
    Attributes:
        backend_base_url (str): The base URL of the upload backend
        auth_token (Optional[str]): Authentication token for API requests
        max_workers (int): Maximum number of concurrent upload workers
        timeout (int): Request timeout in seconds
        max_retries (int): Maximum number of retry attempts
        retry_delay (float): Delay between retries in seconds
    """

    def __init__(
        self,
        backend_base_url: str,
        auth_token: Optional[str] = None,
        max_workers: int = 3,
        timeout: int = DEFAULT_TIMEOUT,
        max_retries: int = DEFAULT_MAX_RETRIES,
        retry_delay: float = DEFAULT_RETRY_DELAY,
        enable_logging: bool = True
    ):
        """
        Initialize the chunked uploader.
        
        Args:
            backend_base_url: The base URL of the upload backend
            auth_token: Optional authentication token
            max_workers: Maximum number of concurrent upload workers
            timeout: Request timeout in seconds
            max_retries: Maximum number of retry attempts
            retry_delay: Delay between retries in seconds
            enable_logging: Whether to enable detailed logging
        """
        if not backend_base_url:
            raise ValidationError("backend_base_url cannot be empty")
        
        if max_workers < 1:
            raise ValidationError("max_workers must be at least 1")
        
        if timeout < 1:
            raise ValidationError("timeout must be at least 1 second")
        
        if max_retries < 0:
            raise ValidationError("max_retries cannot be negative")
        
        if retry_delay < 0:
            raise ValidationError("retry_delay cannot be negative")
        
        self.backend_base_url = backend_base_url.rstrip("/")
        self.auth_token = auth_token
        self.max_workers = max_workers
        self.timeout = timeout
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        
        # Configure logging
        if enable_logging:
            logging.basicConfig(
                level=logging.INFO,
                format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
        
        logger.info(f"Initialized ChunkedUploader with backend: {self.backend_base_url}")

    def _headers(self) -> Dict[str, str]:
        """Generate request headers with authentication if provided."""
        headers = {"Content-Type": "application/json"}
        if self.auth_token:
            headers["Authorization"] = f"Bearer {self.auth_token}"
        return headers

    def start_upload(self, file_name: str, file_size: int, path: str) -> Dict:
        """
        Initiate a chunked upload session.
        
        Args:
            file_name: Name of the file to upload
            file_size: Size of the file in bytes
            path: Destination path for the file
            
        Returns:
            Dict containing upload session information
            
        Raises:
            AuthenticationError: If authentication fails
            UploadInitiationError: If upload initiation fails
        """
        if not file_name:
            raise ValidationError("file_name cannot be empty")
        
        if file_size <= 0:
            raise ValidationError("file_size must be positive")
        
        if not path:
            raise ValidationError("path cannot be empty")
        
        payload = {
            "file_name": file_name,
            "file_size": file_size,
            "path": path
        }
        
        logger.info(f"Initiating upload for '{file_name}' ({file_size} bytes) to '{path}'")
        
        try:
            response = requests.post(
                f"{self.backend_base_url}/file-upload/initiate",
                json=payload,
                headers=self._headers(),
                timeout=self.timeout
            )
            response.raise_for_status()
            
            result = response.json()
            logger.info(f"Upload initiated successfully for '{file_name}'")
            return result
            
        except requests.exceptions.HTTPError as e:
            status = e.response.status_code
            if status == 401:
                raise AuthenticationError(
                    f"Authentication failed while initiating upload for '{file_name}'. "
                    "Check your auth token."
                )
            elif status == 400:
                raise UploadInitiationError(
                    f"Bad request while initiating upload for '{file_name}': {e.response.text}"
                )
            else:
                raise UploadInitiationError(
                    f"HTTP {status} error during upload initiation for '{file_name}': {e.response.text}"
                )
        except Exception as e:
            raise UploadInitiationError(f"Failed to initiate upload for '{file_name}': {e}")

    def upload_chunk(
        self,
        file_path: str,
        chunk_index: int,
        chunk_size: int,
        part_url: str,
        progress_cb: Optional[Callable[[int, int], None]] = None,
        attempt: int = 1
    ) -> Dict:
        """
        Upload a single chunk of the file.
        
        Args:
            file_path: Path to the file being uploaded
            chunk_index: Index of the chunk (0-based)
            chunk_size: Size of the chunk in bytes
            part_url: URL for uploading this chunk
            progress_cb: Optional progress callback
            attempt: Current attempt number (for retries)
            
        Returns:
            Dict containing chunk upload result
            
        Raises:
            ChunkUploadError: If chunk upload fails after all retries
        """
        try:
            logger.debug(f"Uploading chunk {chunk_index + 1} ({chunk_size} bytes)")
            
            with open(file_path, "rb") as f:
                f.seek(chunk_index * chunk_size)
                data = f.read(chunk_size)

            response = requests.put(part_url, data=data, timeout=self.timeout)
            response.raise_for_status()
            
            etag = response.headers.get("ETag", "").replace('"', "")
            if not etag:
                raise ChunkUploadError(f"Missing ETag in response for chunk {chunk_index + 1}")

            if progress_cb:
                progress_cb(chunk_index, len(data))

            logger.debug(f"Successfully uploaded chunk {chunk_index + 1}")
            return {"part_number": chunk_index + 1, "part_e_tag": etag}

        except Exception as e:
            if attempt < self.max_retries:
                delay = self.retry_delay * (2 ** (attempt - 1))
                logger.warning(
                    f"Chunk {chunk_index + 1} failed (attempt {attempt}): {e}. "
                    f"Retrying in {delay}s..."
                )
                time.sleep(delay)
                return self.upload_chunk(
                    file_path, chunk_index, chunk_size, part_url, progress_cb, attempt + 1
                )
            else:
                raise ChunkUploadError(
                    f"Chunk {chunk_index + 1} failed after {self.max_retries} attempts "
                    f"(size: {chunk_size} bytes): {e}"
                )

    def upload_chunks(
        self,
        file_path: str,
        upload_info: Dict,
        progress_cb: Optional[Callable[[int], None]] = None
    ) -> List[Dict]:
        """
        Upload all chunks of the file concurrently.
        
        Args:
            file_path: Path to the file being uploaded
            upload_info: Upload session information from start_upload
            progress_cb: Optional progress callback for overall progress
            
        Returns:
            List of chunk upload results
            
        Raises:
            ChunkUploadError: If any chunk upload fails
        """
        upload_data = upload_info.get("data", upload_info)
        parts = upload_data.get("parts", [])
        total_chunks = len(parts)

        if total_chunks == 0:
            raise ValidationError("No chunks available to upload.")

        file_size = os.path.getsize(file_path)
        chunk_size = math.ceil(file_size / total_chunks)
        completed = [0] * total_chunks

        logger.info(f"Starting upload of {total_chunks} chunks for file: {file_path}")

        def wrapped_progress_cb(chunk_idx: int, chunk_bytes: int):
            completed[chunk_idx] = chunk_bytes
            if progress_cb:
                percent = int(sum(completed) * 100 / file_size)
                progress_cb(percent)

        results = [None] * total_chunks
        
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            future_to_idx = {
                executor.submit(
                    self.upload_chunk,
                    file_path,
                    i,
                    chunk_size,
                    parts[i].get("part_url") or parts[i].get("partUrl"),
                    lambda idx, b: wrapped_progress_cb(idx, b),
                ): i
                for i in range(total_chunks)
            }
            
            for future in as_completed(future_to_idx):
                idx = future_to_idx[future]
                try:
                    results[idx] = future.result()
                    logger.debug(f"Completed chunk {idx + 1}/{total_chunks}")
                except Exception as e:
                    logger.error(f"Chunk {idx + 1} failed: {e}")
                    raise ChunkUploadError(f"Upload failed at chunk {idx + 1}: {e}")
        
        logger.info(f"Successfully uploaded all {total_chunks} chunks")
        return results

    def complete_upload(self, upload_id: str, object_key: str, etags: List[Dict]) -> Dict:
        """
        Complete the chunked upload session.
        
        Args:
            upload_id: Upload session ID
            object_key: Object key for the uploaded file
            etags: List of chunk ETags from upload_chunks
            
        Returns:
            Dict containing completion result
            
        Raises:
            UploadCompletionError: If upload completion fails
        """
        if not upload_id:
            raise ValidationError("upload_id cannot be empty")
        
        if not object_key:
            raise ValidationError("object_key cannot be empty")
        
        if not etags:
            raise ValidationError("etags cannot be empty")
        
        logger.info(f"Completing upload session: {upload_id}")
        
        try:
            response = requests.post(
                f"{self.backend_base_url}/file-upload/complete",
                json={
                    "upload_id": upload_id,
                    "object_key": object_key,
                    "parts": etags
                },
                headers=self._headers(),
                timeout=self.timeout
            )
            response.raise_for_status()
            
            result = response.json()
            logger.info(f"Upload completed successfully: {upload_id}")
            return result
            
        except Exception as e:
            raise UploadCompletionError(f"Failed to complete upload: {e}")

    def _download_url_to_temp_file(self, url: str) -> Tuple[str, str]:
        """
        Download a file from URL to a temporary file.
        
        Args:
            url: URL to download from
            
        Returns:
            Tuple of (temp_file_path, filename)
            
        Raises:
            FileDownloadError: If download fails
        """
        if not url.strip():
            raise ValidationError("Empty file URL provided for upload.")

        logger.info(f"Downloading file from URL: {url}")

        try:
            response = requests.get(url, stream=True, timeout=self.timeout)
            response.raise_for_status()

            parsed = urlparse(url)
            filename = os.path.basename(parsed.path) or "downloaded_file"
            ext = os.path.splitext(filename)[1]

            # Try to determine file extension from query parameters
            if not ext:
                query_params = parse_qs(parsed.query)
                for key in ["ext", "format", "fm", "type"]:
                    val = query_params.get(key, [""])[0].lower()
                    if val.isalnum():
                        ext = "." + val
                        break

            # Try to determine extension from content type
            if not ext:
                content_type = response.headers.get("Content-Type", "")
                ext = mimetypes.guess_extension(content_type) or ""

            if ext and not filename.endswith(ext):
                filename += ext

            fd, temp_file_path = tempfile.mkstemp(suffix=ext)

            with os.fdopen(fd, "wb") as out:
                for chunk in response.iter_content(CHUNK_SIZE):
                    out.write(chunk)

            logger.info(f"Successfully downloaded file to: {temp_file_path}")
            return temp_file_path, filename

        except requests.exceptions.RequestException as e:
            raise FileDownloadError(f"Failed to download file from URL '{url}': {e}")
        except Exception as e:
            raise FileDownloadError(f"Unexpected error downloading from URL '{url}': {e}")

    def upload(
        self,
        file_or_url: Union[str, os.PathLike],
        path: str,
        progress_cb: Optional[Callable[[int], None]] = None
    ) -> Optional[str]:
        """
        Upload a file or URL to the specified path.
        
        This is the main method that orchestrates the complete upload process:
        1. Downloads URL to temp file if needed
        2. Initiates upload session
        3. Uploads all chunks concurrently
        4. Completes the upload
        5. Cleans up temporary files
        
        Args:
            file_or_url: Local file path or URL to upload
            path: Destination path for the uploaded file
            progress_cb: Optional callback for upload progress (0-100)
            
        Returns:
            The final path of the uploaded file, or None if upload failed
            
        Raises:
            ValidationError: If input validation fails
            FileNotFoundError: If local file doesn't exist
            PermissionError: If file access is denied
            ChunkUploaderError: For other upload-related errors
        """
        if not file_or_url:
            raise ValidationError("No file path or URL provided.")

        if not path:
            raise ValidationError("No destination path provided.")

        is_url = isinstance(file_or_url, str) and file_or_url.startswith(("http://", "https://"))
        temp_file_path: Optional[str] = None

        try:
            # Handle URL or local file
            if is_url:
                logger.info(f"Processing URL upload: {file_or_url}")
                temp_file_path, filename = self._download_url_to_temp_file(file_or_url)
                file_path = temp_file_path
            else:
                logger.info(f"Processing local file upload: {file_or_url}")
                file_path = str(file_or_url)
                filename = os.path.basename(file_path)

            # Validate file
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"File not found: {file_path}")

            file_size = os.path.getsize(file_path)
            if file_size == 0:
                raise ValidationError(f"File '{filename}' is empty.")

            logger.info(f"Starting upload process for '{filename}' ({file_size} bytes)")

            # Initiate upload
            upload_info = self.start_upload(filename, file_size, path)
            upload_data = upload_info.get("data", upload_info)

            # Upload chunks
            etags = self.upload_chunks(file_path, upload_info, progress_cb)

            # Complete upload
            complete_resp = self.complete_upload(
                upload_data.get("upload_id"),
                upload_data.get("object_key"),
                etags,
            )

            final_path = complete_resp.get("data", complete_resp).get("path")
            logger.info(f"Upload completed successfully. Final path: {final_path}")
            return final_path

        except FileNotFoundError:
            raise FileNotFoundError(f"Local file '{file_path}' not found.")
        except PermissionError:
            raise PermissionError(f"Permission denied while accessing '{file_path}'.")
        except Exception as e:
            logger.error(f"Upload failed for '{file_or_url}': {e}")
            raise

        finally:
            # Clean up temporary file
            if temp_file_path and os.path.exists(temp_file_path):
                try:
                    os.remove(temp_file_path)
                    logger.debug(f"Cleaned up temporary file: {temp_file_path}")
                except Exception as e:
                    logger.warning(f"Failed to clean up temporary file {temp_file_path}: {e}")


# Convenience function for quick uploads
def upload_file(
    backend_url: str,
    file_or_url: Union[str, os.PathLike],
    path: str,
    auth_token: Optional[str] = None,
    progress_cb: Optional[Callable[[int], None]] = None,
    **kwargs
) -> Optional[str]:
    """
    Convenience function for quick file uploads.
    
    Args:
        backend_url: Backend base URL
        file_or_url: File path or URL to upload
        path: Destination path
        auth_token: Optional authentication token
        progress_cb: Optional progress callback
        **kwargs: Additional arguments for ChunkedUploader
        
    Returns:
        Final path of uploaded file
    """
    uploader = ChunkedUploader(backend_url, auth_token, **kwargs)
    return uploader.upload(file_or_url, path, progress_cb)