"""boto3/botocore patching for AutoDebug capture.

This module captures AWS API calls at the botocore level,
initially supporting S3 read operations (GetObject, HeadObject, ListObjectsV2).
"""

from __future__ import annotations

import base64
import hashlib
import io
import json
from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from autodebug_probe.recorder import Recorder

# Original functions
_original_make_request: Optional[Callable] = None

# Global recorder reference
_recorder: Optional["Recorder"] = None

# Supported S3 operations (read-only initially)
SUPPORTED_S3_OPS = {
    "GetObject",
    "HeadObject",
    "ListObjectsV2",
    "ListObjects",
    "ListBuckets",
    "GetObjectMetadata",
}

# Unsupported S3 operations (write operations)
UNSUPPORTED_S3_OPS = {
    "PutObject",
    "DeleteObject",
    "DeleteObjects",
    "CopyObject",
    "CreateMultipartUpload",
    "UploadPart",
    "CompleteMultipartUpload",
    "AbortMultipartUpload",
}

# Other supported AWS services (read-only)
SUPPORTED_SERVICES = {
    "s3": SUPPORTED_S3_OPS,
}

# Services that are not supported
UNSUPPORTED_SERVICES = {
    "dynamodb",
    "sqs",
    "sns",
    "lambda",
    "kinesis",
    "ec2",
    "rds",
    "iam",
}


def _get_recorder() -> "Recorder":
    """Get the global recorder instance."""
    if _recorder is None:
        from autodebug_probe.recorder import get_recorder
        return get_recorder()
    return _recorder


def _extract_bucket_key(request_dict: Dict[str, Any]) -> tuple:
    """Extract bucket and key from request dict."""
    # S3 requests have bucket in url_path or as parameter
    bucket = None
    key = None

    if "url_path" in request_dict:
        path = request_dict["url_path"]
        # Path format: /bucket/key or /bucket
        parts = path.strip("/").split("/", 1)
        if parts:
            bucket = parts[0]
            if len(parts) > 1:
                key = parts[1]

    # Also check body/context for bucket/key
    if "body" in request_dict and isinstance(request_dict["body"], dict):
        body = request_dict["body"]
        bucket = body.get("Bucket", bucket)
        key = body.get("Key", key)

    # Check context
    if "context" in request_dict and isinstance(request_dict["context"], dict):
        ctx = request_dict["context"]
        if "input_params" in ctx:
            params = ctx["input_params"]
            bucket = params.get("Bucket", bucket)
            key = params.get("Key", key)

    return bucket, key


def _capture_response_body(response: tuple) -> Dict[str, Any]:
    """Capture response body for recording."""
    result = {}

    if not response or len(response) < 2:
        return result

    http_response, response_dict = response

    # Capture status code
    if hasattr(http_response, "status_code"):
        result["status_code"] = http_response.status_code

    # Capture headers
    if hasattr(http_response, "headers"):
        # Convert to dict, filtering sensitive headers
        headers = {}
        for key, value in http_response.headers.items():
            key_lower = key.lower()
            if key_lower not in {"authorization", "x-amz-security-token"}:
                headers[key_lower] = value
        result["headers"] = headers

    # Capture response body (for GetObject)
    if response_dict and isinstance(response_dict, dict):
        if "Body" in response_dict:
            body = response_dict["Body"]
            if hasattr(body, "read"):
                # StreamingBody - read content
                try:
                    content = body.read()
                    if len(content) <= 1024 * 100:  # Max 100KB
                        try:
                            result["body"] = content.decode("utf-8")
                            result["body_encoding"] = "utf-8"
                        except UnicodeDecodeError:
                            result["body"] = base64.b64encode(content).decode("ascii")
                            result["body_encoding"] = "base64"
                    else:
                        result["body_truncated"] = True
                        result["body_size"] = len(content)

                    result["body_hash"] = hashlib.sha256(content).hexdigest()[:16]

                    # Replace the consumed StreamingBody with a seekable
                    # BytesIO wrapper so the caller can still read the data.
                    # StreamingBody has no seek() method, so the old
                    # body.seek(0) approach silently failed, leaving the
                    # caller with an empty/consumed stream (data loss).
                    response_dict["Body"] = io.BytesIO(content)

                except Exception:
                    result["body_error"] = "Failed to read body"

        # Capture other response fields
        for field in ["Contents", "KeyCount", "IsTruncated", "ContentLength", "ContentType", "ETag", "LastModified"]:
            if field in response_dict:
                value = response_dict[field]
                # Serialize datetime objects
                if hasattr(value, "isoformat"):
                    value = value.isoformat()
                result[field.lower()] = value

    return result


def _patched_make_request(self, operation_model, request_dict):
    """Patched Endpoint.make_request that captures AWS calls."""
    recorder = _get_recorder()

    service_name = operation_model.service_model.service_name
    operation_name = operation_model.name

    # Check if this service/operation is supported
    is_supported = False
    if service_name in SUPPORTED_SERVICES:
        supported_ops = SUPPORTED_SERVICES[service_name]
        is_supported = operation_name in supported_ops

    # Check if explicitly unsupported
    is_unsupported = service_name in UNSUPPORTED_SERVICES or (
        service_name == "s3" and operation_name in UNSUPPORTED_S3_OPS
    )

    if is_unsupported:
        recorder.record_event({
            "type": "AWS_UNSUPPORTED",
            "service": service_name,
            "operation": operation_name,
            "message": f"AWS {service_name}.{operation_name} not supported for replay",
        })

    # Make the actual request
    response = _original_make_request(self, operation_model, request_dict)

    # Record the call if supported
    if is_supported:
        bucket, key = _extract_bucket_key(request_dict)
        response_data = _capture_response_body(response)

        payload = {
            "service": service_name,
            "operation": operation_name,
            "bucket": bucket,
            "key": key,
            "response": response_data,
            "request_hash": hashlib.sha256(
                json.dumps({"bucket": bucket, "key": key, "op": operation_name}, sort_keys=True).encode()
            ).hexdigest()[:16],
        }

        recorder.record_aws(payload)

    return response


def install() -> None:
    """Install boto3/botocore capture patches."""
    global _original_make_request

    try:
        import botocore.endpoint
    except ImportError:
        return

    if _original_make_request is not None:
        return  # Already patched

    _original_make_request = botocore.endpoint.Endpoint.make_request
    botocore.endpoint.Endpoint.make_request = _patched_make_request


def uninstall() -> None:
    """Restore original botocore functions."""
    global _original_make_request

    try:
        import botocore.endpoint
    except ImportError:
        return

    if _original_make_request is not None:
        botocore.endpoint.Endpoint.make_request = _original_make_request
        _original_make_request = None


def get_boto3_info() -> Dict[str, Any]:
    """Get boto3/botocore detection info for manifest."""
    result = {"detected": False}

    try:
        import boto3
        result["detected"] = True
        result["boto3_version"] = boto3.__version__
    except ImportError:
        pass

    try:
        import botocore
        result["botocore_version"] = botocore.__version__
    except ImportError:
        pass

    if result["detected"]:
        result["supported_services"] = list(SUPPORTED_SERVICES.keys())
        result["supported_s3_ops"] = list(SUPPORTED_S3_OPS)
        result["unsupported_s3_ops"] = list(UNSUPPORTED_S3_OPS)

    return result
