from __future__ import annotations

from typing import Any

from ..constants import DEFAULT_BUCKET
from ..services.quilt_service import QuiltService
from ..utils import generate_signed_url, parse_s3_uri
from .auth_helpers import AuthorizationContext, check_s3_authorization

# Helpers


def _normalize_bucket(uri_or_name: str) -> str:
    if uri_or_name.startswith("s3://"):
        return uri_or_name[5:].split("/", 1)[0]
    return uri_or_name


def _authorize_s3(
    tool_name: str,
    tool_args: dict[str, Any],
    *,
    context: dict[str, Any],
) -> tuple[AuthorizationContext | None, dict[str, Any] | None]:
    auth_ctx = check_s3_authorization(tool_name, tool_args)
    if not auth_ctx.authorized or auth_ctx.s3_client is None:
        error_payload = auth_ctx.error_response()
        error_payload.update(context)
        return None, error_payload
    return auth_ctx, None


def _attach_auth_metadata(payload: dict[str, Any], auth_ctx: AuthorizationContext | None) -> dict[str, Any]:
    if auth_ctx and auth_ctx.auth_type:
        payload.setdefault("auth_type", auth_ctx.auth_type)
    return payload


def bucket_objects_list(
    bucket: str = DEFAULT_BUCKET,
    prefix: str = "",
    max_keys: int = 100,
    continuation_token: str = "",
    include_signed_urls: bool = True,
) -> dict[str, Any]:
    """List objects in an S3 bucket with optional prefix filtering - S3 bucket exploration and object retrieval tasks

    Args:
        bucket: S3 bucket name or s3:// URI (default: DEFAULT_BUCKET)
        prefix: Filter objects by prefix (default: "")
        max_keys: Maximum number of objects to return, 1-1000 (default: 100)
        continuation_token: Token for paginating through large result sets (default: "")
        include_signed_urls: Include presigned download URLs for each object (default: True)

    Returns:
        Dict with bucket info, objects list, and pagination details.

    Next step:
        Use the returned S3 metadata to answer the user's question or pass identifiers into the next bucket tool.

    Example:
        ```python
        from quilt_mcp.tools import buckets

        result = buckets.bucket_objects_list()
        # Next step: Use the returned S3 metadata to answer the user's question or pass identifiers into the next bucket tool.
        ```
    """
    bkt = _normalize_bucket(bucket)
    max_keys = max(1, min(max_keys, 1000))
    auth_ctx, error = _authorize_s3(
        "bucket_objects_list",
        {"bucket": bkt},
        context={"bucket": bkt, "prefix": prefix},
    )
    if error:
        return error
    client = auth_ctx.s3_client
    params: dict[str, Any] = {"Bucket": bkt, "MaxKeys": max_keys}
    if prefix:
        params["Prefix"] = prefix
    if continuation_token:
        params["ContinuationToken"] = continuation_token
    try:
        resp = client.list_objects_v2(**params)
    except Exception as e:
        return {"error": f"Failed to list objects: {e}", "bucket": bkt}
    objects: list[dict[str, Any]] = []
    for item in resp.get("Contents", []) or []:
        key = item.get("Key")
        s3_uri = f"s3://{bkt}/{key}"
        obj_data = {
            "key": key,
            "s3_uri": s3_uri,
            "size": item.get("Size"),
            "last_modified": str(item.get("LastModified")),
            "etag": item.get("ETag"),
            "storage_class": item.get("StorageClass"),
        }
        if include_signed_urls:
            signed_url = generate_signed_url(s3_uri)
            if signed_url:
                obj_data["download_url"] = signed_url
        objects.append(obj_data)
    return _attach_auth_metadata(
        {
            "bucket": bkt,
            "prefix": prefix,
            "objects": objects,
            "truncated": resp.get("IsTruncated", False),
            "next_token": resp.get("NextContinuationToken", ""),
            "key_count": resp.get("KeyCount", len(objects)),
            "max_keys": max_keys,
        },
        auth_ctx,
    )


