from __future__ import annotations

from typing import Any

from pydantic import BaseModel, Field


class AgentConfig(BaseModel):
    """Per-agent-type configuration stored in Restaurant.settings.agents.<type>."""

    enabled: bool = True
    custom_prompt_suffix: str = Field(default="", max_length=2000)


class AgentSettings(BaseModel):
    """Top-level agents key in Restaurant.settings."""

    reservation: AgentConfig = Field(default_factory=AgentConfig)
    faq: AgentConfig = Field(default_factory=AgentConfig)
    # Takeaway is opt-in: disabled unless explicitly enabled in settings.
    takeaway: AgentConfig = Field(default_factory=lambda: AgentConfig(enabled=False))


def get_agent_config(settings: dict[str, Any] | None, agent_type: str) -> AgentConfig:
    """Extract and validate agent config from restaurant settings JSONB.

    Returns defaults when absent or invalid.
    """
    if not settings or "agents" not in settings:
        return AgentConfig()
    try:
        agent_settings = AgentSettings.model_validate(settings["agents"])
        return getattr(agent_settings, agent_type, AgentConfig())
    except Exception:
        return AgentConfig()


def get_enabled_agents(settings: dict[str, Any] | None) -> set[str]:
    """Return the set of agent types that are enabled for this restaurant.

    Checks ``settings.agents.<type>.enabled`` first (structured config).
    Falls back to flat flags like ``takeaway_enabled`` when the ``agents``
    sub-object is absent or doesn't cover a capability.
    """
    # Flat-flag overrides: map flat settings keys to agent types.
    # When the structured `agents` key is absent, these are the authority.
    _FLAT_FLAGS: dict[str, str] = {
        "takeaway_enabled": "takeaway",
    }

    if not settings:
        return {"reservation", "faq"}

    if "agents" in settings:
        try:
            agent_settings = AgentSettings.model_validate(settings["agents"])
            return {
                name
                for name in ("reservation", "faq", "takeaway")
                if getattr(agent_settings, name).enabled
            }
        except Exception:
            pass  # Fall through to flat-flag logic

    # No structured agents config — derive from flat flags.
    # faq and reservation are always enabled; takeaway is opt-in.
    enabled: set[str] = {"reservation", "faq"}
    for flag_key, agent_type in _FLAT_FLAGS.items():
        if settings.get(flag_key) is True:
            enabled.add(agent_type)
    return enabled
