"""
Memory AI State handler - mirrors Redis AI state handler functionality.

Provides AI agent state cache operations using in-memory storage.
"""

import logging
from datetime import datetime, timezone
from typing import Any

from pydantic import BaseModel

from ....domain.interfaces.cache_interfaces import IAIStateCache
from ..storage_manager import storage_manager
from .utils.key_factory import default_key_factory

logger = logging.getLogger("MemoryAIState")


class MemoryAIState(IAIStateCache):
    """
    Memory-based AI state cache handler.

    Mirrors RedisAIState functionality using in-memory storage.
    Maintains the same API for seamless cache backend switching.

    Stores data in memory for testing and development.
    """

    def __init__(self, tenant: str, user_id: str):
        """
        Initialize Memory AI state handler.

        Args:
            tenant: Tenant identifier
            user_id: User identifier
        """
        if not tenant or not user_id:
            raise ValueError(
                f"Missing required parameters: tenant={tenant}, user_id={user_id}"
            )

        self.tenant = tenant
        self.user_id = user_id
        self.keys = default_key_factory

    def _key(self, agent_name: str) -> str:
        """Build AI state key using KeyFactory (same as Redis)."""
        return self.keys.aistate(self.tenant, agent_name, self.user_id)

    # ---- Public API matching RedisAIState ----
    async def get(
        self, agent_name: str, models: type[BaseModel] | None = None
    ) -> dict[str, Any] | None:
        """
        Get AI agent state data.

        Args:
            agent_name: AI agent name
            models: Optional BaseModel class for deserialization

        Returns:
            AI agent state data or None if not found
        """
        key = self._key(agent_name)
        return await storage_manager.get(
            "ai_states", self.tenant, self.user_id, key, models
        )

    async def upsert(
        self,
        agent_name: str,
        data: dict[str, Any] | BaseModel,
        ttl: int | None = None,
    ) -> bool:
        """
        Create or update AI agent state data.

        Args:
            agent_name: AI agent name
            data: State data to store
            ttl: Time to live in seconds

        Returns:
            True if successful, False otherwise
        """
        key = self._key(agent_name)
        return await storage_manager.set(
            "ai_states", self.tenant, self.user_id, key, data, ttl
        )

    async def delete(self, agent_name: str) -> int:
        """
        Delete AI agent state data.

        Args:
            agent_name: AI agent name

        Returns:
            1 if deleted, 0 if didn't exist
        """
        key = self._key(agent_name)
        success = await storage_manager.delete(
            "ai_states", self.tenant, self.user_id, key
        )
        return 1 if success else 0

    async def exists(self, agent_name: str) -> bool:
        """
        Check if AI agent state exists.

        Args:
            agent_name: AI agent name

        Returns:
            True if exists, False otherwise
        """
        key = self._key(agent_name)
        return await storage_manager.exists("ai_states", self.tenant, self.user_id, key)

    async def get_field(self, agent_name: str, field: str) -> Any | None:
        """
        Get a specific field from AI agent state.

        Args:
            agent_name: AI agent name
            field: Field name

        Returns:
            Field value or None if not found
        """
        state_data = await self.get(agent_name)
        if state_data is None:
            return None

        if isinstance(state_data, dict):
            return state_data.get(field)
        else:
            # BaseModel instance
            return getattr(state_data, field, None)

    async def update_field(
        self,
        agent_name: str,
        field: str,
        value: Any,
        ttl: int | None = None,
    ) -> bool:
        """
        Update a specific field in AI agent state.

        Args:
            agent_name: AI agent name
            field: Field name
            value: New value
            ttl: Time to live in seconds

        Returns:
            True if successful, False otherwise
        """
        state_data = await self.get(agent_name)
        if state_data is None:
            state_data = {}

        if isinstance(state_data, BaseModel):
            state_data = state_data.model_dump()

        state_data[field] = value
        return await self.upsert(agent_name, state_data, ttl)

    async def increment_field(
        self,
        agent_name: str,
        field: str,
        increment: int = 1,
        ttl: int | None = None,
    ) -> int | None:
        """
        Atomically increment an integer field in AI agent state.

        Args:
            agent_name: AI agent name
            field: Field name
            increment: Amount to increment by
            ttl: Time to live in seconds

        Returns:
            New value after increment or None on error
        """
        state_data = await self.get(agent_name)
        if state_data is None:
            state_data = {}

        if isinstance(state_data, BaseModel):
            state_data = state_data.model_dump()

        current_value = state_data.get(field, 0)
        if not isinstance(current_value, int | float):
            logger.warning(
                f"Cannot increment non-numeric field '{field}': {current_value}"
            )
            return None

        new_value = int(current_value) + increment
        state_data[field] = new_value

        success = await self.upsert(agent_name, state_data, ttl)
        return new_value if success else None

    async def append_to_list(
        self,
        agent_name: str,
        field: str,
        value: Any,
        ttl: int | None = None,
    ) -> bool:
        """
        Append value to a list field in AI agent state.

        Args:
            agent_name: AI agent name
            field: Field name containing list
            value: Value to append
            ttl: Time to live in seconds

        Returns:
            True if successful, False otherwise
        """
        state_data = await self.get(agent_name)
        if state_data is None:
            state_data = {}

        if isinstance(state_data, BaseModel):
            state_data = state_data.model_dump()

        current_list = state_data.get(field, [])
        if not isinstance(current_list, list):
            current_list = []

        current_list.append(value)
        state_data[field] = current_list

        return await self.upsert(agent_name, state_data, ttl)

    async def get_ttl(self, agent_name: str) -> int:
        """
        Get remaining time to live for AI agent state.

        Args:
            agent_name: AI agent name identifier

        Returns:
            Remaining TTL in seconds, -1 if no expiry, -2 if doesn't exist
        """
        key = self._key(agent_name)
        return await storage_manager.get_ttl(
            "ai_states", self.tenant, self.user_id, key
        )

    async def renew_ttl(self, agent_name: str, ttl: int) -> bool:
        """
        Renew time to live for AI agent state.

        Args:
            agent_name: AI agent name identifier
            ttl: New time to live in seconds

        Returns:
            True if successful, False otherwise
        """
        key = self._key(agent_name)
        return await storage_manager.set_ttl(
            "ai_states", self.tenant, self.user_id, key, ttl
        )

    async def merge(
        self,
        agent_name: str,
        state_data: dict[str, Any],
        ttl: int | None = None,
        models: type[BaseModel] | None = None,
    ) -> dict[str, Any] | None:
        """
        Merge new data with existing AI agent state.

        Args:
            agent_name: AI agent name
            state_data: New state data to merge
            ttl: Optional TTL override
            models: Optional mapping for BaseModel deserialization

        Returns:
            Final merged state or None on failure
        """
        # Get existing state with optional BaseModel deserialization
        existing_state = await self.get(agent_name, models=models) or {}

        if isinstance(existing_state, BaseModel):
            existing_state = existing_state.model_dump()

        # Merge new data with existing
        new_state = {
            **existing_state,
            **state_data,
            "agent_type": agent_name,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }

        # Save merged state
        success = await self.upsert(agent_name, new_state, ttl)
        return new_state if success else None