def bucket_object_info(s3_uri: str) -> dict[str, Any]:
    """Get metadata information for a specific S3 object - S3 bucket exploration and object retrieval tasks

    Args:
        s3_uri: Full S3 URI, optionally with versionId query parameter
                Examples:
                - "s3://bucket-name/path/to/object"
                - "s3://bucket-name/path/to/object?versionId=abc123"

    Returns:
        Dict with object metadata including size, content type, etag, and modification date.
        For versioned objects, includes version-specific metadata.

    Version-specific Error Responses:
        - InvalidVersionId: Generic error with operation details
        - NoSuchVersion: "Version {versionId} not found for {s3_uri}"
        - AccessDenied (with versionId): "Access denied for version {versionId} of {s3_uri}"

    Next step:
        Use the returned S3 metadata to answer the user's question or pass identifiers into the next bucket tool.

    Example:
        ```python
        from quilt_mcp.tools import buckets

        result = buckets.bucket_object_info(
            s3_uri="example_value",
        )
        # Next step: Use the returned S3 metadata to answer the user's question or pass identifiers into the next bucket tool.
        ```
    """
    try:
        bucket, key, version_id = parse_s3_uri(s3_uri)
    except ValueError as e:
        return {"error": str(e)}

    auth_ctx, error = _authorize_s3(
        "bucket_object_info",
        {"bucket": bucket, "key": key},
        context={"bucket": bucket, "key": key},
    )
    if error:
        return error
    client = auth_ctx.s3_client
    try:
        # Build params dict and conditionally add VersionId
        params = {"Bucket": bucket, "Key": key}
        if version_id:
            params["VersionId"] = version_id
        head = client.head_object(**params)
    except Exception as e:
        # Handle version-specific errors
        if hasattr(e, 'response') and 'Error' in e.response:
            error_code = e.response['Error']['Code']
            if error_code == 'NoSuchVersion':
                return {"error": f"Version {version_id} not found for {s3_uri}", "bucket": bucket, "key": key}
            elif error_code == 'AccessDenied' and version_id:
                return {"error": f"Access denied for version {version_id} of {s3_uri}", "bucket": bucket, "key": key}
        return {"error": f"Failed to head object: {e}", "bucket": bucket, "key": key}
    return _attach_auth_metadata(
        {
            "bucket": bucket,
            "key": key,
            "size": head.get("ContentLength"),
            "content_type": head.get("ContentType"),
            "etag": head.get("ETag"),
            "last_modified": str(head.get("LastModified")),
            "metadata": head.get("Metadata", {}),
            "storage_class": head.get("StorageClass"),
            "cache_control": head.get("CacheControl"),
        },
        auth_ctx,
    )


