"""
MemoryAgent - High-level memory-aware agent abstraction.

Provides smart memory operations for AI agents.
"""
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Callable
import httpx


@dataclass
class RecallResult:
    """Result of a recall operation."""
    content: str
    confidence: float
    needs_clarification: bool = False
    fallback_question: Optional[str] = None
    episodes: List[Dict[str, Any]] = field(default_factory=list)


@dataclass
class ContextSummary:
    """Summarized context for agent."""
    text: str
    token_count: int
    episodes_included: int
    truncated: bool = False


class MemoryAgent:
    """
    High-level memory-aware agent abstraction.

    Provides intelligent memory operations with auto-importance detection,
    confidence-based retrieval, and context management.

    Usage:
        from hieromem.agent import MemoryAgent

        agent = MemoryAgent(
            memory_url="http://localhost:8000",
            api_key="sk-...",
            confidence_threshold=0.7
        )

        # Auto-decide what to remember
        agent.remember_if_useful("User prefers dark mode", context=history)

        # Recall with fallback
        result = agent.recall_or_ask("user's timezone", "What timezone are you in?")

        # Summarize context within token budget
        context = agent.summarize_context("help with Python", max_tokens=2000)
    """

    def __init__(
        self,
        memory_url: str = "http://localhost:8000",
        api_key: Optional[str] = None,
        namespace: str = "agent",
        confidence_threshold: float = 0.7,
        importance_llm: Optional[Callable[[str, str], float]] = None,
    ) -> None:
        """
        Initialize MemoryAgent.

        Args:
            memory_url: HIEROMEM API URL
            api_key: API key for authentication
            namespace: Default namespace for operations
            confidence_threshold: Minimum confidence for recall
            importance_llm: Optional LLM function to assess importance
        """
        self.memory_url = memory_url
        self.api_key = api_key
        self.namespace = namespace
        self.confidence_threshold = confidence_threshold
        self.importance_llm = importance_llm

        self._client = httpx.Client(
            base_url=memory_url,
            headers={"Authorization": f"Bearer {api_key}"} if api_key else {},
            timeout=30.0,
        )

    def remember_if_useful(
        self,
        content: str,
        context: Optional[str] = None,
        importance_threshold: float = 0.5,
    ) -> Optional[str]:
        """
        Store content only if deemed useful/important.

        Uses heuristics or LLM to assess importance before storing.

        Args:
            content: The content to potentially remember
            context: Conversation context for importance assessment
            importance_threshold: Minimum importance score (0-1)

        Returns:
            Episode ID if stored, None if skipped
        """
        importance = self._assess_importance(content, context)

        if importance < importance_threshold:
            return None

        try:
            response = self._client.post(
                "/memory/insert",
                json={
                    "text": content,
                    "namespace": self.namespace,
                    "meta": {
                        "importance": importance,
                        "type": "agent_memory",
                    },
                },
            )
            response.raise_for_status()
            result: Optional[str] = response.json().get("episode_id")
            return result
        except httpx.HTTPError:
            return None

    def recall_or_ask(
        self,
        query: str,
        fallback_question: str,
        k: int = 5,
    ) -> RecallResult:
        """
        Retrieve memories with fallback to clarification question.

        If confidence is below threshold, returns a clarification question
        instead of unreliable results.

        Args:
            query: The search query
            fallback_question: Question to ask if recall fails
            k: Number of results to retrieve

        Returns:
            RecallResult with content or clarification flag
        """
        try:
            response = self._client.post(
                "/memory/retrieve",
                json={
                    "query": query,
                    "k": k,
                    "namespace": self.namespace,
                    "mode": "graph_hybrid",
                    "min_confidence": 0.0,  # Get all, filter locally
                },
            )
            response.raise_for_status()
            episodes = response.json().get("episodes", [])

            if not episodes:
                return RecallResult(
                    content="",
                    confidence=0.0,
                    needs_clarification=True,
                    fallback_question=fallback_question,
                    episodes=[],
                )

            # Check confidence of top result
            top_episode = episodes[0]
            confidence = top_episode.get("meta", {}).get("score", 0.5)

            if confidence < self.confidence_threshold:
                return RecallResult(
                    content="",
                    confidence=confidence,
                    needs_clarification=True,
                    fallback_question=fallback_question,
                    episodes=episodes,
                )

            # Build content from episodes
            content = self._format_episodes(episodes)

            return RecallResult(
                content=content,
                confidence=confidence,
                needs_clarification=False,
                episodes=episodes,
            )

        except httpx.HTTPError:
            return RecallResult(
                content="",
                confidence=0.0,
                needs_clarification=True,
                fallback_question=fallback_question,
                episodes=[],
            )

    def summarize_context(
        self,
        query: str,
        max_tokens: int = 2000,
        k: int = 20,
    ) -> ContextSummary:
        """
        Retrieve and summarize context within token budget.

        Retrieves relevant memories and truncates to fit within
        the specified token limit.

        Args:
            query: Query to find relevant context
            max_tokens: Maximum tokens in summary
            k: Number of memories to consider

        Returns:
            ContextSummary with formatted text
        """
        try:
            response = self._client.post(
                "/memory/retrieve",
                json={
                    "query": query,
                    "k": k,
                    "namespace": self.namespace,
                    "mode": "graph_hybrid",
                },
            )
            response.raise_for_status()
            episodes = response.json().get("episodes", [])

            if not episodes:
                return ContextSummary(
                    text="",
                    token_count=0,
                    episodes_included=0,
                )

            # Build context within token budget
            context_parts = []
            total_tokens = 0
            episodes_included = 0
            truncated = False

            for ep in episodes:
                content = ep.get("meta", {}).get("text", "")
                # Rough token estimate: ~4 chars per token
                episode_tokens = len(content) // 4

                if total_tokens + episode_tokens > max_tokens:
                    truncated = True
                    break

                context_parts.append(content)
                total_tokens += episode_tokens
                episodes_included += 1

            return ContextSummary(
                text="\n\n".join(context_parts),
                token_count=total_tokens,
                episodes_included=episodes_included,
                truncated=truncated,
            )

        except httpx.HTTPError:
            return ContextSummary(
                text="",
                token_count=0,
                episodes_included=0,
            )

    def forget_safely(
        self,
        episode_id: str,
        reason: Optional[str] = None,
    ) -> bool:
        """
        Delete a memory with audit trail.

        Args:
            episode_id: ID of episode to delete
            reason: Reason for deletion (for audit)

        Returns:
            True if deleted successfully
        """
        try:
            # First, store audit record
            self._client.post(
                "/memory/insert",
                json={
                    "text": f"Deleted episode {episode_id}",
                    "namespace": f"{self.namespace}_audit",
                    "meta": {
                        "type": "deletion_audit",
                        "deleted_episode_id": episode_id,
                        "reason": reason or "User requested",
                    },
                },
            )

            # Then delete
            response = self._client.delete(
                f"/memory/episode/{episode_id}",
            )
            response.raise_for_status()
            return True

        except httpx.HTTPError:
            return False

    def _assess_importance(
        self,
        content: str,
        context: Optional[str] = None,
    ) -> float:
        """
        Assess importance of content.

        Uses LLM if provided, otherwise falls back to heuristics.
        """
        if self.importance_llm:
            return self.importance_llm(content, context or "")

        # Simple heuristics
        score = 0.3  # Base score

        # Preference indicators
        if any(kw in content.lower() for kw in ["prefer", "like", "want", "need"]):
            score += 0.3

        # Factual indicators
        if any(kw in content.lower() for kw in ["is", "are", "was", "name", "email"]):
            score += 0.2

        # Length bonus (longer = more substantial)
        if len(content) > 50:
            score += 0.1

        return min(score, 1.0)

    def _format_episodes(self, episodes: List[Dict]) -> str:
        """Format episodes as readable text."""
        parts = []
        for ep in episodes:
            content = ep.get("meta", {}).get("text", "")
            if content:
                parts.append(content)
        return "\n\n".join(parts)

    def __del__(self) -> None:
        if hasattr(self, "_client") and self._client:
            self._client.close()
