"""Supabase-backed session manager for Strands SDK.

This module provides a persistent session storage implementation using Supabase
as the backend database. It extends the Strands SDK RepositorySessionManager,
which means it can be passed directly to Agent(session_manager=...) without
any additional wrapping.

The manager handles all CRUD operations for session data and provides proper error
handling with custom exceptions for common failure scenarios.
"""
import os
from datetime import datetime
from typing import Any

from strands.session import RepositorySessionManager
from strands.session.session_repository import SessionRepository, Session, SessionAgent, SessionMessage
from supabase import create_client, Client


class SupabaseSessionManagerError(Exception):
    """Base exception for Supabase Session Manager errors."""
    pass


class ConnectionError(SupabaseSessionManagerError):
    """Raised when connection to Supabase fails."""
    pass


class SessionNotFoundError(SupabaseSessionManagerError):
    """Raised when session is not found."""
    pass


class AgentNotFoundError(SupabaseSessionManagerError):
    """Raised when agent is not found."""
    pass


class OperationError(SupabaseSessionManagerError):
    """Raised when a database operation fails."""
    pass


class SupabaseSessionManager(RepositorySessionManager, SessionRepository):
    """Session manager using Supabase for persistent storage.

    This class extends RepositorySessionManager and implements SessionRepository,
    so it can be passed directly to Agent(session_manager=...) without wrapping.

    The manager requires a Supabase project with the appropriate schema (see the
    migrations folder for schema setup). All timestamp fields use ISO 8601 format
    strings for compatibility with the Strands SDK.

    Args:
        session_id: Unique identifier for the session
        supabase_url: Supabase project URL (optional, defaults to SUPABASE_URL env var)
        supabase_key: Supabase API key (optional, defaults to SUPABASE_KEY env var)

    Raises:
        ValueError: If Supabase credentials are not provided
        ConnectionError: If connection to Supabase fails

    Example:
        >>> from strands import Agent
        >>> from supabase_strands import SupabaseSessionManager
        >>>
        >>> session_manager = SupabaseSessionManager(
        ...     session_id="user-123",
        ...     supabase_url="https://your-project.supabase.co",
        ...     supabase_key="your-api-key",
        ... )
        >>> agent = Agent(session_manager=session_manager)
    """

    def __init__(
        self,
        session_id: str,
        supabase_url: str | None = None,
        supabase_key: str | None = None,
    ):
        url = supabase_url or os.environ.get("SUPABASE_URL")
        key = supabase_key or os.environ.get("SUPABASE_KEY")

        if not url or not key:
            raise ValueError(
                "Supabase credentials required. Provide supabase_url and supabase_key "
                "or set SUPABASE_URL and SUPABASE_KEY environment variables."
            )

        try:
            self.client: Client = create_client(url, key)
        except Exception as e:
            raise ConnectionError(
                f"Failed to create Supabase client. Check your credentials and network connection. Error: {e}"
            ) from e

        # Initialize RepositorySessionManager with self as the repository
        super().__init__(session_id=session_id, session_repository=self)

    # ── SessionRepository implementation ──────────────────────────────────

    def create_session(self, session: Session) -> Session:
        """Create a new session in the database."""
        data = {
            "session_id": session.session_id,
            "session_type": session.session_type,
            "created_at": session.created_at,
            "updated_at": session.updated_at,
        }
        try:
            self.client.table("sessions").insert(data).execute()
            return session
        except Exception as e:
            error_msg = str(e).lower()
            if "duplicate" in error_msg or "unique" in error_msg:
                raise OperationError(
                    f"Session with ID '{session.session_id}' already exists. Error: {e}"
                ) from e
            raise OperationError(
                f"Failed to create session '{session.session_id}'. Error: {e}"
            ) from e

    def read_session(self, session_id: str) -> Session | None:
        """Read a session from the database."""
        try:
            result = self.client.table("sessions").select("*").eq("session_id", session_id).execute()
            if not result.data:
                return None
            row = result.data[0]
            return Session(
                session_id=row["session_id"],
                session_type=row["session_type"],
                created_at=row["created_at"],
                updated_at=row["updated_at"],
            )
        except Exception as e:
            raise OperationError(
                f"Failed to read session '{session_id}'. Error: {e}"
            ) from e

    def delete_session(self, session_id: str) -> None:
        """Delete a session and all associated data (cascades via FK)."""
        try:
            self.client.table("sessions").delete().eq("session_id", session_id).execute()
        except Exception as e:
            raise OperationError(
                f"Failed to delete session '{session_id}'. Error: {e}"
            ) from e

    def create_agent(self, session_id: str, agent: SessionAgent) -> SessionAgent:
        """Create a new agent in a session."""
        data = {
            "session_id": session_id,
            "agent_id": agent.agent_id,
            "state": agent.state,
            "conversation_manager_state": agent.conversation_manager_state,
            "_internal_state": agent._internal_state,
            "created_at": agent.created_at,
            "updated_at": agent.updated_at,
        }
        try:
            self.client.table("agents").insert(data).execute()
            return agent
        except Exception as e:
            error_msg = str(e).lower()
            if "duplicate" in error_msg or "unique" in error_msg:
                raise OperationError(
                    f"Agent '{agent.agent_id}' already exists in session '{session_id}'. Error: {e}"
                ) from e
            if "foreign key" in error_msg:
                raise SessionNotFoundError(
                    f"Session '{session_id}' does not exist. Create the session first. Error: {e}"
                ) from e
            raise OperationError(
                f"Failed to create agent '{agent.agent_id}' in session '{session_id}'. Error: {e}"
            ) from e

    def read_agent(self, session_id: str, agent_id: str) -> SessionAgent | None:
        """Read an agent from a session."""
        try:
            result = (
                self.client.table("agents")
                .select("*")
                .eq("session_id", session_id)
                .eq("agent_id", agent_id)
                .execute()
            )
            if not result.data:
                return None
            row = result.data[0]
            return SessionAgent(
                agent_id=row["agent_id"],
                state=row["state"],
                conversation_manager_state=row["conversation_manager_state"],
                _internal_state=row["_internal_state"],
                created_at=row["created_at"],
                updated_at=row["updated_at"],
            )
        except Exception as e:
            raise OperationError(
                f"Failed to read agent '{agent_id}' from session '{session_id}'. Error: {e}"
            ) from e

    def update_agent(self, session_id: str, agent: SessionAgent) -> SessionAgent:
        """Update an existing agent's state."""
        data = {
            "state": agent.state,
            "conversation_manager_state": agent.conversation_manager_state,
            "_internal_state": agent._internal_state,
            "updated_at": datetime.utcnow().isoformat(),
        }
        try:
            self.client.table("agents").update(data).eq("session_id", session_id).eq("agent_id", agent.agent_id).execute()
            return agent
        except Exception as e:
            raise OperationError(
                f"Failed to update agent '{agent.agent_id}' in session '{session_id}'. Error: {e}"
            ) from e

    def create_message(self, session_id: str, agent_id: str, message: SessionMessage) -> SessionMessage:
        """Create a new message for an agent."""
        data = {
            "session_id": session_id,
            "agent_id": agent_id,
            "message_id": message.message_id,
            "message": message.message,
            "redact_message": message.redact_message,
            "created_at": message.created_at,
            "updated_at": message.updated_at,
        }
        try:
            self.client.table("messages").insert(data).execute()
            return message
        except Exception as e:
            error_msg = str(e).lower()
            if "duplicate" in error_msg or "unique" in error_msg:
                raise OperationError(
                    f"Message '{message.message_id}' already exists for agent '{agent_id}'. Error: {e}"
                ) from e
            if "foreign key" in error_msg:
                raise AgentNotFoundError(
                    f"Agent '{agent_id}' does not exist in session '{session_id}'. Error: {e}"
                ) from e
            raise OperationError(
                f"Failed to create message '{message.message_id}' for agent '{agent_id}'. Error: {e}"
            ) from e

    def read_message(self, session_id: str, agent_id: str, message_id: int) -> SessionMessage | None:
        """Read a specific message from an agent."""
        try:
            result = (
                self.client.table("messages")
                .select("*")
                .eq("session_id", session_id)
                .eq("agent_id", agent_id)
                .eq("message_id", message_id)
                .execute()
            )
            if not result.data:
                return None
            row = result.data[0]
            return SessionMessage(
                message_id=row["message_id"],
                message=row["message"],
                redact_message=row["redact_message"],
                created_at=row["created_at"],
                updated_at=row["updated_at"],
            )
        except Exception as e:
            raise OperationError(
                f"Failed to read message '{message_id}' for agent '{agent_id}'. Error: {e}"
            ) from e

    def update_message(self, session_id: str, agent_id: str, message: SessionMessage) -> SessionMessage:
        """Update an existing message."""
        data = {
            "message": message.message,
            "redact_message": message.redact_message,
            "updated_at": datetime.utcnow().isoformat(),
        }
        try:
            self.client.table("messages").update(data).eq("session_id", session_id).eq("agent_id", agent_id).eq("message_id", message.message_id).execute()
            return message
        except Exception as e:
            raise OperationError(
                f"Failed to update message '{message.message_id}' for agent '{agent_id}'. Error: {e}"
            ) from e

    def list_messages(
        self, session_id: str, agent_id: str, limit: int | None = None, offset: int = 0
    ) -> list[SessionMessage]:
        """List all messages for an agent with optional pagination."""
        try:
            query = (
                self.client.table("messages")
                .select("*")
                .eq("session_id", session_id)
                .eq("agent_id", agent_id)
                .order("message_id")
                .range(offset, offset + (limit - 1) if limit else 999999)
            )
            result = query.execute()
            return [
                SessionMessage(
                    message_id=row["message_id"],
                    message=row["message"],
                    redact_message=row["redact_message"],
                    created_at=row["created_at"],
                    updated_at=row["updated_at"],
                )
                for row in result.data
            ]
        except Exception as e:
            raise OperationError(
                f"Failed to list messages for agent '{agent_id}' in session '{session_id}'. Error: {e}"
            ) from e