def bucket_object_text(s3_uri: str, max_bytes: int = 65536, encoding: str = "utf-8") -> dict[str, Any]:
    """Read text content from an S3 object - S3 bucket exploration and object retrieval tasks

    Args:
        s3_uri: Full S3 URI, optionally with versionId query parameter
                Examples:
                - "s3://bucket-name/path/to/file.txt"
                - "s3://bucket-name/path/to/file.txt?versionId=xyz"
        max_bytes: Maximum bytes to read (default: 65536)
        encoding: Text encoding to use (default: "utf-8")

    Returns:
        Dict with decoded text content and metadata.
        For versioned objects, retrieves content from the specific version.

    Version-specific Error Responses:
        - InvalidVersionId: Generic error with operation details
        - NoSuchVersion: "Version {versionId} not found for {s3_uri}"
        - AccessDenied (with versionId): "Access denied for version {versionId} of {s3_uri}"

    Next step:
        Use the returned S3 metadata to answer the user's question or pass identifiers into the next bucket tool.

    Example:
        ```python
        from quilt_mcp.tools import buckets

        result = buckets.bucket_object_text(
            s3_uri="example_value",
        )
        # Next step: Use the returned S3 metadata to answer the user's question or pass identifiers into the next bucket tool.
        ```
    """
    try:
        bucket, key, version_id = parse_s3_uri(s3_uri)
    except ValueError as e:
        return {"error": str(e)}

    auth_ctx, error = _authorize_s3(
        "bucket_object_text",
        {"bucket": bucket, "key": key},
        context={"bucket": bucket, "key": key},
    )
    if error:
        return error
    client = auth_ctx.s3_client
    try:
        # Build params dict and conditionally add VersionId
        params = {"Bucket": bucket, "Key": key}
        if version_id:
            params["VersionId"] = version_id
        obj = client.get_object(**params)
        body = obj["Body"].read(max_bytes + 1)
    except Exception as e:
        # Handle version-specific errors
        if hasattr(e, 'response') and 'Error' in e.response:
            error_code = e.response['Error']['Code']
            if error_code == 'NoSuchVersion':
                return _attach_auth_metadata(
                    {"error": f"Version {version_id} not found for {s3_uri}", "bucket": bucket, "key": key},
                    auth_ctx,
                )
            elif error_code == 'AccessDenied' and version_id:
                return _attach_auth_metadata(
                    {"error": f"Access denied for version {version_id} of {s3_uri}", "bucket": bucket, "key": key},
                    auth_ctx,
                )
        return _attach_auth_metadata(
            {"error": f"Failed to get object: {e}", "bucket": bucket, "key": key},
            auth_ctx,
        )
    truncated = len(body) > max_bytes
    if truncated:
        body = body[:max_bytes]
    try:
        text = body.decode(encoding, errors="replace")
    except Exception as e:
        return _attach_auth_metadata(
            {"error": f"Decode failed: {e}", "bucket": bucket, "key": key},
            auth_ctx,
        )
    return _attach_auth_metadata(
        {
            "bucket": bucket,
            "key": key,
            "encoding": encoding,
            "truncated": truncated,
            "max_bytes": max_bytes,
            "text": text,
        },
        auth_ctx,
    )


def bucket_objects_put(bucket: str, items: list[dict[str, Any]]) -> dict[str, Any]:
    """Upload multiple objects to an S3 bucket - S3 bucket exploration and object retrieval tasks

    Args:
        bucket: S3 bucket name or s3:// URI
        items: List of objects to upload, each with 'key' and either 'text' or 'data' (base64)
               Optional: 'content_type', 'encoding' (for text), 'metadata' dict

    Returns:
        Dict with upload results and summary statistics.

    Next step:
        Use the returned S3 metadata to answer the user's question or pass identifiers into the next bucket tool.

    Example:
        ```python
        from quilt_mcp.tools import buckets

        result = buckets.bucket_objects_put(
            bucket="example-bucket",
            items=["example"],
        )
        # Next step: Use the returned S3 metadata to answer the user's question or pass identifiers into the next bucket tool.
        ```
    """
    import base64

    bkt = _normalize_bucket(bucket)
    if not items:
        return {"error": "items list is empty", "bucket": bkt}
    auth_ctx, error = _authorize_s3(
        "bucket_objects_put",
        {"bucket": bkt},
        context={"bucket": bkt},
    )
    if error:
        return error
    client = auth_ctx.s3_client
    results: list[dict[str, Any]] = []
    for idx, item in enumerate(items):
        key = item.get("key")
        if not key or not isinstance(key, str):
            results.append({"index": idx, "error": "missing key"})
            continue
        text = item.get("text")
        data_b64 = item.get("data")
        if (text is None) == (data_b64 is None):
            results.append({"key": key, "error": "provide exactly one of text or data"})
            continue
        if text is not None:
            encoding = item.get("encoding", "utf-8")
            try:
                body = text.encode(encoding)
            except Exception as e:
                results.append({"key": key, "error": f"encode failed: {e}"})
                continue
        else:
            try:
                body = base64.b64decode(str(data_b64), validate=True)
            except Exception as e:
                results.append({"key": key, "error": f"base64 decode failed: {e}"})
                continue
        put_kwargs: dict[str, Any] = {"Bucket": bkt, "Key": key, "Body": body}
        if item.get("content_type"):
            put_kwargs["ContentType"] = item["content_type"]
        if item.get("metadata") and isinstance(item.get("metadata"), dict):
            put_kwargs["Metadata"] = item["metadata"]
        try:
            resp = client.put_object(**put_kwargs)
            results.append(
                {
                    "key": key,
                    "etag": resp.get("ETag"),
                    "size": len(body),
                    "content_type": put_kwargs.get("ContentType"),
                }
            )
        except Exception as e:
            results.append({"key": key, "error": str(e)})
    successes = sum(1 for r in results if "etag" in r)
    return _attach_auth_metadata(
        {
            "bucket": bkt,
            "requested": len(items),
            "uploaded": successes,
            "results": results,
        },
        auth_ctx,
    )


