"""
S3/Object Storage Client with Agendex Governance.

Routes all object storage operations through Agendex for policy evaluation.
"""
from __future__ import annotations

import copy
import logging
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional

from .context import get_reasoning

if TYPE_CHECKING:
    from .client import AgendexClient

logger = logging.getLogger("agendex.s3")


class GovernedS3Client:
    """
    S3 client that routes all operations through Agendex governance.

    The agent never has direct S3 credentials - all operations go through
    Agendex proxy which holds credentials and enforces policy.

    Example:
        from agendex import AgendexClient
        from agendex.s3 import GovernedS3Client

        client = AgendexClient()
        s3 = GovernedS3Client(client, task="monthly_report")

        # This operation will be governed
        result = s3.get_object("reports-bucket", "monthly/2025-01.txt")

        # With reasoning for audit trail
        result = s3.get_object(
            "reports-bucket",
            "monthly/2025-01.txt",
            reasoning="Fetching Q1 report for analysis"
        )

        # Task switching (returns clone, safe for concurrent use)
        with s3.with_task("audit") as s3_audit:
            s3_audit.get_object("audit-bucket", "logs/access.log")
    """

    def __init__(
        self,
        agendex: AgendexClient,
        task: str,
        actions: Optional[Dict[str, str]] = None,
        on_event: Optional[Callable[[Dict], None]] = None,
        default_context: Optional[Dict[str, Any]] = None,
    ):
        """
        Initialize governed S3 client.

        Args:
            agendex: The AgendexClient instance
            task: Task context for all operations (used for policy scoping)
            actions: Override action names dict (e.g., {"get": "s3.fetch", "put": "s3.write"})
            on_event: Optional callback for governance events
            default_context: Default context merged into all operations (e.g., user_prompt)
        """
        self.agendex = agendex
        self._task = task
        self._actions = actions or {}
        self.on_event = on_event
        self.default_context = default_context or {}

    @property
    def task(self) -> str:
        """Current task context."""
        return self._task

    def _get_action(self, operation: str) -> str:
        """
        Get action name for S3 operation (custom → global config → default).

        Args:
            operation: One of "get", "put", "list"
        """
        if operation in self._actions:
            return self._actions[operation]
        from .config import get_global_config
        return get_global_config().get_action(f"s3.{operation}")

    def set_task(self, task: str) -> None:
        """
        Change the task context.

        Note: For concurrent/async code, prefer with_task() which returns
        a clone instead of mutating.
        """
        self._task = task

    def with_task(self, task: str) -> "_TaskContext":
        """
        Context manager for temporary task switch.

        Returns a clone of this client with the new task.
        Safe for concurrent/async usage - original is never mutated.

        Example:
            with s3.with_task("audit") as s3_audit:
                s3_audit.get_object("audit-bucket", "logs/access.log")
            # Original s3 still has original task
        """
        return _TaskContext(self, task)

    def get_object(
        self,
        bucket: str,
        key: str,
        context: Optional[Dict[str, Any]] = None,
        reasoning: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Get an object from S3 through governance.

        Args:
            bucket: S3 bucket name
            key: Object key/path
            context: Optional additional context for governance
            reasoning: Optional reasoning/explanation for audit trail
                       (auto-captured from contextvar if not provided)

        Returns:
            Object content and metadata from S3 adapter

        Raises:
            DeniedError: If access is denied by policy
            PendingApprovalError: If access requires approval
        """
        # Auto-capture reasoning from contextvar if not provided
        if reasoning is None:
            reasoning = get_reasoning()

        # Extract key prefix for policy matching
        key_prefix = ""
        if "/" in key:
            key_prefix = key.rsplit("/", 1)[0] + "/"

        # Merge contexts: default → operation-specific → reasoning
        merged_context = {**self.default_context, **(context or {})}
        if reasoning:
            merged_context["reasoning"] = reasoning

        result = self.agendex.invoke(
            action=self._get_action("get"),
            params={
                "bucket": bucket,
                "key": key,
            },
            task=self._task,
            context=merged_context if merged_context else None,
            resources=[{
                "type": "s3",
                "bucket": bucket,
                "key": key,
                "key_prefix": key_prefix,
            }],
            on_event=self.on_event,
        )

        return result

    def put_object(
        self,
        bucket: str,
        key: str,
        content: bytes,
        context: Optional[Dict[str, Any]] = None,
        reasoning: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Put an object to S3 through governance.

        Args:
            bucket: S3 bucket name
            key: Object key/path
            content: Object content
            context: Optional additional context for governance
            reasoning: Optional reasoning/explanation for audit trail
                       (auto-captured from contextvar if not provided)

        Returns:
            Result from S3 adapter
        """
        # Auto-capture reasoning from contextvar if not provided
        if reasoning is None:
            reasoning = get_reasoning()

        key_prefix = ""
        if "/" in key:
            key_prefix = key.rsplit("/", 1)[0] + "/"

        # Merge contexts: default → operation-specific → reasoning
        merged_context = {**self.default_context, **(context or {})}
        if reasoning:
            merged_context["reasoning"] = reasoning

        result = self.agendex.invoke(
            action=self._get_action("put"),
            params={
                "bucket": bucket,
                "key": key,
                "content_length": len(content),
            },
            task=self._task,
            context=merged_context if merged_context else None,
            resources=[{
                "type": "s3",
                "bucket": bucket,
                "key": key,
                "key_prefix": key_prefix,
                "operation": "write",
            }],
            on_event=self.on_event,
        )

        return result

    def list_objects(
        self,
        bucket: str,
        prefix: str = "",
        context: Optional[Dict[str, Any]] = None,
        reasoning: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        List objects in S3 through governance.

        Args:
            bucket: S3 bucket name
            prefix: Key prefix to filter by
            context: Optional additional context for governance
            reasoning: Optional reasoning/explanation for audit trail
                       (auto-captured from contextvar if not provided)

        Returns:
            List of objects from S3 adapter
        """
        # Auto-capture reasoning from contextvar if not provided
        if reasoning is None:
            reasoning = get_reasoning()

        # Merge contexts: default → operation-specific → reasoning
        merged_context = {**self.default_context, **(context or {})}
        if reasoning:
            merged_context["reasoning"] = reasoning

        result = self.agendex.invoke(
            action=self._get_action("list"),
            params={
                "bucket": bucket,
                "prefix": prefix,
            },
            task=self._task,
            context=merged_context if merged_context else None,
            resources=[{
                "type": "s3",
                "bucket": bucket,
                "key_prefix": prefix,
                "operation": "list",
            }],
            on_event=self.on_event,
        )

        return result


class _TaskContext:
    """
    Context manager for temporary task switching.

    Returns a shallow clone instead of mutating the original,
    making it safe for concurrent/async usage.
    """

    def __init__(self, client: GovernedS3Client, task: str):
        self.client = client
        self.new_task = task
        self.clone: Optional[GovernedS3Client] = None

    def __enter__(self) -> GovernedS3Client:
        self.clone = copy.copy(self.client)
        self.clone._task = self.new_task
        return self.clone

    def __exit__(self, *args: Any) -> None:
        self.clone = None

