from __future__ import annotations

import asyncio
import inspect
import json
import logging
import time
from collections.abc import Iterator
from io import IOBase
from typing import Any, BinaryIO, Callable

from mrp.auth import Keypair, _b64url_decode
from mrp.client import Client
from mrp.crypto import decrypt_message, encrypt_message
from mrp.errors import RequestError
from mrp.models import (
    ACLEntry,
    AgentInfo,
    Attachment,
    Blob,
    CapabilitiesResult,
    DiscoveredAgent,
    Message,
    PollResult,
    SendResult,
)
from mrp.schemas import (
    REQUEST_CONTENT_TYPE,
    RESPONSE_CONTENT_TYPE,
    EVENT_CONTENT_TYPE,
    STATUS_CONTENT_TYPE,
    MRPRequest,
    create_request,
    create_ok_response,
    create_error_response,
    create_event,
    create_status,
    parse_request,
)
from mrp.websocket import ReconnectingWSClient

logger = logging.getLogger("mrp.agent")


class Agent:
    """High-level MRP agent with automatic lifecycle management."""

    def __init__(
        self,
        relay: str,
        key_file: str | None = None,
        keypair: Keypair | None = None,
        name: str | None = None,
        capabilities: list[str] | None = None,
        metadata: dict[str, str] | None = None,
        visibility: str | None = None,
        inbox_policy: str | None = None,
    ):
        if keypair is not None:
            self._keypair = keypair
        elif key_file is not None:
            self._keypair = Keypair.from_file(key_file, create=True)
        else:
            self._keypair = Keypair.generate()

        self._client = Client(relay, self._keypair)
        self._relay = relay
        self._name = name
        self._capabilities = capabilities
        self._metadata = metadata
        self._visibility = visibility
        self._inbox_policy = inbox_policy
        self._registered = False
        self._on_message_handler: Callable[[Message], dict | None] | None = None
        self._action_handlers: dict[str, Callable[[MRPRequest, Message], Any]] = {}

    @property
    def public_key(self) -> str:
        """Base64url public key string."""
        return self._keypair.public_key_b64

    @property
    def client(self) -> Client:
        """Access the underlying low-level client."""
        return self._client

    def _ensure_registered(self) -> None:
        """Register/update profile if not done yet."""
        if not self._registered:
            self.register()

    def register(self) -> AgentInfo:
        """Register/update profile on relay."""
        info = self._client.update_agent(
            self._keypair.public_key_b64,
            display_name=self._name,
            capabilities=self._capabilities,
            metadata=self._metadata,
            visibility=self._visibility,
            inbox_policy=self._inbox_policy,
        )
        self._registered = True
        return info

    # --- Discovery ---

    def discover(self, capability: str) -> list[DiscoveredAgent]:
        """Find agents by capability. Auto-paginates."""
        self._ensure_registered()
        return self._client.discover_all(capability)

    def discover_search(
        self,
        capabilities: list[str] | None = None,
        capability_prefix: str | None = None,
        name: str | None = None,
    ) -> list[DiscoveredAgent]:
        """Find agents by multiple criteria. Auto-paginates."""
        self._ensure_registered()
        return self._client.discover_search_all(
            capabilities=capabilities,
            capability_prefix=capability_prefix,
            name=name,
        )

    def list_capabilities(self) -> CapabilitiesResult:
        """List all capability tags with agent counts."""
        self._ensure_registered()
        return self._client.list_capabilities()

    # --- Messaging (sync, HTTP-based) ---

    def send(
        self,
        to: str,
        body: dict,
        content_type: str = "application/json",
        thread: str | None = None,
        in_reply_to: str | None = None,
        expires_in: int | None = None,
        attachments: list[str] | None = None,
        encrypt: bool = False,
    ) -> SendResult:
        """Send a message. If encrypt=True, E2E encrypt the body."""
        self._ensure_registered()

        actual_body: dict | None = body
        actual_ct = content_type

        if encrypt:
            recipient_pub = _b64url_decode(to)
            encrypted = encrypt_message(
                self._keypair.seed,
                recipient_pub,
                json.dumps(body).encode(),
                content_type,
            )
            actual_body = encrypted
            actual_ct = "application/x-m2m-encrypted"

        attach_list = None
        if attachments:
            attach_list = [Attachment(blob_id=bid) for bid in attachments]

        return self._client.send_message(
            recipient_key=to,
            body=actual_body,
            content_type=actual_ct,
            thread_id=thread,
            in_reply_to=in_reply_to,
            expires_in=expires_in,
            attachments=attach_list,
        )

    def reply(
        self,
        message: Message,
        body: dict,
        encrypt: bool = False,
    ) -> SendResult:
        """Reply to a message. Auto-sets in_reply_to and recipient."""
        return self.send(
            to=message.sender_key,
            body=body,
            in_reply_to=message.message_id,
            thread=message.thread_id,
            encrypt=encrypt,
        )

    # --- Structured Messaging ---

    def request(
        self,
        to: str,
        action: str,
        params: dict[str, Any],
        thread: str | None = None,
        response_format: str | None = None,
        encrypt: bool = False,
    ) -> SendResult:
        """Send a structured MRP request message."""
        body = create_request(action, params, response_format)
        return self.send(
            to=to,
            body=body.to_dict(),
            content_type=REQUEST_CONTENT_TYPE,
            thread=thread,
            encrypt=encrypt,
        )

    def respond(
        self,
        message: Message,
        result: dict[str, Any],
    ) -> SendResult:
        """Send an ok response to an MRP request message."""
        body = create_ok_response(result)
        return self.send(
            to=message.sender_key,
            body=body.to_dict(),
            content_type=RESPONSE_CONTENT_TYPE,
            in_reply_to=message.message_id,
            thread=message.thread_id,
        )

    def respond_error(
        self,
        message: Message,
        code: str,
        msg: str,
    ) -> SendResult:
        """Send an error response to an MRP request message."""
        body = create_error_response(code, msg)
        return self.send(
            to=message.sender_key,
            body=body.to_dict(),
            content_type=RESPONSE_CONTENT_TYPE,
            in_reply_to=message.message_id,
            thread=message.thread_id,
        )

    def emit_event(
        self,
        to: str,
        event: str,
        data: dict[str, Any] | None = None,
        thread: str | None = None,
        encrypt: bool = False,
    ) -> SendResult:
        """Send an MRP event message."""
        body = create_event(event, data)
        return self.send(
            to=to,
            body=body.to_dict(),
            content_type=EVENT_CONTENT_TYPE,
            thread=thread,
            encrypt=encrypt,
        )

    def send_status(
        self,
        to: str,
        progress: float | None = None,
        stage: str | None = None,
        detail: str | None = None,
        thread: str | None = None,
        encrypt: bool = False,
    ) -> SendResult:
        """Send an MRP status update message."""
        body = create_status(progress=progress, stage=stage, detail=detail)
        return self.send(
            to=to,
            body=body.to_dict(),
            content_type=STATUS_CONTENT_TYPE,
            thread=thread,
            encrypt=encrypt,
        )

    # --- Action Router ---

    def on_request(self, action: str, handler: Callable[[MRPRequest, Message], Any]) -> None:
        """Register a handler for a specific request action.

        The handler receives (MRPRequest, Message) and should return a dict.
        If the handler raises RequestError, an error response is auto-sent.
        """
        self._action_handlers[action] = handler

    def _dispatch_message(self, msg: Message) -> dict | None:
        """Dispatch a message through action routing then generic handler.

        Returns a reply dict if the generic handler produces one, else None.
        Action-routed handlers auto-send their own responses.
        """
        # Try action routing first for request messages
        if msg.content_type == REQUEST_CONTENT_TYPE and self._action_handlers:
            req = parse_request(msg.body)
            if req:
                handler = self._action_handlers.get(req.action)
                if handler:
                    try:
                        if inspect.iscoroutinefunction(handler):
                            result = asyncio.get_event_loop().run_until_complete(handler(req, msg))
                        else:
                            result = handler(req, msg)
                        self.respond(msg, result)
                    except RequestError as e:
                        self.respond_error(msg, e.code, e.message)
                    return None

        # Fall back to generic message handler
        if self._on_message_handler:
            return self._on_message_handler(msg)
        return None

    async def _dispatch_message_async(self, msg: Message) -> dict | None:
        """Async version of _dispatch_message for WebSocket mode."""
        # Try action routing first for request messages
        if msg.content_type == REQUEST_CONTENT_TYPE and self._action_handlers:
            req = parse_request(msg.body)
            if req:
                handler = self._action_handlers.get(req.action)
                if handler:
                    try:
                        if inspect.iscoroutinefunction(handler):
                            result = await handler(req, msg)
                        else:
                            result = handler(req, msg)
                        self.respond(msg, result)
                    except RequestError as e:
                        self.respond_error(msg, e.code, e.message)
                    return None

        # Fall back to generic message handler
        if self._on_message_handler:
            return self._on_message_handler(msg)
        return None

    def receive(
        self,
        limit: int = 50,
        sender: str | None = None,
    ) -> list[Message]:
        """One-shot poll for new messages."""
        self._ensure_registered()
        result = self._client.poll_messages(limit=limit, sender_key=sender)
        return result.messages

    def messages(
        self,
        poll_interval: float = 2.0,
        max_empty_interval: float = 30.0,
    ) -> Iterator[Message]:
        """Blocking iterator that polls for messages with smart backoff."""
        self._ensure_registered()
        interval = poll_interval

        while True:
            msgs = self.receive()
            if msgs:
                interval = poll_interval
                yield from msgs
            else:
                interval = min(interval * 2, max_empty_interval)
            time.sleep(interval)

    # --- Messaging (async, WebSocket-based) ---

    def on_message(self, handler: Callable[[Message], dict | None]) -> None:
        """Register a message handler for WebSocket mode."""
        self._on_message_handler = handler

    def run(self, *, mode: str = "websocket") -> None:
        """Block and process messages.

        mode="websocket": maintain WebSocket, dispatch to on_message handler.
        mode="poll": poll loop, dispatch to on_message handler.
        """
        self._ensure_registered()

        if mode == "poll":
            self._run_poll()
        elif mode == "websocket":
            asyncio.run(self._run_ws())
        else:
            raise ValueError(f"Unknown mode: {mode}")

    def _run_poll(self) -> None:
        """Poll loop that dispatches messages through action routing and on_message handler."""
        for msg in self.messages():
            reply = self._dispatch_message(msg)
            if reply is not None:
                self.reply(msg, reply)

    async def _run_ws(self) -> None:
        """WebSocket loop that dispatches messages through action routing and on_message handler."""
        ws = ReconnectingWSClient(self._relay, self._keypair)

        async def _dispatch(msg: Message) -> dict | None:
            return await self._dispatch_message_async(msg)

        ws.on_message(_dispatch)
        try:
            await ws.run()
        finally:
            await ws.close()

    # --- Blobs ---

    def upload(self, data: bytes | BinaryIO, content_type: str = "application/octet-stream") -> Blob:
        """Upload a blob."""
        self._ensure_registered()
        return self._client.upload_blob(data, content_type)

    def download(self, blob_id: str) -> tuple[bytes, str]:
        """Download a blob."""
        self._ensure_registered()
        return self._client.download_blob(blob_id)

    # --- Thread ---

    def thread(self, thread_id: str) -> list[Message]:
        """Get all messages in a thread (auto-paginates)."""
        self._ensure_registered()
        all_msgs: list[Message] = []
        cursor = None
        while True:
            result = self._client.get_thread(thread_id, cursor=cursor)
            all_msgs.extend(result.messages)
            if not result.next_cursor:
                break
            cursor = result.next_cursor
        return all_msgs

    # --- ACL ---

    def allow(self, peer_key: str) -> ACLEntry:
        """Allow a peer to message this agent."""
        self._ensure_registered()
        return self._client.set_acl(self._keypair.public_key_b64, peer_key, "allow")

    def block(self, peer_key: str) -> ACLEntry:
        """Block a peer from messaging this agent."""
        self._ensure_registered()
        return self._client.set_acl(self._keypair.public_key_b64, peer_key, "block")

    def unblock(self, peer_key: str) -> None:
        """Remove a peer from the ACL."""
        self._ensure_registered()
        self._client.delete_acl(self._keypair.public_key_b64, peer_key)

    def list_acl(self, entry_type: str | None = None) -> list[ACLEntry]:
        """List ACL entries, optionally filtered by type."""
        self._ensure_registered()
        return self._client.list_acl(self._keypair.public_key_b64, entry_type)

    # --- Lifecycle ---

    def delete(self) -> None:
        """Delete this agent from the relay."""
        self._client.delete_agent(self._keypair.public_key_b64)

    def close(self) -> None:
        """Close connections."""
        self._client.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    def decrypt(self, message: Message) -> tuple[bytes, str]:
        """Decrypt an E2E encrypted message."""
        if message.content_type != "application/x-m2m-encrypted":
            raise ValueError("Message is not encrypted")
        sender_pub = _b64url_decode(message.sender_key)
        return decrypt_message(self._keypair.seed, sender_pub, message.body)