def bucket_object_fetch(s3_uri: str, max_bytes: int = 65536, base64_encode: bool = True) -> dict[str, Any]:
    """Fetch binary or text data from an S3 object - S3 bucket exploration and object retrieval tasks

    Args:
        s3_uri: Full S3 URI, optionally with versionId query parameter
                Examples:
                - "s3://bucket-name/path/to/file"
                - "s3://bucket-name/path/to/file?versionId=xyz"
        max_bytes: Maximum bytes to read (default: 65536)
        base64_encode: Return binary data as base64 (default: True)

    Returns:
        Dict with object data as base64 or text, plus metadata.
        For versioned objects, fetches data from the specific version.

    Version-specific Error Responses:
        - InvalidVersionId: Generic error with operation details
        - NoSuchVersion: "Version {versionId} not found for {s3_uri}"
        - AccessDenied (with versionId): "Access denied for version {versionId} of {s3_uri}"

    Next step:
        Use the returned S3 metadata to answer the user's question or pass identifiers into the next bucket tool.

    Example:
        ```python
        from quilt_mcp.tools import buckets

        result = buckets.bucket_object_fetch(
            s3_uri="example_value",
        )
        # Next step: Use the returned S3 metadata to answer the user's question or pass identifiers into the next bucket tool.
        ```
    """
    import base64

    try:
        bucket, key, version_id = parse_s3_uri(s3_uri)
    except ValueError as e:
        return {"error": str(e)}

    auth_ctx, error = _authorize_s3(
        "bucket_object_fetch",
        {"bucket": bucket, "key": key},
        context={"bucket": bucket, "key": key},
    )
    if error:
        return error
    client = auth_ctx.s3_client
    try:
        # Build params dict and conditionally add VersionId
        params = {"Bucket": bucket, "Key": key}
        if version_id:
            params["VersionId"] = version_id
        obj = client.get_object(**params)
        body = obj["Body"].read(max_bytes + 1)
        content_type = obj.get("ContentType")
    except Exception as e:
        # Handle version-specific errors
        if hasattr(e, 'response') and 'Error' in e.response:
            error_code = e.response['Error']['Code']
            if error_code == 'NoSuchVersion':
                return _attach_auth_metadata(
                    {"error": f"Version {version_id} not found for {s3_uri}", "bucket": bucket, "key": key},
                    auth_ctx,
                )
            elif error_code == 'AccessDenied' and version_id:
                return _attach_auth_metadata(
                    {"error": f"Access denied for version {version_id} of {s3_uri}", "bucket": bucket, "key": key},
                    auth_ctx,
                )
        return _attach_auth_metadata(
            {"error": f"Failed to get object: {e}", "bucket": bucket, "key": key},
            auth_ctx,
        )
    truncated = len(body) > max_bytes
    if truncated:
        body = body[:max_bytes]
    if base64_encode:
        data = base64.b64encode(body).decode("ascii")
        return _attach_auth_metadata(
            {
                "bucket": bucket,
                "key": key,
                "truncated": truncated,
                "max_bytes": max_bytes,
                "base64": True,
                "data": data,
                "content_type": content_type,
                "size": len(body),
            },
            auth_ctx,
        )
    try:
        text = body.decode("utf-8")
        return _attach_auth_metadata(
            {
                "bucket": bucket,
                "key": key,
                "truncated": truncated,
                "max_bytes": max_bytes,
                "base64": False,
                "text": text,
                "content_type": content_type,
                "size": len(body),
            },
            auth_ctx,
        )
    except Exception:
        data = base64.b64encode(body).decode("ascii")
        return _attach_auth_metadata(
            {
                "bucket": bucket,
                "key": key,
                "truncated": truncated,
                "max_bytes": max_bytes,
                "base64": True,
                "data": data,
                "content_type": content_type,
                "size": len(body),
                "note": "Binary data returned as base64 after decode failure",
            },
            auth_ctx,
        )


