"""Asynchronous storage client implementation."""

from __future__ import annotations

import io
from pathlib import Path
from typing import Any, BinaryIO, Callable

from kamimind_storage._internal.auth import AuthHandler
from kamimind_storage._internal.http_client import AsyncHttpClient
from kamimind_storage.constants import (
    DEFAULT_MAX_RETRIES,
    DEFAULT_MULTIPART_THRESHOLD,
    DEFAULT_PART_SIZE,
    DEFAULT_TIMEOUT,
)
from kamimind_storage.exceptions import UploadError, ValidationError
from kamimind_storage.models import (
    DownloadInfo,
    InitiateUploadResponse,
    PublicUrlInfo,
    UploadProgress,
    UploadSession,
)
from kamimind_storage.utils.file_utils import (
    calculate_total_parts,
    get_content_type,
    get_file_size,
    sanitize_filename,
)
from kamimind_storage.utils.validation import validate_file_path, validate_session_id


class AsyncStorageClient:
    """Asynchronous client for Kamimind Storage Service."""

    def __init__(
        self,
        base_url: str,
        access_token: str,
        *,
        timeout: float = DEFAULT_TIMEOUT,
        max_retries: int = DEFAULT_MAX_RETRIES,
        multipart_threshold: int = DEFAULT_MULTIPART_THRESHOLD,
        part_size: int = DEFAULT_PART_SIZE,
    ) -> None:
        """Initialize the async storage client.

        Args:
            base_url: Base URL of the storage service.
            access_token: JWT access token for authentication.
            timeout: Request timeout in seconds.
            max_retries: Maximum retry attempts for failed requests.
            multipart_threshold: File size threshold for multipart uploads.
            part_size: Size of each part for multipart uploads.
        """
        self._base_url = base_url.rstrip("/")
        self._auth_handler = AuthHandler(access_token)
        self._http_client = AsyncHttpClient(base_url, self._auth_handler, timeout)
        self._max_retries = max_retries
        self._multipart_threshold = multipart_threshold
        self._part_size = part_size

    async def __aenter__(self) -> AsyncStorageClient:
        """Enter async context manager."""
        await self._http_client.__aenter__()
        return self

    async def __aexit__(self, *args: Any) -> None:
        """Exit async context manager."""
        await self._http_client.__aexit__(*args)

    def set_access_token(self, token: str) -> None:
        """Update the access token.

        Args:
            token: New JWT access token.
        """
        self._auth_handler.set_access_token(token)

    async def upload_file(
        self,
        file: str | Path | BinaryIO,
        *,
        filename: str | None = None,
        content_type: str | None = None,
        metadata: dict[str, Any] | None = None,
        is_public: bool = False,
        on_progress: Callable[[UploadProgress], None] | None = None,
    ) -> UploadSession:
        """Upload a file to the storage service.

        Args:
            file: File path, Path object, or file-like object.
            filename: Override filename (optional).
            content_type: MIME type (auto-detected if not provided).
            metadata: Custom metadata to attach.
            is_public: Make file publicly accessible.
            on_progress: Callback for upload progress.

        Returns:
            UploadSession with session details.
        """
        # Handle file path
        if isinstance(file, (str, Path)):
            path = validate_file_path(file)
            filename = filename or path.name
            content_type = content_type or get_content_type(filename)
            file_size = get_file_size(path)

            with open(path, "rb") as f:
                return await self._upload_stream(
                    f,
                    filename=filename,
                    file_size=file_size,
                    content_type=content_type,
                    metadata=metadata,
                    is_public=is_public,
                    on_progress=on_progress,
                )

        # Handle file-like object
        if not filename:
            raise ValidationError("Filename is required for file-like objects")

        content_type = content_type or get_content_type(filename)
        file_size = get_file_size(file)

        return await self._upload_stream(
            file,
            filename=filename,
            file_size=file_size,
            content_type=content_type,
            metadata=metadata,
            is_public=is_public,
            on_progress=on_progress,
        )

    async def upload_bytes(
        self,
        data: bytes,
        filename: str,
        *,
        content_type: str | None = None,
        metadata: dict[str, Any] | None = None,
        is_public: bool = False,
        on_progress: Callable[[UploadProgress], None] | None = None,
    ) -> UploadSession:
        """Upload bytes data as a file.

        Args:
            data: Bytes to upload.
            filename: Name for the file.
            content_type: MIME type (auto-detected if not provided).
            metadata: Custom metadata to attach.
            is_public: Make file publicly accessible.
            on_progress: Callback for upload progress.

        Returns:
            UploadSession with session details.
        """
        content_type = content_type or get_content_type(filename)

        return await self._upload_stream(
            io.BytesIO(data),
            filename=filename,
            file_size=len(data),
            content_type=content_type,
            metadata=metadata,
            is_public=is_public,
            on_progress=on_progress,
        )

    async def _upload_stream(
        self,
        stream: BinaryIO,
        filename: str,
        file_size: int,
        content_type: str,
        metadata: dict[str, Any] | None,
        is_public: bool,
        on_progress: Callable[[UploadProgress], None] | None,
    ) -> UploadSession:
        """Internal method to upload a stream."""
        filename = sanitize_filename(filename)
        use_multipart = file_size > self._multipart_threshold

        # Initiate upload
        response = await self._http_client.post(
            "initiate",
            json={
                "fileName": filename,
                "fileSize": file_size,
                "contentType": content_type,
                "metadata": metadata,
                "useMultipart": use_multipart,
                "isPublic": is_public,
            },
        )

        init_response = InitiateUploadResponse.model_validate(response["result"])

        if use_multipart:
            return await self._upload_multipart(
                stream=stream,
                session_id=init_response.session_id,
                file_size=file_size,
                on_progress=on_progress,
            )
        else:
            return await self._upload_single(
                stream=stream,
                session_id=init_response.session_id,
                upload_url=init_response.upload_url,
                file_size=file_size,
                on_progress=on_progress,
            )

    async def _upload_single(
        self,
        stream: BinaryIO,
        session_id: str,
        upload_url: str | None,
        file_size: int,
        on_progress: Callable[[UploadProgress], None] | None,
    ) -> UploadSession:
        """Upload a single file (non-multipart)."""
        if not upload_url:
            raise UploadError("No upload URL provided", session_id=session_id)

        content = stream.read()

        # Upload to presigned URL
        await self._http_client.put_presigned(upload_url, content)

        if on_progress:
            on_progress(
                UploadProgress(
                    session_id=session_id,
                    bytes_uploaded=file_size,
                    total_bytes=file_size,
                    percentage=100.0,
                )
            )

        # Complete the upload
        await self._http_client.post(f"{session_id}/complete")
        # Fetch full session details
        return await self.get_upload_status(session_id)

    async def _upload_multipart(
        self,
        stream: BinaryIO,
        session_id: str,
        file_size: int,
        on_progress: Callable[[UploadProgress], None] | None,
    ) -> UploadSession:
        """Upload using multipart."""
        total_parts = calculate_total_parts(file_size, self._part_size)
        bytes_uploaded = 0
        parts_metadata: list[dict[str, Any]] = []

        for part_number in range(1, total_parts + 1):
            chunk = stream.read(self._part_size)
            if not chunk:
                break

            # Get presigned URL for this part
            part_response = await self._http_client.post(
                f"{session_id}/part",
                json={"partNumber": part_number},
            )
            upload_url = part_response["result"]["uploadUrl"]

            # Upload the part
            response = await self._http_client.put_presigned(upload_url, chunk)
            etag = response.headers.get("ETag", "").strip('"')

            # Mark part as complete
            await self._http_client.post(
                f"{session_id}/part/complete",
                json={"partNumber": part_number, "etag": etag},
            )

            parts_metadata.append({"partNumber": part_number, "etag": etag})
            bytes_uploaded += len(chunk)

            if on_progress:
                on_progress(
                    UploadProgress(
                        session_id=session_id,
                        bytes_uploaded=bytes_uploaded,
                        total_bytes=file_size,
                        percentage=(bytes_uploaded / file_size) * 100,
                        current_part=part_number,
                        total_parts=total_parts,
                    )
                )

        # Complete the multipart upload
        await self._http_client.post(
            f"{session_id}/complete",
            json={"parts": parts_metadata},
        )
        # Fetch full session details
        return await self.get_upload_status(session_id)

    async def get_upload_status(self, session_id: str) -> UploadSession:
        """Get the status of an upload session.

        Args:
            session_id: The upload session ID.

        Returns:
            UploadSession with current status.
        """
        validate_session_id(session_id)
        response = await self._http_client.get(f"{session_id}/status")
        return UploadSession.model_validate(response["result"])

    async def abort_upload(self, session_id: str) -> None:
        """Abort an ongoing upload session.

        Args:
            session_id: The upload session ID to abort.
        """
        validate_session_id(session_id)
        await self._http_client.post(f"{session_id}/abort")

    async def get_download_url(self, session_id: str) -> DownloadInfo:
        """Get a presigned download URL for a file.

        Args:
            session_id: The upload session ID.

        Returns:
            DownloadInfo with download URL and file details.
        """
        validate_session_id(session_id)
        response = await self._http_client.get(f"{session_id}/download")
        return DownloadInfo.model_validate(response["result"])

    async def make_public(self, session_id: str) -> PublicUrlInfo:
        """Make a file publicly accessible and get the public URL.

        Args:
            session_id: The upload session ID.

        Returns:
            PublicUrlInfo with the public URL.
        """
        validate_session_id(session_id)
        response = await self._http_client.post(f"{session_id}/make-public")
        result = response["result"]

        return PublicUrlInfo.model_validate(result)

    async def toggle_public_access(self, session_id: str, is_public: bool) -> PublicUrlInfo:
        """Toggle public access for a file.

        Args:
            session_id: The upload session ID.
            is_public: Whether to make the file public.

        Returns:
            PublicUrlInfo with updated status.
        """
        validate_session_id(session_id)
        response = await self._http_client.patch(
            f"{session_id}/public-download-urls",
            json={"isPublic": is_public},
        )
        result = response["result"]

        return PublicUrlInfo.model_validate(result)

    def get_public_url(self, shortcut: str) -> str:
        """Get the public URL for a file by its shortcut code.

        Args:
            shortcut: The shortcut code.

        Returns:
            Full public URL.
        """
        return f"{self._base_url}/api/v2/storage/upload/public/{shortcut}"
