"""
Synrix Agent Runtime — LangChain Integration
Drop-in replacement for LangChain memory modules.

Usage:
    from synrix_runtime.integrations.langchain_memory import SynrixMemory
    memory = SynrixMemory(agent_id="my_chain")
    chain = ConversationChain(llm=llm, memory=memory)
"""

import time
import json
from typing import Dict, List, Any, Optional

try:
    from langchain.memory import BaseMemory
    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False
    class BaseMemory:
        """Stub when LangChain not installed."""
        memory_variables = []
        def save_context(self, inputs, outputs): pass
        def load_memory_variables(self, inputs): return {}
        def clear(self): pass


class SynrixMemory(BaseMemory):
    """
    Drop-in replacement for LangChain memory backed by Synrix.

    Usage:
        from synrix_runtime.integrations.langchain_memory import SynrixMemory
        memory = SynrixMemory(agent_id="my_chain")
        # Use with LangChain: ConversationChain(llm=llm, memory=memory)
    """

    def __init__(self, agent_id: str = "langchain_default", memory_key: str = "history", backend=None, **kwargs):
        self.agent_id = agent_id
        self.memory_key = memory_key
        self._memory_variables = [memory_key]

        if backend is not None:
            self.backend = backend
        else:
            from synrix.agent_backend import get_synrix_backend
            from synrix_runtime.config import SynrixConfig
            config = SynrixConfig.from_env()
            self.backend = get_synrix_backend(**config.get_backend_kwargs())
        self._message_count = 0

    @property
    def memory_variables(self) -> List[str]:
        return self._memory_variables

    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        """Save conversation turn to Synrix."""
        ts = int(time.time() * 1000000)
        self._message_count += 1

        entry = {
            "inputs": inputs,
            "outputs": outputs,
            "turn": self._message_count,
            "timestamp": time.time(),
        }

        start = time.perf_counter_ns()
        self.backend.write(
            f"langchain:{self.agent_id}:messages:{ts}",
            entry,
            metadata={"type": "langchain_message", "agent_id": self.agent_id}
        )
        latency_us = (time.perf_counter_ns() - start) / 1000

        # Update summary
        self.backend.write(
            f"langchain:{self.agent_id}:summary",
            {"total_turns": self._message_count, "last_updated": time.time()},
            metadata={"type": "langchain_summary"}
        )

    def load_memory_variables(self, inputs: Dict[str, Any] = None) -> Dict[str, Any]:
        """Load conversation history from Synrix."""
        start = time.perf_counter_ns()
        results = self.backend.query_prefix(f"langchain:{self.agent_id}:messages:", limit=100)
        latency_us = (time.perf_counter_ns() - start) / 1000

        messages = []
        for r in results:
            data = r.get("data", {})
            val = data.get("value", data)
            if isinstance(val, dict):
                messages.append(val)

        messages.sort(key=lambda x: x.get("turn", 0))

        # Build conversation string
        history_parts = []
        for msg in messages:
            inp = msg.get("inputs", {})
            out = msg.get("outputs", {})
            human_input = inp.get("input", inp.get("human_input", str(inp)))
            ai_output = out.get("output", out.get("response", str(out)))
            history_parts.append(f"Human: {human_input}")
            history_parts.append(f"AI: {ai_output}")

        return {self.memory_key: "\n".join(history_parts)}

    def clear(self) -> None:
        """Clear is a no-op — Synrix preserves all history."""
        pass

    def get_full_history(self) -> List[dict]:
        """Get complete conversation history."""
        results = self.backend.query_prefix(f"langchain:{self.agent_id}:messages:", limit=500)
        messages = []
        for r in results:
            data = r.get("data", {})
            val = data.get("value", data)
            if isinstance(val, dict):
                messages.append(val)
        messages.sort(key=lambda x: x.get("turn", 0))
        return messages

    def restore_from_crash(self) -> int:
        """Restore message count from Synrix after a crash."""
        messages = self.get_full_history()
        self._message_count = len(messages)
        return self._message_count

    def export_conversation(self) -> dict:
        """Export full conversation with metadata."""
        messages = self.get_full_history()
        summary = self.backend.read(f"langchain:{self.agent_id}:summary")

        return {
            "agent_id": self.agent_id,
            "total_turns": len(messages),
            "messages": messages,
            "summary": summary.get("data", {}).get("value") if summary else None,
            "exported_at": time.time(),
        }

    def store_entity(self, entity_name: str, entity_data: dict):
        """Store an entity in the knowledge base."""
        self.backend.write(
            f"langchain:{self.agent_id}:entities:{entity_name}",
            entity_data,
            metadata={"type": "langchain_entity"}
        )

    def get_entity(self, entity_name: str) -> Optional[dict]:
        """Retrieve an entity."""
        result = self.backend.read(f"langchain:{self.agent_id}:entities:{entity_name}")
        if result:
            data = result.get("data", {})
            return data.get("value", data)
        return None

    def get_all_entities(self) -> List[dict]:
        """Get all stored entities."""
        results = self.backend.query_prefix(f"langchain:{self.agent_id}:entities:", limit=200)
        entities = []
        for r in results:
            key = r.get("key", "").replace(f"langchain:{self.agent_id}:entities:", "")
            data = r.get("data", {})
            val = data.get("value", data)
            entities.append({"name": key, "data": val})
        return entities
