"""
communication/channel.py - Message Channel
===========================================
LocalChannel: same-machine agent-to-agent via asyncio queues.
RemoteChannel: cross-machine communication via WebSocket (AICQ relay).
"""
from __future__ import annotations

import asyncio
import json
import time
import uuid
import logging
from dataclasses import dataclass, field, asdict
from typing import Optional, List, Callable, Awaitable, Any

from .crypto import (
    derive_shared_key,
    generate_x25519_keypair,
    x25519_key_exchange,
    encrypt_message,
    decrypt_message,
    sign,
    verify,
    serialize_public_key,
    deserialize_public_key,
)

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


# =========================================================================
# Message Data Model
# =========================================================================

@dataclass
class Message:
    """A communication message between agents."""

    id: str = ""
    from_agent: str = ""          # Sender 的数字 aid
    to_agent: str = ""            # Recipient 的数字 aid
    content: str = ""             # Plaintext content
    timestamp: float = 0.0
    encrypted: bool = False
    signature: str = ""           # Hex signature (from sender)
    msg_type: str = "text"        # text | handshake | heartbeat | ack

    def __post_init__(self):
        if not self.id:
            self.id = uuid.uuid4().hex[:16]
        if self.timestamp == 0.0:
            self.timestamp = time.time()

    def to_dict(self) -> dict:
        d = asdict(self)
        return d

    @classmethod
    def from_dict(cls, data: dict) -> "Message":
        return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), ensure_ascii=False)


# =========================================================================
# Message Callback
# =========================================================================

MessageCallback = Callable[[Message], Awaitable[None]]


# =========================================================================
# Local Channel — Same-Machine Agent Communication
# =========================================================================

class LocalChannel:
    """
    LocalChannel enables encrypted message passing between agents on the
    same machine using asyncio queues.

    Each agent instance has its own ``LocalChannel``.  Messages are routed
    through a shared broker (dict of queues) keyed by agent_id.
    """

    # Shared broker across all LocalChannel instances on this process
    _queues: dict[str, asyncio.Queue] = {}
    _callbacks: dict[str, List[MessageCallback]] = {}

    def __init__(
        self,
        agent_id: str,
        ed25519_private_key: bytes,
        ed25519_public_key: bytes,
    ):
        self.agent_id = agent_id
        self._ed_priv = ed25519_private_key
        self._ed_pub = ed25519_public_key

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

        # Ensure our queue exists
        if agent_id not in LocalChannel._queues:
            LocalChannel._queues[agent_id] = asyncio.Queue()
        if agent_id not in LocalChannel._callbacks:
            LocalChannel._callbacks[agent_id] = []

        self._running = False
        self._receiver_task: Optional[asyncio.Task] = None

    # ------------------------------------------------------------------
    # Derive a shared AES key with another agent
    # ------------------------------------------------------------------

    def _derive_key(self, their_ed_pub: bytes) -> bytes:
        their_x_pub = ed25519_to_x25519_public(their_ed_pub)
        shared_secret = None
        # Simple DH: we need the actual x25519 exchange function
        from .crypto import x25519_key_exchange
        shared_secret = x25519_key_exchange(self._x_priv, their_x_pub)
        return derive_shared_key(shared_secret)

    # ------------------------------------------------------------------
    # Send
    # ------------------------------------------------------------------

    async def send(
        self,
        to_agent: str,
        content: str,
        msg_type: str = "text",
        their_public_key: Optional[bytes] = None,
    ) -> Message:
        """Send an encrypted message to a local peer."""
        msg = Message(
            from_agent=self.agent_id,
            to_agent=to_agent,
            content=content,
            msg_type=msg_type,
        )

        # Sign the message
        msg_data = json.dumps({
            "id": msg.id,
            "from": msg.from_agent,
            "to": msg.to_agent,
            "content": msg.content,
            "timestamp": msg.timestamp,
            "type": msg.msg_type,
        }, sort_keys=True).encode("utf-8")
        sig = sign(self._ed_priv, msg_data)
        msg.signature = sig.hex()

        # Note: encryption requires a pre-established shared key.
        # In local mode, messages are already in-process and signed.
        # For real E2EE, the shared key must be established via DH with
        # the peer's X25519 public key (separate from Ed25519 identity).
        msg.encrypted = False

        # Put into recipient's queue
        queue = LocalChannel._queues.get(to_agent)
        if queue is None:
            raise RuntimeError(f"LocalChannel: agent '{to_agent}' has no queue")

        await queue.put(msg)
        return msg

    # ------------------------------------------------------------------
    # Receive
    # ------------------------------------------------------------------

    def on_message(self, callback: MessageCallback):
        """Register a callback for incoming messages."""
        LocalChannel._callbacks.setdefault(self.agent_id, []).append(callback)

    async def _dispatch(self, msg: Message):
        """Dispatch a received message to registered callbacks."""
        for cb in LocalChannel._callbacks.get(self.agent_id, []):
            try:
                await cb(msg)
            except Exception as e:
                logger.error(f"LocalChannel callback error: {e}", exc_info=True)

    # ------------------------------------------------------------------
    # Receiver loop
    # ------------------------------------------------------------------

    async def _receiver_loop(self):
        """Background task: read from queue and dispatch."""
        queue = LocalChannel._queues.get(self.agent_id)
        if not queue:
            return
        while self._running:
            try:
                msg = await asyncio.wait_for(queue.get(), timeout=1.0)
                await self._dispatch(msg)
            except asyncio.TimeoutError:
                continue
            except Exception as e:
                logger.error(f"LocalChannel receiver error: {e}")

    def start(self):
        """Start the receiver loop."""
        if self._running:
            return
        self._running = True
        self._receiver_task = asyncio.create_task(self._receiver_loop())

    async def stop(self):
        """Stop the receiver loop."""
        self._running = False
        if self._receiver_task:
            self._receiver_task.cancel()
            try:
                await self._receiver_task
            except asyncio.CancelledError:
                pass
            self._receiver_task = None

    # ------------------------------------------------------------------
    # Static: register a new agent queue
    # ------------------------------------------------------------------

    @staticmethod
    def register_agent(agent_id: str):
        """Register a new agent in the shared broker."""
        if agent_id not in LocalChannel._queues:
            LocalChannel._queues[agent_id] = asyncio.Queue()
        if agent_id not in LocalChannel._callbacks:
            LocalChannel._callbacks[agent_id] = []


