"""
Synrix Agent Runtime — AutoGen Integration
Persistent conversation history for AutoGen agents.

Usage:
    from synrix_runtime.integrations.autogen_memory import SynrixAutoGenMemory
    memory = SynrixAutoGenMemory(group_id="my_autogen_group")
"""

import time
import json
from typing import Dict, List, Optional, Any
from synrix.agent_backend import get_synrix_backend


class SynrixAutoGenMemory:
    """Persistent conversation history for AutoGen agents."""

    def __init__(self, group_id: str = "default", backend=None):
        self.group_id = group_id
        if backend is not None:
            self.backend = backend
        else:
            from synrix_runtime.config import SynrixConfig
            config = SynrixConfig.from_env()
            self.backend = get_synrix_backend(**config.get_backend_kwargs())

    def store_message(self, sender: str, recipient: str, content: str, timestamp: float = None):
        """Store a message in the conversation history."""
        if timestamp is None:
            timestamp = time.time()
        ts = int(timestamp * 1000000)

        entry = {
            "sender": sender,
            "recipient": recipient,
            "content": content,
            "timestamp": timestamp,
        }

        start = time.perf_counter_ns()
        self.backend.write(
            f"autogen:{self.group_id}:messages:{sender}:{recipient}:{ts}",
            entry,
            metadata={"type": "autogen_message", "sender": sender, "recipient": recipient}
        )
        latency_us = (time.perf_counter_ns() - start) / 1000
        return {"latency_us": latency_us}

    def get_conversation_history(self, agent_pair: tuple = None, limit: int = 100) -> list:
        """Get conversation history, optionally between a specific pair."""
        if agent_pair:
            sender, recipient = agent_pair
            prefix = f"autogen:{self.group_id}:messages:{sender}:{recipient}:"
        else:
            prefix = f"autogen:{self.group_id}:messages:"

        results = self.backend.query_prefix(prefix, limit=limit)
        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("timestamp", 0))
        return messages

    def search_conversations(self, query_text: str) -> list:
        """Search through all conversations for matching content."""
        all_messages = self.get_conversation_history(limit=500)
        matches = []
        query_lower = query_text.lower()
        for msg in all_messages:
            content = msg.get("content", "")
            if query_lower in str(content).lower():
                matches.append(msg)
        return matches

    def get_agent_knowledge(self, agent_name: str) -> list:
        """Get all messages sent by a specific agent."""
        results = self.backend.query_prefix(f"autogen:{self.group_id}:messages:{agent_name}:", limit=200)
        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("timestamp", 0))
        return messages

    def export_conversation(self, agent_pair: tuple = None, format: str = "json") -> str:
        """Export conversation history."""
        messages = self.get_conversation_history(agent_pair)
        if format == "json":
            return json.dumps({"group_id": self.group_id, "messages": messages, "exported_at": time.time()}, indent=2)
        else:
            lines = []
            for msg in messages:
                lines.append(f"[{msg.get('sender')} -> {msg.get('recipient')}] {msg.get('content')}")
            return "\n".join(lines)

    def get_stats(self) -> dict:
        """Get conversation statistics."""
        messages = self.get_conversation_history(limit=500)
        agents = set()
        for msg in messages:
            agents.add(msg.get("sender", ""))
            agents.add(msg.get("recipient", ""))
        agents.discard("")

        return {
            "group_id": self.group_id,
            "total_messages": len(messages),
            "unique_agents": len(agents),
            "agents": list(agents),
        }
