"""Main client for zymemory API."""

import requests
from typing import List, Dict, Any, Optional
from .models import (
    Memory,
    SearchResult,
    MemoryList,
    MemoryLink,
    User,
    Representation,
    RepresentationVersions,
    MemoryAccess
)
from .exceptions import (
    APIError,
    AuthenticationError,
    ValidationError,
    NetworkError,
    RateLimitError
)


class ZymemoryClient:
    """Client for interacting with the Zymemory API.

    This client provides a Pythonic interface to all Zymemory API endpoints
    including memory management, search, linking, and conversation storage.

    Args:
        api_key: Organization API key
        org_email: Organization email address
        user_token: User-specific API token (optional, can be set later)
        api_url: Base URL for the API (default: production endpoint)
        timeout: Request timeout in seconds (default: 30)

    Example:
        >>> client = ZymemoryClient(
        ...     api_key="your-org-key",
        ...     org_email="org@example.com",
        ...     user_token="user-token"
        ... )
        >>> results = client.search("coffee preferences")
        >>> for memory in results.memories:
        ...     print(memory.content)
    """

    def __init__(
        self,
        api_key: str,
        org_email: str,
        user_token: Optional[str] = None,
        api_url: str = "https://api.heymaple.app",
        timeout: int = 30
    ):
        self.api_key = api_key
        self.org_email = org_email
        self.user_token = user_token
        self.api_url = api_url.rstrip("/")
        self.timeout = timeout

    def _get_headers(self, user_token: Optional[str] = None) -> Dict[str, str]:
        """Build request headers with authentication."""
        token = user_token or self.user_token
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "X-Org-Email": self.org_email,
            "Content-Type": "application/json"
        }
        if token:
            headers["X-User-Token"] = token
        return headers

    def _request(
        self,
        method: str,
        endpoint: str,
        json_data: Optional[Dict[str, Any]] = None,
        params: Optional[Dict[str, Any]] = None,
        user_token: Optional[str] = None
    ) -> Dict[str, Any]:
        """Make HTTP request to API with error handling."""
        url = f"{self.api_url}{endpoint}"
        headers = self._get_headers(user_token)

        try:
            response = requests.request(
                method=method,
                url=url,
                headers=headers,
                json=json_data,
                params=params,
                timeout=self.timeout
            )

            # Handle error responses
            if response.status_code == 401:
                raise AuthenticationError("Invalid API credentials")
            elif response.status_code == 429:
                raise RateLimitError("Rate limit exceeded")
            elif response.status_code >= 400:
                try:
                    error_data = response.json()
                    message = error_data.get("detail", f"API error: {response.status_code}")
                except Exception:
                    message = f"API error: {response.status_code}"
                raise APIError(message, response.status_code, error_data if 'error_data' in locals() else None)

            # Handle 204 No Content responses
            if response.status_code == 204:
                return {"success": True}

            # Try to parse JSON, return empty dict if no content
            if response.text:
                return response.json()
            else:
                return {}

        except requests.exceptions.Timeout:
            raise NetworkError(f"Request timed out after {self.timeout} seconds")
        except requests.exceptions.ConnectionError as e:
            raise NetworkError(f"Connection error: {str(e)}")
        except requests.exceptions.RequestException as e:
            raise NetworkError(f"Network error: {str(e)}")

    # =========================================================================
    # Memory Management
    # =========================================================================

    def create_memory(
        self,
        content: str,
        auto_merge: bool = False,
        user_token: Optional[str] = None
    ) -> Memory:
        """Create a new memory.

        Args:
            content: The memory content/text
            auto_merge: Whether to automatically merge with similar memories
            user_token: Override user token for this request

        Returns:
            Created Memory object

        Example:
            >>> memory = client.create_memory("I love espresso in the morning")
            >>> print(memory.id)
        """
        if not content or not content.strip():
            raise ValidationError("Memory content cannot be empty")

        data = self._request(
            "POST",
            "/memories/create",
            json_data={"content": content, "auto_merge": auto_merge},
            user_token=user_token
        )
        return Memory.from_dict(data)

    def delete_memory(
        self,
        memory_id: int,
        user_token: Optional[str] = None
    ) -> Dict[str, str]:
        """Delete a memory by ID.

        Args:
            memory_id: ID of memory to delete
            user_token: Override user token for this request

        Returns:
            Status dictionary

        Example:
            >>> client.delete_memory(123)
        """
        return self._request(
            "DELETE",
            f"/memories/del/{memory_id}",
            user_token=user_token
        )

    def list_memories(
        self,
        page: int = 1,
        page_size: int = 20,
        user_token: Optional[str] = None
    ) -> MemoryList:
        """List memories with pagination.

        Args:
            page: Page number (1-indexed)
            page_size: Number of items per page
            user_token: Override user token for this request

        Returns:
            MemoryList with pagination info

        Example:
            >>> result = client.list_memories(page=1, page_size=10)
            >>> print(f"Total: {result.total}, Page: {result.page}")
            >>> for memory in result.memories:
            ...     print(memory.content)
        """
        data = self._request(
            "GET",
            "/memories/preview",
            params={"page": page, "page_size": page_size},
            user_token=user_token
        )
        return MemoryList.from_dict(data)

    # =========================================================================
    # Search & Retrieval
    # =========================================================================

    def search(
        self,
        query: str,
        top_k: int = 5,
        user_token: Optional[str] = None
    ) -> SearchResult:
        """Search for relevant memories.

        This searches both structured memories and recent unassigned conversations.

        Args:
            query: Search query text
            top_k: Number of results to return
            user_token: Override user token for this request

        Returns:
            SearchResult with memories and unassigned conversations

        Example:
            >>> results = client.search("coffee preferences", top_k=3)
            >>> for memory in results.memories:
            ...     print(f"Match: {memory.content}")
        """
        if not query or not query.strip():
            raise ValidationError("Search query cannot be empty")

        data = self._request(
            "POST",
            "/memories/rank",
            json_data={"query": query, "page_size": top_k},
            user_token=user_token
        )
        return SearchResult.from_dict(data)

    # =========================================================================
    # Memory Links
    # =========================================================================

    def create_link(
        self,
        source_id: int,
        destination_id: int,
        user_token: Optional[str] = None
    ) -> Dict[str, Any]:
        """Create a link between two memories.

        Args:
            source_id: ID of source memory
            destination_id: ID of destination memory
            user_token: Override user token for this request

        Returns:
            Link creation response

        Example:
            >>> client.create_link(source_id=1, destination_id=2)
        """
        return self._request(
            "POST",
            "/memories/link",
            json_data={"source_id": source_id, "destination_id": destination_id},
            user_token=user_token
        )

    def get_memory_radius(
        self,
        memory_id: int,
        depth: int = 1,
        user_token: Optional[str] = None
    ) -> Dict[str, Any]:
        """Get connected memories within a radius.

        Args:
            memory_id: ID of central memory
            depth: Connection depth (default: 1)
            user_token: Override user token for this request

        Returns:
            Dictionary with 'outgoing' and 'incoming' links

        Example:
            >>> links = client.get_memory_radius(memory_id=1, depth=2)
            >>> print(f"Outgoing: {len(links['outgoing'])}")
            >>> print(f"Incoming: {len(links['incoming'])}")
        """
        return self._request(
            "GET",
            f"/memories/radius/{memory_id}",
            params={"depth": depth},
            user_token=user_token
        )

    # =========================================================================
    # Conversations
    # =========================================================================

    def store_conversation(
        self,
        user_input: str,
        assistant_output: str,
        user_token: Optional[str] = None
    ) -> Dict[str, str]:
        """Store a conversation turn.

        These conversation turns are processed into structured memories
        by the backend's clustering algorithm.

        Args:
            user_input: User's message
            assistant_output: Assistant's response
            user_token: Override user token for this request

        Returns:
            Status dictionary

        Example:
            >>> client.store_conversation(
            ...     user_input="What's the weather?",
            ...     assistant_output="It's sunny today!"
            ... )
        """
        if not user_input or not user_input.strip():
            raise ValidationError("User input cannot be empty")
        if not assistant_output or not assistant_output.strip():
            raise ValidationError("Assistant output cannot be empty")

        # /chat/register expects multipart/form-data with 'req' as JSON string
        import json
        url = f"{self.api_url}/chat/register"
        headers = self._get_headers(user_token)
        # Remove Content-Type to let requests set it for multipart/form-data
        headers.pop("Content-Type", None)

        try:
            response = requests.post(
                url,
                headers=headers,
                data={"req": json.dumps({"input": user_input, "output": assistant_output})},
                timeout=self.timeout
            )

            # Handle error responses
            if response.status_code == 401:
                raise AuthenticationError("Invalid API credentials")
            elif response.status_code == 429:
                raise RateLimitError("Rate limit exceeded")
            elif response.status_code >= 400:
                try:
                    error_data = response.json()
                    message = error_data.get("detail", f"API error: {response.status_code}")
                except Exception:
                    message = f"API error: {response.status_code}"
                raise APIError(message, response.status_code, error_data if 'error_data' in locals() else None)

            # Return successful response
            if response.text:
                return response.json()
            else:
                return {"status": "accepted"}

        except requests.exceptions.Timeout:
            raise NetworkError(f"Request timed out after {self.timeout} seconds")
        except requests.exceptions.ConnectionError as e:
            raise NetworkError(f"Connection error: {str(e)}")
        except requests.exceptions.RequestException as e:
            raise NetworkError(f"Network error: {str(e)}")

    # =========================================================================
    # User Management
    # =========================================================================

    def register_user(
        self,
        external_id: str,
        name: str,
        email: Optional[str] = None
    ) -> User:
        """Register a new end user.

        Args:
            external_id: Your system's user identifier
            name: User's name
            email: User's email (optional)

        Returns:
            User object with generated token

        Example:
            >>> user = client.register_user(
            ...     external_id="user123",
            ...     name="John Doe",
            ...     email="john@example.com"
            ... )
            >>> print(f"User token: {user.token}")
            >>> # Save user.token for future requests
        """
        if not external_id or not external_id.strip():
            raise ValidationError("External ID cannot be empty")
        if not name or not name.strip():
            raise ValidationError("Name cannot be empty")

        payload = {"external_id": external_id, "name": name}
        if email:
            payload["email"] = email

        data = self._request("POST", "/users/register", json_data=payload)
        return User.from_dict(data)

    # =========================================================================
    # Convenience Methods
    # =========================================================================

    def search_memories_only(
        self,
        query: str,
        top_k: int = 5,
        user_token: Optional[str] = None
    ) -> List[Memory]:
        """Search and return only structured memories (no unassigned conversations).

        Args:
            query: Search query
            top_k: Number of results
            user_token: Override user token

        Returns:
            List of Memory objects

        Example:
            >>> memories = client.search_memories_only("coffee")
            >>> for m in memories:
            ...     print(m.content)
        """
        result = self.search(query, top_k, user_token)
        return result.memories

    def get_all_memories(
        self,
        max_pages: int = 10,
        page_size: int = 100,
        user_token: Optional[str] = None
    ) -> List[Memory]:
        """Fetch all memories (up to max_pages).

        Args:
            max_pages: Maximum pages to fetch
            page_size: Items per page
            user_token: Override user token

        Returns:
            List of all Memory objects

        Example:
            >>> all_memories = client.get_all_memories(max_pages=5)
            >>> print(f"Total memories: {len(all_memories)}")
        """
        all_memories = []
        for page in range(1, max_pages + 1):
            result = self.list_memories(page, page_size, user_token)
            all_memories.extend(result.memories)
            if page >= (result.total + page_size - 1) // page_size:
                break
        return all_memories

    # =========================================================================
    # Memory Details & History
    # =========================================================================

    def get_memory_detail(
        self,
        memory_id: int,
        user_token: Optional[str] = None
    ) -> Memory:
        """Get detailed information about a specific memory.

        Args:
            memory_id: ID of memory to fetch
            user_token: Override user token for this request

        Returns:
            Memory object with full details

        Example:
            >>> memory = client.get_memory_detail(123)
            >>> print(memory.content)
        """
        data = self._request(
            "GET",
            f"/memories/detail/{memory_id}",
            user_token=user_token
        )
        return Memory.from_dict(data)

    def get_current_representation(
        self,
        memory_id: int,
        user_token: Optional[str] = None
    ) -> Representation:
        """Get the current representation (user instruction) for a memory.

        Args:
            memory_id: ID of memory
            user_token: Override user token for this request

        Returns:
            Representation object

        Example:
            >>> repr = client.get_current_representation(123)
            >>> print(repr.repr)
        """
        data = self._request(
            "GET",
            f"/memories/history/{memory_id}/current",
            user_token=user_token
        )
        return Representation.from_dict(data)

    def edit_memory(
        self,
        memory_id: int,
        repr_data: Dict[str, Any],
        full_rewrite: bool = False,
        user_token: Optional[str] = None
    ) -> Representation:
        """Edit a memory's representation.

        Args:
            memory_id: ID of memory to edit
            repr_data: Dictionary containing the new representation data
            full_rewrite: If True, replace all fields; if False, merge new info
            user_token: Override user token for this request

        Returns:
            New Representation object

        Example:
            >>> repr = client.edit_memory(
            ...     memory_id=123,
            ...     repr_data={"instruction": "Remember user loves coffee"},
            ...     full_rewrite=False
            ... )
        """
        data = self._request(
            "PUT",
            f"/memories/history/{memory_id}/edit",
            json_data={"repr_data": repr_data},
            params={"full_rewrite": full_rewrite},
            user_token=user_token
        )
        return Representation.from_dict(data)

    def get_all_versions(
        self,
        memory_id: int,
        user_token: Optional[str] = None
    ) -> RepresentationVersions:
        """Get all representation versions for a memory.

        Args:
            memory_id: ID of memory
            user_token: Override user token for this request

        Returns:
            RepresentationVersions object with version history

        Example:
            >>> versions = client.get_all_versions(123)
            >>> print(f"HEAD version: {versions.head}")
            >>> for v in versions.versions:
            ...     print(f"Version {v.version_num} created at {v.created_at}")
        """
        data = self._request(
            "GET",
            f"/memories/history/{memory_id}/all",
            user_token=user_token
        )
        return RepresentationVersions.from_dict(data)

    def get_version(
        self,
        memory_id: int,
        version_num: int,
        user_token: Optional[str] = None
    ) -> Representation:
        """Get a specific version of a memory's representation.

        Args:
            memory_id: ID of memory
            version_num: Version number to fetch
            user_token: Override user token for this request

        Returns:
            Representation object for that version

        Example:
            >>> repr = client.get_version(memory_id=123, version_num=2)
            >>> print(repr.repr)
        """
        data = self._request(
            "GET",
            f"/memories/history/{memory_id}/v/{version_num}",
            user_token=user_token
        )
        return Representation.from_dict(data)

    def revert_to_version(
        self,
        memory_id: int,
        version_num: int,
        user_token: Optional[str] = None
    ) -> Representation:
        """Revert a memory to a previous version.

        Args:
            memory_id: ID of memory
            version_num: Version number to revert to
            user_token: Override user token for this request

        Returns:
            Representation object after revert

        Example:
            >>> repr = client.revert_to_version(memory_id=123, version_num=1)
            >>> print(f"Reverted to version {repr.version_num}")
        """
        data = self._request(
            "POST",
            f"/memories/history/{memory_id}/revert/{version_num}",
            user_token=user_token
        )
        return Representation.from_dict(data)

    def delete_memory_history(
        self,
        memory_id: int,
        user_token: Optional[str] = None
    ) -> Dict[str, str]:
        """Delete all representation history for a memory.

        Args:
            memory_id: ID of memory
            user_token: Override user token for this request

        Returns:
            Status dictionary

        Example:
            >>> client.delete_memory_history(123)
        """
        return self._request(
            "DELETE",
            f"/memories/history/{memory_id}/del",
            user_token=user_token
        )

    # =========================================================================
    # Memory Access Control
    # =========================================================================

    def get_memory_access(
        self,
        memory_id: int,
        user_token: Optional[str] = None
    ) -> MemoryAccess:
        """Get the access control list for a memory.

        Args:
            memory_id: ID of memory
            user_token: Override user token for this request

        Returns:
            MemoryAccess object with tenant list

        Example:
            >>> access = client.get_memory_access(123)
            >>> for tenant in access.tenants:
            ...     print(f"User {tenant.user_id} has role {tenant.role}")
        """
        data = self._request(
            "GET",
            f"/memories/access/{memory_id}",
            user_token=user_token
        )
        return MemoryAccess.from_dict(data)

    def invite_user_to_memory(
        self,
        memory_id: int,
        user_id: int,
        role: int = 0,
        user_token: Optional[str] = None
    ) -> MemoryAccess:
        """Grant a user access to a memory.

        Args:
            memory_id: ID of memory
            user_id: ID of user to grant access
            role: Access role (0=read, 1=write, 2=owner)
            user_token: Override user token for this request

        Returns:
            Updated MemoryAccess object

        Example:
            >>> access = client.invite_user_to_memory(
            ...     memory_id=123,
            ...     user_id=456,
            ...     role=1  # write access
            ... )
        """
        data = self._request(
            "POST",
            f"/memories/access/{memory_id}/invite",
            json_data={"user_id": user_id, "role": role},
            user_token=user_token
        )
        return MemoryAccess.from_dict(data)

    def revoke_memory_access(
        self,
        memory_id: int,
        target_user_id: int,
        user_token: Optional[str] = None
    ) -> Dict[str, str]:
        """Revoke a user's access to a memory.

        Args:
            memory_id: ID of memory
            target_user_id: ID of user to revoke access from
            user_token: Override user token for this request

        Returns:
            Status dictionary

        Example:
            >>> client.revoke_memory_access(memory_id=123, target_user_id=456)
        """
        return self._request(
            "DELETE",
            f"/memories/access/{memory_id}/del/{target_user_id}",
            user_token=user_token
        )
