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.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 resolve_recipient(self, identifier: str) -> str:
        """Resolve a name, alias, or public key to a public key."""
        # 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_discover to find agents first."
        )