def bucket_object_link(s3_uri: str, expiration: int = 3600) -> dict[str, Any]:
    """Generate a presigned URL for downloading an S3 object - S3 bucket exploration and object retrieval tasks

    Args:
        s3_uri: Full S3 URI, optionally with versionId query parameter
                Examples:
                - "s3://bucket-name/path/to/file"
                - "s3://bucket-name/path/to/file?versionId=xyz"
        expiration: URL expiration time in seconds (default: 3600, max: 604800)

    Returns:
        Dict with presigned URL and metadata.
        For versioned objects, generates URL for the specific version.

    Version-specific Error Responses:
        - InvalidVersionId: Generic error with operation details
        - NoSuchVersion: "Version {versionId} not found for {s3_uri}"
        - AccessDenied (with versionId): "Access denied for version {versionId} of {s3_uri}"

    Next step:
        Use the returned S3 metadata to answer the user's question or pass identifiers into the next bucket tool.

    Example:
        ```python
        from quilt_mcp.tools import buckets

        result = buckets.bucket_object_link(
            s3_uri="example_value",
        )
        # Next step: Use the returned S3 metadata to answer the user's question or pass identifiers into the next bucket tool.
        ```
    """
    try:
        bucket, key, version_id = parse_s3_uri(s3_uri)
    except ValueError as e:
        return {"error": str(e)}

    expiration = max(1, min(expiration, 604800))
    auth_ctx, error = _authorize_s3(
        "bucket_object_link",
        {"bucket": bucket, "key": key},
        context={"bucket": bucket, "key": key},
    )
    if error:
        return error
    client = auth_ctx.s3_client
    try:
        # Build params dict and conditionally add VersionId
        params = {"Bucket": bucket, "Key": key}
        if version_id:
            params["VersionId"] = version_id
        url = client.generate_presigned_url("get_object", Params=params, ExpiresIn=expiration)
        return _attach_auth_metadata(
            {
                "bucket": bucket,
                "key": key,
                "presigned_url": url,
                "expires_in": expiration,
            },
            auth_ctx,
        )
    except Exception as e:
        # Handle version-specific errors
        if hasattr(e, 'response') and 'Error' in e.response:
            error_code = e.response['Error']['Code']
            if error_code == 'NoSuchVersion':
                return _attach_auth_metadata(
                    {"error": f"Version {version_id} not found for {s3_uri}", "bucket": bucket, "key": key},
                    auth_ctx,
                )
            elif error_code == 'AccessDenied' and version_id:
                return _attach_auth_metadata(
                    {"error": f"Access denied for version {version_id} of {s3_uri}", "bucket": bucket, "key": key},
                    auth_ctx,
                )
        return _attach_auth_metadata(
            {
                "error": f"Failed to generate presigned URL: {e}",
                "bucket": bucket,
                "key": key,
            },
            auth_ctx,
        )
