from __future__ import annotations

import json
from datetime import datetime, timezone
from io import IOBase
from typing import BinaryIO

import httpx

from mrp.auth import Keypair
from mrp.errors import (
    AuthError,
    ConflictError,
    ForbiddenError,
    MRPError,
    NotFoundError,
    PayloadTooLargeError,
    RateLimitError,
    StorageError,
    ValidationError,
)
from mrp.models import (
    AgentInfo,
    Attachment,
    Blob,
    BlobMetadata,
    CapabilitiesResult,
    DiscoveredAgent,
    DiscoverResult,
    Message,
    MessageStatus,
    PollResult,
    SendResult,
)


_STATUS_MAP: dict[int, type[MRPError]] = {
    401: AuthError,
    403: ForbiddenError,
    404: NotFoundError,
    409: ConflictError,
    413: PayloadTooLargeError,
    422: ValidationError,
    429: RateLimitError,
    507: StorageError,
}


class Client:
    """Low-level MRP HTTP client. Stateless, handles signing."""

    def __init__(self, relay: str, keypair: Keypair):
        self.relay = relay.rstrip("/")
        self.keypair = keypair
        self._session = httpx.Client(timeout=30.0)

    def close(self) -> None:
        self._session.close()

    def __enter__(self):
        return self

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

    def _request(
        self,
        method: str,
        path: str,
        body: dict | bytes | None = None,
        params: dict | list[tuple[str, str]] | None = None,
        headers: dict | None = None,
        raw_body: bool = False,
    ) -> httpx.Response:
        """Make an authenticated request."""
        # Serialize body
        if body is None:
            body_bytes = b""
        elif raw_body or isinstance(body, bytes):
            body_bytes = body
        else:
            body_bytes = json.dumps(body).encode()

        # Generate timestamp
        timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")

        # Sign request
        signature = self.keypair.sign_request(method, path, body_bytes, timestamp)

        # Build headers
        req_headers = {
            "X-M2M-Public-Key": self.keypair.public_key_b64,
            "X-M2M-Timestamp": timestamp,
            "X-M2M-Signature": signature,
        }
        if not raw_body and body is not None and not isinstance(body, bytes):
            req_headers["Content-Type"] = "application/json"
        if headers:
            req_headers.update(headers)

        # Send request
        url = self.relay + path
        resp = self._session.request(
            method,
            url,
            content=body_bytes if body_bytes else None,
            params=params,
            headers=req_headers,
        )

        # Check for errors
        if resp.status_code >= 400:
            self._raise_error(resp)

        return resp

    def _raise_error(self, resp: httpx.Response) -> None:
        """Parse error response and raise appropriate exception."""
        code = "unknown"
        message = f"HTTP {resp.status_code}"
        request_id = None

        try:
            data = resp.json()
            err = data.get("error", {})
            code = err.get("code", code)
            message = err.get("message", message)
            request_id = err.get("request_id")
        except Exception:
            pass

        status = resp.status_code

        if status == 400:
            raise ValidationError(code, message, request_id)

        exc_cls = _STATUS_MAP.get(status, MRPError)
        if exc_cls is RateLimitError:
            raise RateLimitError(
                code,
                message,
                request_id,
                retry_after=float(resp.headers.get("Retry-After", 0)),
                limit=int(resp.headers.get("X-RateLimit-Limit", 0)),
                remaining=int(resp.headers.get("X-RateLimit-Remaining", 0)),
                reset=float(resp.headers.get("X-RateLimit-Reset", 0)),
            )
        raise exc_cls(code, message, request_id)

    # --- Agent Lifecycle ---

    def get_agent(self, public_key: str) -> AgentInfo:
        """GET /v1/agents/{publicKey}"""
        resp = self._request("GET", f"/v1/agents/{public_key}")
        return AgentInfo.from_dict(resp.json())

    def update_agent(
        self,
        public_key: str,
        display_name: str | None = None,
        capabilities: list[str] | None = None,
        metadata: dict[str, str] | None = None,
    ) -> AgentInfo:
        """PATCH /v1/agents/{publicKey}"""
        body: dict = {}
        if display_name is not None:
            body["display_name"] = display_name
        if capabilities is not None:
            body["capabilities"] = capabilities
        if metadata is not None:
            body["metadata"] = metadata
        resp = self._request("PATCH", f"/v1/agents/{public_key}", body=body)
        return AgentInfo.from_dict(resp.json())

    def delete_agent(self, public_key: str) -> None:
        """DELETE /v1/agents/{publicKey} -> 204"""
        self._request("DELETE", f"/v1/agents/{public_key}")

    # --- Messaging ---

    def send_message(
        self,
        recipient_key: str,
        body: dict | None = None,
        content_type: str = "application/json",
        thread_id: str | None = None,
        in_reply_to: str | None = None,
        expires_in: int | None = None,
        attachments: list[Attachment] | None = None,
        metadata: dict[str, str] | None = None,
    ) -> SendResult:
        """POST /v1/messages -> 202 Accepted"""
        payload: dict = {"recipient_key": recipient_key}
        if body is not None:
            payload["body"] = body
        if content_type != "application/json":
            payload["content_type"] = content_type
        if thread_id is not None:
            payload["thread_id"] = thread_id
        if in_reply_to is not None:
            payload["in_reply_to"] = in_reply_to
        if expires_in is not None:
            payload["expires_in"] = expires_in
        if attachments:
            payload["attachments"] = [a.to_dict() for a in attachments]
        if metadata:
            payload["metadata"] = metadata
        resp = self._request("POST", "/v1/messages", body=payload)
        return SendResult.from_dict(resp.json())

    def poll_messages(
        self,
        status: str = "sent",
        limit: int = 50,
        cursor: str | None = None,
        sender_key: str | None = None,
        in_reply_to: str | None = None,
    ) -> PollResult:
        """GET /v1/messages"""
        params: dict = {"status": status, "limit": str(limit)}
        if cursor:
            params["cursor"] = cursor
        if sender_key:
            params["sender_key"] = sender_key
        if in_reply_to:
            params["in_reply_to"] = in_reply_to
        resp = self._request("GET", "/v1/messages", params=params)
        return PollResult.from_dict(resp.json())

    def get_message(self, message_id: str) -> Message:
        """GET /v1/messages/{messageID}"""
        resp = self._request("GET", f"/v1/messages/{message_id}")
        return Message.from_dict(resp.json())

    def get_message_status(self, message_id: str) -> MessageStatus:
        """GET /v1/messages/{messageID}/status"""
        resp = self._request("GET", f"/v1/messages/{message_id}/status")
        return MessageStatus.from_dict(resp.json())

    def get_thread(
        self,
        thread_id: str,
        limit: int = 50,
        cursor: str | None = None,
    ) -> PollResult:
        """GET /v1/messages/threads/{threadID}"""
        params: dict = {"limit": str(limit)}
        if cursor:
            params["cursor"] = cursor
        resp = self._request("GET", f"/v1/messages/threads/{thread_id}", params=params)
        return PollResult.from_dict(resp.json())

    # --- Discovery ---

    def discover(
        self,
        capability: str,
        limit: int = 20,
        cursor: str | None = None,
    ) -> DiscoverResult:
        """GET /v1/discover?capability=X&limit=N&cursor=C"""
        return self.discover_search(
            capabilities=[capability], limit=limit, cursor=cursor
        )

    def discover_search(
        self,
        capabilities: list[str] | None = None,
        capability_prefix: str | None = None,
        name: str | None = None,
        limit: int = 20,
        cursor: str | None = None,
    ) -> DiscoverResult:
        """Advanced discovery search with multiple filters."""
        params: list[tuple[str, str]] = []
        if capabilities:
            for cap in capabilities:
                params.append(("capability", cap))
        if capability_prefix:
            params.append(("capability_prefix", capability_prefix))
        if name:
            params.append(("name", name))
        params.append(("limit", str(limit)))
        if cursor:
            params.append(("cursor", cursor))
        resp = self._request("GET", "/v1/discover", params=params)
        return DiscoverResult.from_dict(resp.json())

    def discover_all(self, capability: str) -> list[DiscoveredAgent]:
        """Paginate through all results for a single capability."""
        return self.discover_search_all(capabilities=[capability])

    def discover_search_all(
        self,
        capabilities: list[str] | None = None,
        capability_prefix: str | None = None,
        name: str | None = None,
    ) -> list[DiscoveredAgent]:
        """Paginate through all results matching search criteria."""
        results: list[DiscoveredAgent] = []
        cursor = None
        while True:
            page = self.discover_search(
                capabilities=capabilities,
                capability_prefix=capability_prefix,
                name=name,
                limit=100,
                cursor=cursor,
            )
            results.extend(page.agents)
            if not page.next_cursor:
                break
            cursor = page.next_cursor
        return results

    def list_capabilities(self) -> CapabilitiesResult:
        """GET /v1/capabilities — list all capability tags with counts."""
        resp = self._request("GET", "/v1/capabilities")
        return CapabilitiesResult.from_dict(resp.json())

    # --- Blobs ---

    def upload_blob(
        self,
        data: bytes | BinaryIO,
        content_type: str = "application/octet-stream",
    ) -> Blob:
        """POST /v1/blobs"""
        if isinstance(data, (IOBase,)):
            data = data.read()
        resp = self._request(
            "POST",
            "/v1/blobs",
            body=data,
            raw_body=True,
            headers={"Content-Type": content_type},
        )
        return Blob.from_dict(resp.json())

    def download_blob(self, blob_id: str) -> tuple[bytes, str]:
        """GET /v1/blobs/{blobID} -> (data, content_type)"""
        resp = self._request("GET", f"/v1/blobs/{blob_id}")
        ct = resp.headers.get("content-type", "application/octet-stream")
        return resp.content, ct

    def get_blob_metadata(self, blob_id: str) -> BlobMetadata:
        """HEAD /v1/blobs/{blobID}"""
        resp = self._request("HEAD", f"/v1/blobs/{blob_id}")
        return BlobMetadata(
            content_type=resp.headers.get("content-type", "application/octet-stream"),
            size=int(resp.headers.get("content-length", 0)),
            hash=resp.headers.get("x-m2m-blob-hash", ""),
        )

    def delete_blob(self, blob_id: str) -> None:
        """DELETE /v1/blobs/{blobID} -> 204"""
        self._request("DELETE", f"/v1/blobs/{blob_id}")
