"""
Agent context and approval state management.

Provides a structured context object to carry conversation history,
metadata, shared state, and optional approval workflow artifacts.
"""

from datetime import datetime
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field

from agentbyte.messages import AssistantMessage, Message, ToolCallRequest, UserMessage


class ToolApprovalRequest(BaseModel):
    """Request for approval before executing a tool call."""

    request_id: str = Field(description="Unique approval request ID")
    tool_call_id: str = Field(description="Tool call ID awaiting approval")
    tool_name: str = Field(description="Tool name")
    parameters: Dict[str, Any] = Field(description="Tool call parameters")
    original_tool_call: ToolCallRequest = Field(description="Original tool call")

    def create_response(
        self,
        approved: bool,
        reason: Optional[str] = None,
    ) -> "ToolApprovalResponse":
        """Create a response object for this approval request."""
        return ToolApprovalResponse(
            request_id=self.request_id,
            tool_call_id=self.tool_call_id,
            approved=approved,
            reason=reason,
        )


class ToolApprovalResponse(BaseModel):
    """Decision for a pending tool approval request."""

    request_id: str = Field(description="Approval request ID")
    tool_call_id: str = Field(description="Tool call ID")
    approved: bool = Field(description="Whether execution is approved")
    reason: Optional[str] = Field(default=None, description="Optional decision reason")


