from __future__ import annotations

import asyncio
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 (
    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,
    ):
        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._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,
        )
        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 asyncio.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 asyncio.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

    # --- 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)
