"""
communication/manager.py - Communication Manager
================================================
Central manager for agent-to-agent encrypted communication.
Orchestrates crypto keys, peer store, local/remote channels, and offline queue.
"""
from __future__ import annotations

import asyncio
import logging
import time
from pathlib import Path
from typing import Optional, List, Callable, Awaitable

from .crypto import (
    generate_ed25519_keypair,
    generate_agent_id,
    serialize_public_key,
    serialize_private_key,
    deserialize_public_key,
    deserialize_private_key,
    decrypt_message,
    verify,
    generate_x25519_keypair,
    x25519_key_exchange,
    derive_shared_key,
)
from .peer import Peer, PeerStore
from .channel import (
    Message,
    MessageCallback,
    LocalChannel,
    RemoteChannel,
)

logger = logging.getLogger("myagent.communication")


class CommunicationManager:
    """
    Manages all aspects of agent-to-agent encrypted communication.

    Responsibilities:
      - Initialize/load Ed25519 identity keys
      - Manage the peer store (add/remove/list)
      - Run local channel (same-machine messaging)
      - Run remote channel (cross-machine via AICQ WebSocket)
      - Route incoming messages to registered callbacks
      - Offline message queue with automatic retry
    """

    def __init__(
        self,
        config,
        data_dir: str | Path = "",
    ):
        """
        Args:
            config: A ``CommunicationConfig`` dataclass instance.
            data_dir: Base data directory for persistent storage.
        """
        self.config = config
        self._data_dir = Path(data_dir) if data_dir else Path.home() / ".myagent" / "data"
        self._permission_checker = None  # callable(perm_name) -> bool, set by app

        # Identity keys (Ed25519)
        self._ed_priv: bytes = b""
        self._ed_pub: bytes = b""
        self.agent_id: str = ""

        # DH keys (X25519, standalone from Ed25519)
        self._x_priv: bytes = b""
        self._x_pub: bytes = b""

        # Components
        self._peer_store: Optional[PeerStore] = None
        self._local_channel: Optional[LocalChannel] = None
        self._remote_channel: Optional[RemoteChannel] = None

        # State
        self._running = False
        self._message_callbacks: List[MessageCallback] = []
        self._received_messages: List[Message] = []  # In-memory recent messages
        self._max_stored_messages = 500

        # Paths
        self._peers_path = self._data_dir / "communication" / "peers.json"
        self._queue_db_path = self._data_dir / "communication" / "message_queue.db"

    # ==================================================================
    # Initialization
    # ==================================================================

    def initialize(self):
        """Load or generate identity keys, initialize peer store."""
        # Ensure directories
        self._peers_path.parent.mkdir(parents=True, exist_ok=True)
        self._queue_db_path.parent.mkdir(parents=True, exist_ok=True)

        # Load or generate Ed25519 identity keys
        if self.config.agent_id and self.config.private_key:
            self._ed_priv = deserialize_private_key(self.config.private_key)
            self._ed_pub = deserialize_public_key(self.config.agent_id)
            self.agent_id = generate_agent_id(self._ed_pub)
        else:
            self._ed_priv, self._ed_pub = generate_ed25519_keypair()
            self.agent_id = generate_agent_id(self._ed_pub)
            self.config.agent_id = serialize_public_key(self._ed_pub)
            self.config.private_key = serialize_private_key(self._ed_priv)
            logger.info(f"CommunicationManager: generated new identity, agent_id={self.agent_id}")

        # Generate standalone X25519 keypair for DH
        self._x_priv, self._x_pub = generate_x25519_keypair()

        # Initialize peer store
        self._peer_store = PeerStore(self._peers_path)

        logger.info(
            f"CommunicationManager initialized: agent_id={self.agent_id}, "
            f"peers={self._peer_store.peer_count()}, "
            f"has_crypto={True}"
        )

    async def start(self):
        """Start communication channels (must be called from async context)."""
        if not self.config.enabled:
            logger.info("CommunicationManager: disabled by config")
            return

        if self._running:
            return

        self._running = True

        # Initialize local channel
        LocalChannel.register_agent(self.agent_id)
        self._local_channel = LocalChannel(
            agent_id=self.agent_id,
            ed25519_private_key=self._ed_priv,
            ed25519_public_key=self._ed_pub,
        )
        self._local_channel.on_message(self._on_local_message)
        self._local_channel.start()

        # Initialize remote channel if server URL configured
        if self.config.server_url:
            self._remote_channel = RemoteChannel(
                agent_id=self.agent_id,
                ed25519_private_key=self._ed_priv,
                ed25519_public_key=self._ed_pub,
                server_url=self.config.server_url,
                queue_db_path=str(self._queue_db_path),
            )
            self._remote_channel.on_message(self._on_remote_message)
            self._remote_channel.start()

        logger.info("CommunicationManager started")

    async def stop(self):
        """Stop all communication channels."""
        self._running = False
        if self._local_channel:
            await self._local_channel.stop()
        if self._remote_channel:
            await self._remote_channel.stop()
        logger.info("CommunicationManager stopped")

    # ==================================================================
    # Message Handlers
    # ==================================================================

    async def _on_local_message(self, msg: Message):
        """Handle a message received via LocalChannel."""
        await self._dispatch_message(msg)

    async def _on_remote_message(self, msg: Message):
        """Handle a message received via RemoteChannel."""
        # If encrypted, attempt to decrypt
        if msg.encrypted and msg.from_agent:
            peer = self._peer_store.get_peer(msg.from_agent) if self._peer_store else None
            if peer and peer.public_key:
                try:
                    # Use the peer's stored X25519 public key for decryption
                    their_x_pub_hex = peer.permissions.get("x25519_public_key", "")
                    if their_x_pub_hex:
                        their_x_pub = bytes.fromhex(their_x_pub_hex)
                        shared = x25519_key_exchange(self._x_priv, their_x_pub)
                        shared_key = derive_shared_key(shared)
                        payload = decrypt_message(shared_key, msg.content)
                        their_ed_pub = deserialize_public_key(peer.public_key)
                        if payload.get("signature"):
                            sig = bytes.fromhex(payload["signature"])
                            if verify(their_ed_pub, msg.content.encode("utf-8"), sig):
                                msg.content = payload.get("content", msg.content)
                                msg.encrypted = False
                            else:
                                logger.warning(f"Remote message from {msg.from_agent}: signature verification failed")
                        else:
                            msg.content = payload.get("content", msg.content)
                            msg.encrypted = False
                    else:
                        logger.debug(f"No X25519 key for peer {msg.from_agent}, cannot decrypt")
                except Exception as e:
                    logger.warning(f"Remote message decrypt failed: {e}")

        await self._dispatch_message(msg)

    async def _dispatch_message(self, msg: Message):
        """Store message and dispatch to callbacks."""
        # Store in memory
        self._received_messages.append(msg)
        if len(self._received_messages) > self._max_stored_messages:
            self._received_messages = self._received_messages[-self._max_stored_messages:]

        # Notify callbacks
        for cb in self._message_callbacks:
            try:
                await cb(msg)
            except Exception as e:
                logger.error(f"Message callback error: {e}", exc_info=True)

    # ==================================================================
    # Send Message
    # ==================================================================

    async def send_message(
        self,
        to_agent: str,
        content: str,
        msg_type: str = "text",
    ) -> Message:
        """
        Send an encrypted message to another agent.

        Args:
            to_agent: Recipient agent ID.
            content: Plaintext message content.
            msg_type: Message type (text, heartbeat, ack, etc.)

        Returns:
            The sent Message object.

        Raises:
            RuntimeError: If the peer is not found.
        """
        if not self._running or not self.config.enabled:
            raise RuntimeError("Communication is not enabled")

        peer = self._peer_store.get_peer(to_agent) if self._peer_store else None
        if not peer:
            raise RuntimeError(f"Peer '{to_agent}' not found")

        their_pub = deserialize_public_key(peer.public_key)

        # Try local first, then remote
        msg = None
        # Check if the peer is on this machine (has a local queue)
        if to_agent in LocalChannel._queues:
            # 本机通信权限检查
            if not self._check_permission("local_comm"):
                raise RuntimeError("[权限] 当前 Agent 没有本机通信权限")
            msg = await self._local_channel.send(
                to_agent=to_agent,
                content=content,
                msg_type=msg_type,
                their_public_key=their_pub,
            )
        elif self._remote_channel:
            # 跨电脑通信权限检查
            if not self._check_permission("remote_comm"):
                raise RuntimeError("[权限] 当前 Agent 没有跨电脑通信权限，请先在权限管理中开启")
            their_x_pub_hex = peer.permissions.get("x25519_public_key", "")
            msg = await self._remote_channel.send(
                to_agent=to_agent,
                content=content,
                their_ed_pub_hex=peer.public_key,
                their_x25519_pub_hex=their_x_pub_hex,
                msg_type=msg_type,
            )
        else:
            # [v1.23.58] 修复：无通道时消息直接丢失（原代码检查 _remote_channel 但此处必为 None）
            # 改为记录警告并抛出异常，让调用方知道消息未送达
            logger.warning(
                f"无法发送消息给 {to_agent}：目标不在本机且无远程通道。"
                f"请确认对方已上线或检查远程通信配置。"
            )
            raise RuntimeError(
                f"无法投递消息给 '{to_agent}'：目标不在本机（LocalChannel 无队列），"
                f"且远程通道未配置或未连接。消息未送达。"
            )

        return msg

    # ==================================================================
    # Peer Management
    # ==================================================================

    def add_peer(
        self,
        agent_id: str,
        public_key: str,
        display_name: str = "",
    ) -> bool:
        """
        Add a known peer.

        Args:
            agent_id: The peer's short agent ID (or full hex public key).
            public_key: The peer's Ed25519 public key in hex.
            display_name: Optional human-readable name.

        Returns:
            True if added, False if already exists.
        """
        if not self._peer_store:
            return False

        peer = Peer(
            agent_id=agent_id,
            public_key=public_key,
            display_name=display_name or agent_id,
            added_at=time.time(),
        )
        return self._peer_store.add_peer(peer)

    def remove_peer(self, agent_id: str) -> bool:
        """Remove a peer by agent_id."""
        if not self._peer_store:
            return False
        return self._peer_store.remove_peer(agent_id)

    def get_peer(self, agent_id: str) -> Optional[Peer]:
        """Get a peer by agent_id."""
        if not self._peer_store:
            return None
        return self._peer_store.get_peer(agent_id)

    def get_peer_list(self) -> List[Peer]:
        """Return all known peers."""
        if not self._peer_store:
            return []
        return self._peer_store.list_peers()

    # ==================================================================
    # Message History
    # ==================================================================

    def get_messages(
        self,
        limit: int = 50,
        from_agent: str = "",
        to_agent: str = "",
    ) -> List[dict]:
        """Get recent received messages with optional filtering."""
        msgs = list(self._received_messages)
        if from_agent:
            msgs = [m for m in msgs if m.from_agent == from_agent]
        if to_agent:
            msgs = [m for m in msgs if m.to_agent == to_agent]
        msgs = msgs[-limit:]
        return [m.to_dict() for m in msgs]

    # ==================================================================
    # Callbacks
    # ==================================================================

    def on_message(self, callback: MessageCallback):
        """Register a callback for incoming messages."""
        self._message_callbacks.append(callback)

    def set_permission_checker(self, checker):
        """设置权限检查回调（由应用启动时注入）"""
        self._permission_checker = checker

    def _check_permission(self, permission: str) -> bool:
        """检查权限，无权限检查器时默认通过"""
        if self._permission_checker is None:
            return True
        try:
            return bool(self._permission_checker(permission))
        except Exception as e:
            # [v1.23.58] 权限检查器异常时记录错误并拒绝（而非默认放行）
            logger.error(f"权限检查器异常（permission={permission}）: {e}")
            return False

    # ==================================================================
    # Status
    # ==================================================================

    def get_status(self) -> dict:
        """Return current communication status."""
        remote_connected = False
        if self._remote_channel:
            remote_connected = self._remote_channel._connected

        return {
            "enabled": self.config.enabled,
            "running": self._running,
            "agent_id": self.agent_id,
            "public_key": serialize_public_key(self._ed_pub) if self._ed_pub else "",
            "x25519_public_key": self._x_pub.hex() if self._x_pub else "",
            "peers_count": self._peer_store.peer_count() if self._peer_store else 0,
            "online_peers": len(self._peer_store.get_online_peers()) if self._peer_store else 0,
            "remote_connected": remote_connected,
            "server_url": self.config.server_url,
            "local_channel": self._local_channel is not None,
            "remote_channel": self._remote_channel is not None,
            "messages_received": len(self._received_messages),
        }