# =========================================================================
# Remote Channel — Cross-Machine via WebSocket
# =========================================================================

class RemoteChannel:
    """
    RemoteChannel connects to the AICQ relay server via WebSocket for
    cross-machine agent communication.

    Features:
    - Encrypted messaging via AES-256-GCM
    - Offline message queue (SQLite)
    - Automatic reconnection
    - Heartbeat for online status
    """

    def __init__(
        self,
        agent_id: str,
        ed25519_private_key: bytes,
        ed25519_public_key: bytes,
        server_url: str = "wss://aicq.online/ws",
        queue_db_path: str = "",
    ):
        self.agent_id = agent_id
        self._ed_priv = ed25519_private_key
        self._ed_pub = ed25519_public_key
        self.server_url = server_url
        self._queue_db_path = queue_db_path

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

        self._ws = None
        self._running = False
        self._connected = False
        self._recv_task: Optional[asyncio.Task] = None
        self._heartbeat_task: Optional[asyncio.Task] = None
        self._retry_task: Optional[asyncio.Task] = None
        self._callbacks: List[MessageCallback] = []
        self._shared_keys: dict[str, bytes] = {}  # agent_id -> AES key
        self._pending_handshakes: dict[str, asyncio.Future] = {}

        # Offline queue
        self._queue_conn = None

    # ------------------------------------------------------------------
    # Offline Queue (SQLite)
    # ------------------------------------------------------------------

    def _init_queue_db(self):
        """Initialize the offline message queue SQLite database."""
        if not self._queue_db_path:
            return
        import sqlite3
        self._queue_conn = sqlite3.connect(self._queue_db_path)
        self._queue_conn.execute("""
            CREATE TABLE IF NOT EXISTS queued_messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                from_agent TEXT NOT NULL,
                to_agent TEXT NOT NULL,
                encrypted_content TEXT NOT NULL,
                created_at REAL NOT NULL,
                delivered INTEGER DEFAULT 0
            )
        """)
        self._queue_conn.commit()

    def _queue_message(self, from_agent: str, to_agent: str, encrypted_content: str):
        """Queue a message for later delivery."""
        if not self._queue_conn:
            return
        self._queue_conn.execute(
            "INSERT INTO queued_messages (from_agent, to_agent, encrypted_content, created_at, delivered) VALUES (?, ?, ?, ?, 0)",
            (from_agent, to_agent, encrypted_content, time.time()),
        )
        self._queue_conn.commit()

    def _get_undelivered(self) -> list:
        """Get undelivered messages."""
        if not self._queue_conn:
            return []
        rows = self._queue_conn.execute(
            "SELECT id, from_agent, to_agent, encrypted_content, created_at FROM queued_messages WHERE delivered = 0 ORDER BY created_at ASC"
        ).fetchall()
        return [{"id": r[0], "from": r[1], "to": r[2], "content": r[3], "created": r[4]} for r in rows]

    def _mark_delivered(self, msg_id: int):
        """Mark a message as delivered."""
        if not self._queue_conn:
            return
        self._queue_conn.execute(
            "UPDATE queued_messages SET delivered = 1 WHERE id = ?", (msg_id,)
        )
        self._queue_conn.commit()

    # ------------------------------------------------------------------
    # Key Management
    # ------------------------------------------------------------------

    def _get_or_derive_key(self, their_agent_id: str, their_x25519_pub: bytes) -> bytes:
        """Get or derive a shared key with a peer using their X25519 public key."""
        if their_agent_id in self._shared_keys:
            return self._shared_keys[their_agent_id]
        shared = x25519_key_exchange(self._x_priv, their_x25519_pub)
        key = derive_shared_key(shared)
        self._shared_keys[their_agent_id] = key
        return key

    # ------------------------------------------------------------------
    # WebSocket Connection
    # ------------------------------------------------------------------

    async def connect(self):
        """Connect to the AICQ relay server."""
        try:
            import websockets
            self._ws = await websockets.connect(
                self.server_url,
                ping_interval=30,
                ping_timeout=10,
                close_timeout=5,
            )
            self._connected = True
            logger.info(f"RemoteChannel: connected to {self.server_url}")

            # Send auth/register message (includes both identity and DH key)
            auth_msg = json.dumps({
                "type": "register",
                "agent_id": self.agent_id,
                "public_key": serialize_public_key(self._ed_pub),
                "dh_public_key": self._x_pub.hex(),
            })
            await self._ws.send(auth_msg)

        except Exception as e:
            logger.warning(f"RemoteChannel: connection failed: {e}")
            self._connected = False

    async def disconnect(self):
        """Disconnect from the relay server."""
        self._connected = False
        if self._ws:
            try:
                await self._ws.close()
            except Exception:
                pass
            self._ws = None

    # ------------------------------------------------------------------
    # Send
    # ------------------------------------------------------------------

    async def send(
        self,
        to_agent: str,
        content: str,
        their_ed_pub_hex: str = "",
        their_x25519_pub_hex: str = "",
        msg_type: str = "text",
    ) -> Message:
        """Send an encrypted message to a remote peer."""
        msg = Message(
            from_agent=self.agent_id,
            to_agent=to_agent,
            content=content,
            msg_type=msg_type,
        )

        # Sign the message
        msg_data = json.dumps({
            "id": msg.id,
            "from": msg.from_agent,
            "to": msg.to_agent,
            "content": msg.content,
            "timestamp": msg.timestamp,
            "type": msg.msg_type,
        }, sort_keys=True).encode("utf-8")
        sig = sign(self._ed_priv, msg_data)
        msg.signature = sig.hex()

        # Encrypt using X25519 DH shared key
        if their_x25519_pub_hex:
            try:
                their_x_pub = bytes.fromhex(their_x25519_pub_hex)
                shared_key = self._get_or_derive_key(to_agent, their_x_pub)
                encrypted_payload = encrypt_message(shared_key, {
                    "content": content,
                    "timestamp": msg.timestamp,
                    "msg_type": msg.msg_type,
                    "msg_id": msg.id,
                    "signature": msg.signature,
                })
                msg.content = encrypted_payload
                msg.encrypted = True
            except Exception as e:
                logger.warning(f"RemoteChannel encryption failed: {e}")
                # [v1.23.58] 加密失败不应静默降级为明文发送，改为拒绝发送
                raise RuntimeError(f"消息加密失败（目标: {to_agent}）: {e}") from e
        else:
            msg.encrypted = False

        # Send via WebSocket or queue
        wire_msg = json.dumps({
            "type": msg_type,
            "from": msg.from_agent,
            "to": msg.to_agent,
            "content": msg.content,
            "encrypted": msg.encrypted,
            "timestamp": msg.timestamp,
            "msg_id": msg.id,
            "signature": msg.signature,
        })

        if self._connected and self._ws:
            try:
                await self._ws.send(wire_msg)
            except Exception as e:
                logger.warning(f"RemoteChannel: send failed, queuing: {e}")
                self._queue_message(msg.from_agent, msg.to_agent, wire_msg)
                self._connected = False
        else:
            self._queue_message(msg.from_agent, msg.to_agent, wire_msg)

        return msg

    # ------------------------------------------------------------------
    # Receive Loop
    # ------------------------------------------------------------------

    async def _receiver_loop(self):
        """Background task: receive messages from WebSocket."""
        while self._running:
            if not self._connected:
                try:
                    await self.connect()
                except Exception:
                    await asyncio.sleep(5)
                    continue

            try:
                raw = await self._ws.recv()
                data = json.loads(raw)
                await self._handle_incoming(data)
            except Exception as e:
                if self._running:
                    logger.warning(f"RemoteChannel recv error: {e}")
                    self._connected = False
                await asyncio.sleep(1)

    async def _handle_incoming(self, data: dict):
        """Handle an incoming message from the relay."""
        msg_type = data.get("type", "text")

        if msg_type == "heartbeat_ack":
            return

        msg = Message(
            id=data.get("msg_id", ""),
            from_agent=data.get("from", ""),
            to_agent=data.get("to", ""),
            content=data.get("content", ""),
            timestamp=data.get("timestamp", 0.0),
            encrypted=data.get("encrypted", False),
            signature=data.get("signature", ""),
            msg_type=msg_type,
        )

        # Dispatch to callbacks
        for cb in self._callbacks:
            try:
                await cb(msg)
            except Exception as e:
                logger.error(f"RemoteChannel callback error: {e}")

    # ------------------------------------------------------------------
    # Heartbeat
    # ------------------------------------------------------------------

    async def _heartbeat_loop(self):
        """Periodically send heartbeat messages."""
        while self._running:
            await asyncio.sleep(30)
            if self._connected and self._ws:
                try:
                    await self._ws.send(json.dumps({"type": "heartbeat"}))
                except Exception:
                    self._connected = False

    # ------------------------------------------------------------------
    # Retry Undelivered Messages
    # ------------------------------------------------------------------

    async def _retry_loop(self):
        """Periodically retry sending queued messages."""
        while self._running:
            await asyncio.sleep(10)
            if not self._connected or not self._ws:
                continue
            undelivered = self._get_undelivered()
            for item in undelivered:
                try:
                    await self._ws.send(item["content"])
                    self._mark_delivered(item["id"])
                except Exception:
                    break  # Stop retrying if send fails

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

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

    def start(self):
        """Start the remote channel (receiver, heartbeat, retry loops)."""
        if self._running:
            return
        self._running = True
        self._init_queue_db()
        self._recv_task = asyncio.create_task(self._receiver_loop())
        self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
        self._retry_task = asyncio.create_task(self._retry_loop())

    async def stop(self):
        """Stop the remote channel."""
        self._running = False
        await self.disconnect()
        if self._recv_task:
            self._recv_task.cancel()
            try:
                await self._recv_task
            except asyncio.CancelledError:
                pass
        if self._heartbeat_task:
            self._heartbeat_task.cancel()
            try:
                await self._heartbeat_task
            except asyncio.CancelledError:
                pass
        if self._retry_task:
            self._retry_task.cancel()
            try:
                await self._retry_task
            except asyncio.CancelledError:
                pass
        if self._queue_conn:
            self._queue_conn.close()
            self._queue_conn = None