"""
MIT License Copyright (c) 2025-present June

Permission is hereby granted, free of
charge, to any person obtaining a copy of this software and associated
documentation files (the "Software"), to deal in the Software without
restriction, including without limitation the rights to use, copy, modify, merge,
publish, distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to the
following conditions:

The above copyright notice and this permission notice
(including the next paragraph) shall be included in all copies or substantial
portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO
EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

import json
import logging
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    from messages.dialogs import DialogInstance

logger = logging.getLogger(__name__)


class SessionStore(ABC):
    """Abstract base class for session storage.

    Defines the interface for storing and retrieving dialog session state
    and OTP (one-time password) data. Implementations can use different backends
    (in-memory, Redis, database, etc.).
    """

    @abstractmethod
    async def get(self, key: str) -> Optional[List[dict]]:
        """Retrieve session stack for a given key (user ID).

        Args:
            key: The user ID or session key.

        Returns:
            List of dialog instance dicts, or None if not found.
        """
        pass

    @abstractmethod
    async def set(self, key: str, stack: List["DialogInstance"]) -> None:
        """Store session stack for a given key.

        Args:
            key: The user ID or session key.
            stack: List of DialogInstance objects to store.
        """
        pass

    @abstractmethod
    async def delete(self, key: str) -> None:
        """Delete session for a given key.

        Args:
            key: The user ID or session key to delete.
        """
        pass

    @abstractmethod
    async def set_otp(self, user_id: str, otp_data: Dict[str, Any], ttl_seconds: Optional[int] = None) -> None:
        """Store OTP data for a user with optional TTL.

        Args:
            user_id: The user ID.
            otp_data: Dictionary containing OTP information.
            ttl_seconds: Time to live in seconds (optional, backend-dependent).
        """
        pass

    @abstractmethod
    async def get_otp(self, user_id: str) -> Optional[Dict[str, Any]]:
        """Retrieve OTP data for a user.

        Args:
            user_id: The user ID.

        Returns:
            OTP data dictionary, or None if not found or expired.
        """
        pass

    @abstractmethod
    async def delete_otp(self, user_id: str) -> None:
        """Delete OTP data for a user.

        Args:
            user_id: The user ID.
        """
        pass

    @abstractmethod
    async def set_awaiting_flag(self, user_id: str, flag: bool, ttl_seconds: int = 600) -> None:
        """Set or clear a flag indicating the user is awaiting an async response.

        Args:
            user_id: The user ID.
            flag: True to set awaiting, False to clear.
            ttl_seconds: Time to live in seconds (backend-dependent).
        """
        pass

    @abstractmethod
    async def get_awaiting_flag(self, user_id: str) -> bool:
        """Check whether the user is currently awaiting an async response.

        Args:
            user_id: The user ID.

        Returns:
            True if awaiting flag is set, False otherwise.
        """
        pass

    @abstractmethod
    async def set_correlation_id(self, user_id: str, correlation_id: str, ttl_seconds: int = 600) -> None:
        """Store a correlation ID for async response tracking.

        Args:
            user_id: The user ID.
            correlation_id: The correlation ID to store.
            ttl_seconds: Time to live in seconds (backend-dependent).
        """
        pass

    @abstractmethod
    async def get_correlation_id(self, user_id: str) -> Optional[str]:
        """Retrieve the stored correlation ID for a user.

        Args:
            user_id: The user ID.

        Returns:
            The correlation ID string, or None if not found.
        """
        pass


class InMemoryStore(SessionStore):
    """In-memory session storage (default, non-persistent).

    Stores session state and OTP data in memory. Data is lost when the process terminates.
    Suitable for development, testing, or single-instance deployments.
    For production, use a persistent backend like Redis or a database.
    """

    def __init__(self):
        """Initialize the in-memory store."""
        self._sessions: Dict[str, Any] = {}
        self._otp_store: Dict[str, Dict[str, Any]] = {}

    async def get(self, key: str) -> Optional[List[dict]]:
        """Retrieve stored session for a user.

        Args:
            key: The user ID.

        Returns:
            List of dialog instance dicts, or None if not stored.
        """
        return self._sessions.get(key)

    async def set(self, key: str, stack: List["DialogInstance"]) -> None:
        """Store session for a user.

        Args:
            key: The user ID.
            stack: List of DialogInstance objects.
        """
        self._sessions[key] = [{"id": item.id, "state": item.state} for item in stack]

    async def delete(self, key: str) -> None:
        """Delete session for a user.

        Args:
            key: The user ID.
        """
        self._sessions.pop(key, None)

    async def set_otp(self, user_id: str, otp_data: Dict[str, Any], ttl_seconds: Optional[int] = None) -> None:
        """Store OTP data for a user. TTL is ignored in memory store (no expiry).

        Args:
            user_id: The user ID.
            otp_data: OTP information dictionary.
            ttl_seconds: Ignored (in-memory store doesn't expire).
        """
        self._otp_store[user_id] = otp_data

    async def get_otp(self, user_id: str) -> Optional[Dict[str, Any]]:
        """Retrieve OTP data for a user.

        Args:
            user_id: The user ID.

        Returns:
            OTP data dictionary, or None if not found.
        """
        return self._otp_store.get(user_id)

    async def delete_otp(self, user_id: str) -> None:
        """Delete OTP data for a user.

        Args:
            user_id: The user ID.
        """
        self._otp_store.pop(user_id, None)

    async def set_awaiting_flag(self, user_id: str, flag: bool, ttl_seconds: int = 600) -> None:
        """Set or clear the awaiting flag (TTL ignored in-memory).

        Args:
            user_id: The user ID.
            flag: True to set, False to clear.
            ttl_seconds: Ignored.
        """
        key = f"awaiting:{user_id}"
        if flag:
            self._sessions[key] = True
        else:
            self._sessions.pop(key, None)

    async def get_awaiting_flag(self, user_id: str) -> bool:
        """Check the awaiting flag.

        Args:
            user_id: The user ID.

        Returns:
            True if flag is set, False otherwise.
        """
        return bool(self._sessions.get(f"awaiting:{user_id}"))

    async def set_correlation_id(self, user_id: str, correlation_id: str, ttl_seconds: int = 600) -> None:
        """Store a correlation ID (TTL ignored in-memory).

        Args:
            user_id: The user ID.
            correlation_id: The correlation ID.
            ttl_seconds: Ignored.
        """
        self._sessions[f"correlation:{user_id}"] = correlation_id

    async def get_correlation_id(self, user_id: str) -> Optional[str]:
        """Retrieve the stored correlation ID.

        Args:
            user_id: The user ID.

        Returns:
            The correlation ID, or None if not found.
        """
        return self._sessions.get(f"correlation:{user_id}")


try:
    import redis.asyncio as aioredis

    class RedisStore(SessionStore):
        """Redis-backed session storage for production/multi-instance deployments.

        Uses redis.asyncio for async operations. Install with:
            pip install wasup.py[redis]

        Args:
            host: Redis host.
            port: Redis port.
            password: Optional Redis password.
            prefix: Key prefix for session data.
            otp_prefix: Key prefix for OTP data.
            ssl: Whether to use SSL.
        """

        def __init__(
            self,
            host: str,
            port: int,
            password: str = None,
            prefix: str = "wa_session:",
            otp_prefix: str = "wa_otp:",
            ssl: bool = False,
        ):
            self._prefix = prefix
            self._otp_prefix = otp_prefix
            self.client = aioredis.Redis(host=host, port=port, password=password, ssl=ssl, decode_responses=True)

        async def get(self, key: str) -> Optional[List[dict]]:
            raw = await self.client.get(f"{self._prefix}{key}")
            if raw is None:
                return None
            return json.loads(raw)

        async def set(self, key: str, stack: List["DialogInstance"]) -> None:
            data = json.dumps([{"id": i.id, "state": i.state} for i in stack])
            await self.client.set(f"{self._prefix}{key}", data)

        async def delete(self, key: str) -> None:
            await self.client.delete(f"{self._prefix}{key}")

        async def set_otp(self, user_id: str, otp_data: Dict[str, Any], ttl_seconds: Optional[int] = None) -> None:
            kwargs = {}
            if ttl_seconds is not None:
                kwargs["ex"] = ttl_seconds
            await self.client.set(f"{self._otp_prefix}{user_id}", json.dumps(otp_data), **kwargs)

        async def get_otp(self, user_id: str) -> Optional[Dict[str, Any]]:
            raw = await self.client.get(f"{self._otp_prefix}{user_id}")
            if raw is None:
                return None
            return json.loads(raw)

        async def delete_otp(self, user_id: str) -> None:
            await self.client.delete(f"{self._otp_prefix}{user_id}")

        async def set_awaiting_flag(self, user_id: str, flag: bool, ttl_seconds: int = 600) -> None:
            key = f"awaiting_response:{user_id}"
            if flag:
                await self.client.set(key, "1", ex=ttl_seconds)
            else:
                await self.client.delete(key)

        async def get_awaiting_flag(self, user_id: str) -> bool:
            val = await self.client.get(f"awaiting_response:{user_id}")
            return val is not None

        async def set_correlation_id(self, user_id: str, correlation_id: str, ttl_seconds: int = 600) -> None:
            await self.client.set(f"correlation:{user_id}", correlation_id, ex=ttl_seconds)

        async def get_correlation_id(self, user_id: str) -> Optional[str]:
            return await self.client.get(f"correlation:{user_id}")

except ImportError:
    pass