class AgentContext(BaseModel):
    """Unified context object for agent execution."""

    messages: List[Message] = Field(
        default_factory=list,
        description="Conversation history",
    )
    metadata: Dict[str, Any] = Field(
        default_factory=dict,
        description="Request/session metadata",
    )
    shared_state: Dict[str, Any] = Field(
        default_factory=dict,
        description="State shared across orchestrated entities",
    )
    environment: Dict[str, Any] = Field(
        default_factory=dict,
        description="Environment/config payload",
    )
    session_id: Optional[str] = Field(
        default=None,
        description="Optional session identifier",
    )
    user_id: Optional[str] = Field(
        default=None,
        description="Owning user identifier",
    )
    title: Optional[str] = Field(
        default=None,
        description="Human-readable session label",
    )
    tags: List[str] = Field(
        default_factory=list,
        description="Organisational labels",
    )
    created_at: datetime = Field(
        default_factory=datetime.now,
        description="Context creation timestamp",
    )
    updated_at: Optional[datetime] = Field(
        default=None,
        description="Stamped by the store on every save(); not set by callers",
    )

    pending_approval_requests: List[ToolApprovalRequest] = Field(
        default_factory=list,
        description="Pending tool approval requests",
    )
    approval_responses: Dict[str, ToolApprovalResponse] = Field(
        default_factory=dict,
        description="Approval responses by tool_call_id",
    )
    pending_tool_calls: Dict[str, ToolCallRequest] = Field(
        default_factory=dict,
        description="Tool calls waiting for approval",
    )

    def add_message(self, message: Message) -> None:
        """Append a message to conversation history."""
        self.messages.append(message)

    def clear_messages(self) -> None:
        """Clear message history while keeping state and metadata."""
        self.messages.clear()

    def reset(self) -> None:
        """Reset context payload except environment and session identity."""
        self.messages.clear()
        self.shared_state.clear()
        self.metadata.clear()

    @property
    def message_count(self) -> int:
        """Number of messages currently stored."""
        return len(self.messages)

    @property
    def is_empty(self) -> bool:
        """Whether no messages are stored."""
        return len(self.messages) == 0

    @property
    def waiting_for_approval(self) -> bool:
        """Whether there are pending approval requests."""
        return len(self.pending_approval_requests) > 0

    def get_last_user_message(self) -> Optional[UserMessage]:
        """Get latest user message if present."""
        for message in reversed(self.messages):
            if isinstance(message, UserMessage):
                return message
        return None

    def get_last_assistant_message(self) -> Optional[AssistantMessage]:
        """Get latest assistant message if present."""
        for message in reversed(self.messages):
            if isinstance(message, AssistantMessage):
                return message
        return None

    def add_approval_request(
        self,
        tool_call: ToolCallRequest,
        tool_name: str,
    ) -> ToolApprovalRequest:
        """Create and register a new approval request."""
        request = ToolApprovalRequest(
            request_id=f"approval_{tool_call.call_id}",
            tool_call_id=tool_call.call_id,
            tool_name=tool_name,
            parameters=tool_call.parameters,
            original_tool_call=tool_call,
        )
        self.pending_approval_requests.append(request)
        self.pending_tool_calls[tool_call.call_id] = tool_call
        return request

    def add_approval_response(self, response: ToolApprovalResponse) -> None:
        """Store approval response and clear matching pending request."""
        self.approval_responses[response.tool_call_id] = response
        self.pending_approval_requests = [
            request
            for request in self.pending_approval_requests
            if request.tool_call_id != response.tool_call_id
        ]

    def approve_tool(
        self,
        tool_call_id: str,
        reason: Optional[str] = None,
    ) -> ToolApprovalResponse:
        """Approve a pending tool call by ID."""
        if tool_call_id not in self.pending_tool_calls:
            raise ValueError(f"No pending tool call found for id '{tool_call_id}'")

        response = ToolApprovalResponse(
            request_id=f"approval_{tool_call_id}",
            tool_call_id=tool_call_id,
            approved=True,
            reason=reason,
        )
        self.add_approval_response(response)
        return response

    def deny_tool(
        self,
        tool_call_id: str,
        reason: Optional[str] = None,
    ) -> ToolApprovalResponse:
        """Deny a pending tool call by ID."""
        if tool_call_id not in self.pending_tool_calls:
            raise ValueError(f"No pending tool call found for id '{tool_call_id}'")

        response = ToolApprovalResponse(
            request_id=f"approval_{tool_call_id}",
            tool_call_id=tool_call_id,
            approved=False,
            reason=reason,
        )
        self.add_approval_response(response)
        return response

    def get_approval_response(self, tool_call_id: str) -> Optional[ToolApprovalResponse]:
        """Get approval response for specific tool call."""
        return self.approval_responses.get(tool_call_id)

    def consume_approval_response(
        self,
        tool_call_id: str,
    ) -> Optional[ToolApprovalResponse]:
        """Pop stored approval response for a tool call and clear pending state."""
        response = self.approval_responses.pop(tool_call_id, None)
        if response is None:
            return None

        self.pending_tool_calls.pop(tool_call_id, None)
        self.pending_approval_requests = [
            request
            for request in self.pending_approval_requests
            if request.tool_call_id != tool_call_id
        ]
        return response

    def get_approved_tool_calls(self) -> List[ToolCallRequest]:
        """Return approved tool calls and clear them from pending maps."""
        approved_calls: List[ToolCallRequest] = []
        processed_call_ids: List[str] = []

        for tool_call_id, response in self.approval_responses.items():
            if response.approved and tool_call_id in self.pending_tool_calls:
                approved_calls.append(self.pending_tool_calls[tool_call_id])
                processed_call_ids.append(tool_call_id)

        for tool_call_id in processed_call_ids:
            del self.approval_responses[tool_call_id]
            del self.pending_tool_calls[tool_call_id]

        return approved_calls

    def get_rejected_tool_calls(self) -> List[tuple[str, ToolCallRequest]]:
        """Return rejected tool calls with call IDs and clear them."""
        rejected_calls: List[tuple[str, ToolCallRequest]] = []
        processed_call_ids: List[str] = []

        for tool_call_id, response in self.approval_responses.items():
            if not response.approved and tool_call_id in self.pending_tool_calls:
                rejected_calls.append((tool_call_id, self.pending_tool_calls[tool_call_id]))
                processed_call_ids.append(tool_call_id)

        for tool_call_id in processed_call_ids:
            del self.approval_responses[tool_call_id]
            del self.pending_tool_calls[tool_call_id]

        return rejected_calls

    @classmethod
    def from_messages(cls, messages: List[Message]) -> "AgentContext":
        """Create context with pre-seeded message history."""
        return cls(messages=messages)


__all__ = [
    "AgentContext",
    "ToolApprovalRequest",
    "ToolApprovalResponse",
]
