"""AIMemory Client for the AIMemory REST API Server.

This client is designed to work with the self-hosted AIMemory server
without the need for API keys or project/org authentication.
"""

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

import httpx

logger = logging.getLogger(__name__)


class AIMemoryClient:
    """Client for interacting with the AIMemory REST API Server.

    This client matches the server API routes:
        - POST   /memories              - Create memories
        - GET    /memories              - Get all memories
        - GET    /memories/{memory_id}  - Get specific memory
        - PUT    /memories/{memory_id}  - Update a memory
        - DELETE /memories/{memory_id}  - Delete a memory
        - DELETE /memories              - Delete all memories
        - POST   /search                - Search memories
        - GET    /memories/{memory_id}/history - Get memory history
        - POST   /reset                 - Reset all memories

    Example:
        >>> from aimemory_client import AIMemoryClient
        >>> client = AIMemoryClient(host="http://localhost:8080")
        >>> client.add("I love pizza", user_id="user123")
        >>> results = client.search("What food do I like?", user_id="user123")

    Context manager example:
        >>> with AIMemoryClient(host="http://localhost:8080") as client:
        ...     client.add("I love sushi", user_id="user123")

    Attributes:
        host (str): The base URL for the AIMemory server.
        client (httpx.Client): The HTTP client used for making API requests.
    """

    def __init__(
        self,
    ):
        """Initialize the AIMemoryClient.

        Args:
            host: The base URL for the AIMemory server. Defaults to "http://localhost:8080".
            timeout: Request timeout in seconds. Defaults to 300.
            client: A custom httpx.Client instance. If provided, it will be used
                    instead of creating a new one.
        """
        self.host = "https://ai.miraiiapp.com/docs"

        self.client = httpx.Client(
                base_url=self.host,
                timeout=300,
            )
            

    def add(
        self,
        messages: Any,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Create new memories.

        Args:
            messages: A list of message dictionaries, a single message dictionary,
                     or a string. If a string is provided, it will be converted to
                     a user message.
            user_id: User identifier.
            agent_id: Agent identifier.
            run_id: Run identifier.
            metadata: Optional metadata to attach to the memory.

        Returns:
            A dictionary containing the API response with created memories.

        Raises:
            ValueError: If no identifier (user_id, agent_id, run_id) is provided.
            httpx.HTTPStatusError: If the API request fails.

        Example:
            >>> client.add("I love pizza", user_id="user123")
            >>> client.add([{"role": "user", "content": "Hello"}], user_id="user123")
        """
        # Handle different message input formats
        if isinstance(messages, str):
            messages = [{"role": "user", "content": messages}]
        elif isinstance(messages, dict):
            messages = [messages]
        elif not isinstance(messages, list):
            raise ValueError(
                f"messages must be str, dict, or list[dict], got {type(messages).__name__}"
            )

        if not any([user_id, agent_id, run_id]):
            raise ValueError("At least one identifier (user_id, agent_id, run_id) is required.")

        payload = {"messages": messages}
        if user_id:
            payload["user_id"] = user_id
        if agent_id:
            payload["agent_id"] = agent_id
        if run_id:
            payload["run_id"] = run_id
        if metadata:
            payload["metadata"] = metadata

        response = self.client.post("/memories", json=payload)
        response.raise_for_status()
        return response.json()

    def get(self, memory_id: str) -> Dict[str, Any]:
        """Retrieve a specific memory by ID.

        Args:
            memory_id: The ID of the memory to retrieve.

        Returns:
            A dictionary containing the memory data.

        Raises:
            httpx.HTTPStatusError: If the API request fails.

        Example:
            >>> memory = client.get("mem_abc123")
        """
        response = self.client.get(f"/memories/{memory_id}")
        response.raise_for_status()
        return response.json()

    def get_all(
        self,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Retrieve all memories for a given identifier.

        Args:
            user_id: User identifier.
            agent_id: Agent identifier.
            run_id: Run identifier.

        Returns:
            A dictionary containing the list of memories.

        Raises:
            ValueError: If no identifier is provided.
            httpx.HTTPStatusError: If the API request fails.

        Example:
            >>> memories = client.get_all(user_id="user123")
        """
        if not any([user_id, agent_id, run_id]):
            raise ValueError("At least one identifier (user_id, agent_id, run_id) is required.")

        params = {}
        if user_id:
            params["user_id"] = user_id
        if agent_id:
            params["agent_id"] = agent_id
        if run_id:
            params["run_id"] = run_id

        response = self.client.get("/memories", params=params)
        response.raise_for_status()
        return response.json()

    def search(
        self,
        query: str,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
        filters: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Search memories based on a query.

        Args:
            query: The search query string.
            user_id: User identifier.
            agent_id: Agent identifier.
            run_id: Run identifier.
            filters: Optional filters to apply to the search.

        Returns:
            A dictionary containing search results.

        Raises:
            httpx.HTTPStatusError: If the API request fails.

        Example:
            >>> results = client.search("What food do I like?", user_id="user123")
        """
        payload = {"query": query}
        if user_id:
            payload["user_id"] = user_id
        if agent_id:
            payload["agent_id"] = agent_id
        if run_id:
            payload["run_id"] = run_id
        if filters:
            payload["filters"] = filters

        response = self.client.post("/search", json=payload)
        response.raise_for_status()
        return response.json()

    def update(self, memory_id: str, text: str) -> Dict[str, Any]:
        """Update an existing memory.

        Args:
            memory_id: The ID of the memory to update.
            text: The new text content for the memory.

        Returns:
            A dictionary containing the API response.

        Raises:
            httpx.HTTPStatusError: If the API request fails.

        Example:
            >>> client.update("mem_abc123", "Updated memory content")
        """
        payload = {"text": text}
        response = self.client.put(f"/memories/{memory_id}", json=payload)
        response.raise_for_status()
        return response.json()

    def delete(self, memory_id: str) -> Dict[str, str]:
        """Delete a specific memory by ID.

        Args:
            memory_id: The ID of the memory to delete.

        Returns:
            A dictionary containing the success message.

        Raises:
            httpx.HTTPStatusError: If the API request fails.

        Example:
            >>> client.delete("mem_abc123")
        """
        response = self.client.delete(f"/memories/{memory_id}")
        response.raise_for_status()
        return response.json()

    def delete_all(
        self,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
    ) -> Dict[str, str]:
        """Delete all memories for a given identifier.

        Args:
            user_id: User identifier.
            agent_id: Agent identifier.
            run_id: Run identifier.

        Returns:
            A dictionary containing the success message.

        Raises:
            ValueError: If no identifier is provided.
            httpx.HTTPStatusError: If the API request fails.

        Example:
            >>> client.delete_all(user_id="user123")
        """
        if not any([user_id, agent_id, run_id]):
            raise ValueError("At least one identifier (user_id, agent_id, run_id) is required.")

        params = {}
        if user_id:
            params["user_id"] = user_id
        if agent_id:
            params["agent_id"] = agent_id
        if run_id:
            params["run_id"] = run_id

        response = self.client.delete("/memories", params=params)
        response.raise_for_status()
        return response.json()

    def history(self, memory_id: str) -> List[Dict[str, Any]]:
        """Retrieve the history of a specific memory.

        Args:
            memory_id: The ID of the memory to retrieve history for.

        Returns:
            A list of dictionaries containing the memory history.

        Raises:
            httpx.HTTPStatusError: If the API request fails.

        Example:
            >>> history = client.history("mem_abc123")
        """
        response = self.client.get(f"/memories/{memory_id}/history")
        response.raise_for_status()
        return response.json()

    def reset(self) -> Dict[str, str]:
        """Reset all memories in the system.

        Warning: This will delete ALL memories. Use with caution.

        Returns:
            A dictionary containing the success message.

        Raises:
            httpx.HTTPStatusError: If the API request fails.

        Example:
            >>> client.reset()
        """
        response = self.client.post("/reset")
        response.raise_for_status()
        return response.json()

    def close(self):
        """Close the HTTP client connection."""
        self.client.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class AsyncAIMemoryClient:
    """Asynchronous client for interacting with the AIMemory REST API Server.

    This is the async version of AIMemoryClient with the same API.

    Example:
        >>> from aimemory_client import AsyncAIMemoryClient
        >>> async with AsyncAIMemoryClient(host="http://localhost:8080") as client:
        ...     await client.add("I love pizza", user_id="user123")
        ...     results = await client.search("What food do I like?", user_id="user123")
    """

    def __init__(
        self,
        host: str = "http://localhost:8080",
        timeout: int = 300,
        client: Optional[httpx.AsyncClient] = None,
    ):
        """Initialize the AsyncAIMemoryClient.

        Args:
            host: The base URL for the AIMemory server. Defaults to "http://localhost:8080".
            timeout: Request timeout in seconds. Defaults to 300.
            client: A custom httpx.AsyncClient instance. If provided, it will be used
                    instead of creating a new one.
        """
        self.host = host.rstrip("/")

        if client is not None:
            self.client = client
            self.client.base_url = httpx.URL(self.host)
        else:
            self.client = httpx.AsyncClient(
                base_url=self.host,
                timeout=timeout,
            )

    async def add(
        self,
        messages: Any,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Create new memories.

        Args:
            messages: A list of message dictionaries, a single message dictionary,
                     or a string. If a string is provided, it will be converted to
                     a user message.
            user_id: User identifier.
            agent_id: Agent identifier.
            run_id: Run identifier.
            metadata: Optional metadata to attach to the memory.

        Returns:
            A dictionary containing the API response with created memories.

        Raises:
            ValueError: If no identifier (user_id, agent_id, run_id) is provided.
            httpx.HTTPStatusError: If the API request fails.
        """
        # Handle different message input formats
        if isinstance(messages, str):
            messages = [{"role": "user", "content": messages}]
        elif isinstance(messages, dict):
            messages = [messages]
        elif not isinstance(messages, list):
            raise ValueError(
                f"messages must be str, dict, or list[dict], got {type(messages).__name__}"
            )

        if not any([user_id, agent_id, run_id]):
            raise ValueError("At least one identifier (user_id, agent_id, run_id) is required.")

        payload = {"messages": messages}
        if user_id:
            payload["user_id"] = user_id
        if agent_id:
            payload["agent_id"] = agent_id
        if run_id:
            payload["run_id"] = run_id
        if metadata:
            payload["metadata"] = metadata

        response = await self.client.post("/memories", json=payload)
        response.raise_for_status()
        return response.json()

    async def get(self, memory_id: str) -> Dict[str, Any]:
        """Retrieve a specific memory by ID."""
        response = await self.client.get(f"/memories/{memory_id}")
        response.raise_for_status()
        return response.json()

    async def get_all(
        self,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Retrieve all memories for a given identifier."""
        if not any([user_id, agent_id, run_id]):
            raise ValueError("At least one identifier (user_id, agent_id, run_id) is required.")

        params = {}
        if user_id:
            params["user_id"] = user_id
        if agent_id:
            params["agent_id"] = agent_id
        if run_id:
            params["run_id"] = run_id

        response = await self.client.get("/memories", params=params)
        response.raise_for_status()
        return response.json()

    async def search(
        self,
        query: str,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
        filters: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Search memories based on a query."""
        payload = {"query": query}
        if user_id:
            payload["user_id"] = user_id
        if agent_id:
            payload["agent_id"] = agent_id
        if run_id:
            payload["run_id"] = run_id
        if filters:
            payload["filters"] = filters

        response = await self.client.post("/search", json=payload)
        response.raise_for_status()
        return response.json()

    async def update(self, memory_id: str, text: str) -> Dict[str, Any]:
        """Update an existing memory.

        Args:
            memory_id: The ID of the memory to update.
            text: The new text content for the memory.

        Returns:
            A dictionary containing the API response.
        """
        payload = {"text": text}
        response = await self.client.put(f"/memories/{memory_id}", json=payload)
        response.raise_for_status()
        return response.json()

    async def delete(self, memory_id: str) -> Dict[str, str]:
        """Delete a specific memory by ID."""
        response = await self.client.delete(f"/memories/{memory_id}")
        response.raise_for_status()
        return response.json()

    async def delete_all(
        self,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
    ) -> Dict[str, str]:
        """Delete all memories for a given identifier."""
        if not any([user_id, agent_id, run_id]):
            raise ValueError("At least one identifier (user_id, agent_id, run_id) is required.")

        params = {}
        if user_id:
            params["user_id"] = user_id
        if agent_id:
            params["agent_id"] = agent_id
        if run_id:
            params["run_id"] = run_id

        response = await self.client.delete("/memories", params=params)
        response.raise_for_status()
        return response.json()

    async def history(self, memory_id: str) -> List[Dict[str, Any]]:
        """Retrieve the history of a specific memory."""
        response = await self.client.get(f"/memories/{memory_id}/history")
        response.raise_for_status()
        return response.json()

    # async def reset(self) -> Dict[str, str]:
    #     """Reset all memories in the system."""
    #     response = await self.client.post("/reset")
    #     response.raise_for_status()
    #     return response.json()

    async def close(self):
        """Close the HTTP client connection."""
        await self.client.aclose()

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()
