from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime

from mrp.models import DiscoveredAgent, Message


@dataclass
class Contact:
    public_key: str
    display_name: str | None
    capabilities: list[str]
    alias: str | None = None
    last_seen: datetime | None = None


class SessionState:
    """In-memory state for the MCP session."""

    def __init__(self):
        self.contacts: dict[str, Contact] = {}
        self.aliases: dict[str, str] = {}
        self._sender_cache: dict[str, Contact] = {}
        self.threads: dict[str, str] = {}
        self.unread: list[Message] = []
        self.last_poll_cursor: str | None = None

    def add_contact(self, agent: DiscoveredAgent, alias: str | None = None) -> Contact:
        """Add or update a contact. Auto-index by display_name."""
        contact = Contact(
            public_key=agent.public_key,
            display_name=agent.display_name,
            capabilities=agent.capabilities,
            alias=alias,
            last_seen=agent.last_active_at,
        )
        self.contacts[agent.public_key] = contact
        if agent.display_name:
            self.aliases[agent.display_name.lower()] = agent.public_key
        if alias:
            self.aliases[alias.lower()] = agent.public_key
        return contact

    def cache_sender(
        self,
        public_key: str,
        display_name: str | None,
        capabilities: list[str] | None = None,
        last_seen: datetime | None = None,
    ) -> Contact:
        """Cache a sender's info for display purposes only.

        This does NOT add the sender to contacts. It only stores enough
        info to show a friendly name in message formatting.
        """
        contact = Contact(
            public_key=public_key,
            display_name=display_name,
            capabilities=capabilities or [],
            last_seen=last_seen,
        )
        self._sender_cache[public_key] = contact
        return contact

    def resolve_display_name(self, public_key: str) -> str:
        """Get a display name for a public key, checking contacts then sender cache."""
        contact = self.contacts.get(public_key)
        if contact and contact.display_name:
            return contact.display_name
        cached = self._sender_cache.get(public_key)
        if cached and cached.display_name:
            return cached.display_name
        return public_key[:12] + "..."

    def resolve_recipient(self, identifier: str) -> str:
        """Resolve a name, alias, or public key to a public key.

        Only checks contacts (not sender cache). You must explicitly
        add an agent to contacts before you can send to them.
        """
        # Direct public key lookup
        if identifier in self.contacts:
            return identifier
        # Alias or display name lookup (case-insensitive)
        lower = identifier.lower()
        if lower in self.aliases:
            return self.aliases[lower]
        # Fuzzy: check if identifier is a substring of any display name
        for pk, contact in self.contacts.items():
            if contact.display_name and lower in contact.display_name.lower():
                return pk
        raise ValueError(
            f"Unknown recipient: {identifier}. "
            "Use mrp_add_contact to add them to your contacts first."
        )
