"""MCP server for controlling local Javis Slave devices."""
from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass, field
from typing import Any, Awaitable, Callable, Optional
from urllib.parse import quote

import httpx
from mcp.server.fastmcp import FastMCP
from mcp.server.fastmcp.exceptions import ToolError
from mcp.types import ToolAnnotations
from pydantic import BaseModel

log = logging.getLogger("javis_slave_mcp")

DEFAULT_PORT = 8766
# Must match javis_slave.http_server.CONTROL_HEADER: loopback callers send it
# instead of the Bearer token (browsers can't forge it cross-origin).
CONTROL_HEADER = "X-Javis-Control"
_TIMEOUT = httpx.Timeout(10.0, connect=5.0)

# ---- Models ----


class SlaveStatus(BaseModel):
    """Current status of a slave device."""
    device_id: str
    device_name: str
    room: Optional[str] = None
    paired: bool
    engine_running: bool
    audio_state: str
    muted: bool = False
    mic_level: float = 0.0
    wake_word_score: float = 0.0
    last_error: Optional[str] = None


class DeviceInfo(BaseModel):
    """Static device information."""
    device_id: str
    device_name: str
    sample_rate: int
    input_device: Optional[int] = None
    output_device: Optional[int] = None
    wake_word_model: str


class ControlResult(BaseModel):
    """Result of a control operation."""
    success: bool
    message: str = ""
    device_id: str = ""


class SlaveDeviceInfo(BaseModel):
    """Information about a connected slave device."""
    device_id: str
    name: str
    room: Optional[str] = None
    ip_address: Optional[str] = None
    last_seen: Optional[str] = None
    last_seen_ip: Optional[str] = None
    created_at: Optional[str] = None


# ---- Client ----


class SlaveAuthError(Exception):
    """The slave rejected the request (missing or wrong control token)."""


@dataclass
class SlaveClient:
    """HTTP client for a Javis Slave device's control API."""
    device_id: str
    host: str
    port: int = DEFAULT_PORT
    token: str = field(default="", repr=False)
    _http: Optional[httpx.AsyncClient] = field(default=None, init=False, repr=False, compare=False)

    def base_url(self) -> str:
        return f"http://{self.host}:{self.port}"

    def headers(self) -> dict[str, str]:
        h = {CONTROL_HEADER: "1"}
        if self.token:
            h["Authorization"] = f"Bearer {self.token}"
        return h

    def _client(self) -> httpx.AsyncClient:
        # One pooled client per device instead of a new connection per call.
        # trust_env=False: never send LAN control calls (and the token) via a proxy.
        if self._http is None or self._http.is_closed:
            self._http = httpx.AsyncClient(base_url=self.base_url(), headers=self.headers(),
                                           timeout=_TIMEOUT, trust_env=False)
        return self._http

    async def aclose(self) -> None:
        if self._http is not None:
            await self._http.aclose()
            self._http = None

    async def _call(self, method: str, path: str, **params: Any) -> Any:
        r = await self._client().request(method, path, params=params or None)
        if r.status_code == 401:
            raise SlaveAuthError(
                f"{self.base_url()} rejected the request (401). Set the slave's control token "
                "(JAVIS_SLAVE_TOKEN or the device's \"token\"; print it on the slave with "
                "`javis-slave --show-control-token`).")
        r.raise_for_status()
        return r.json()

    async def _control(self, path: str) -> ControlResult:
        data = await self._call("POST", path)
        data["device_id"] = self.device_id
        return ControlResult(**data)

    async def get_status(self) -> SlaveStatus:
        """Fetch current device status."""
        return SlaveStatus(**await self._call("GET", "/mcp/v1/status"))

    async def get_device_info(self) -> DeviceInfo:
        """Fetch device information."""
        return DeviceInfo(**await self._call("GET", "/mcp/v1/device-info"))

    async def mute_mic(self) -> ControlResult:
        return await self._control("/mcp/v1/mic/mute")

    async def unmute_mic(self) -> ControlResult:
        return await self._control("/mcp/v1/mic/unmute")

    async def toggle_mic(self) -> ControlResult:
        return await self._control("/mcp/v1/mic/toggle")

    async def restart_audio(self) -> ControlResult:
        return await self._control("/mcp/v1/restart")

    # ---- Volume Control Methods ----

    async def get_volume(self) -> dict:
        return await self._call("GET", "/mcp/v1/volume")

    async def set_volume(self, level: float) -> dict:
        return await self._call("POST", "/mcp/v1/volume/set", level=level)

    async def increase_volume(self, amount: float = 0.05) -> dict:
        return await self._call("POST", "/mcp/v1/volume/increase", amount=amount)

    async def decrease_volume(self, amount: float = 0.05) -> dict:
        return await self._call("POST", "/mcp/v1/volume/decrease", amount=amount)

    async def mute_output(self) -> dict:
        return await self._call("POST", "/mcp/v1/volume/mute")

    async def unmute_output(self) -> dict:
        return await self._call("POST", "/mcp/v1/volume/unmute")

    # ---- Brightness Control Methods ----

    async def get_brightness(self) -> dict:
        return await self._call("GET", "/mcp/v1/brightness")

    async def set_brightness(self, level: int) -> dict:
        return await self._call("POST", "/mcp/v1/brightness/set", level=level)

    async def increase_brightness(self, amount: int = 5) -> dict:
        return await self._call("POST", "/mcp/v1/brightness/increase", amount=amount)

    async def decrease_brightness(self, amount: int = 5) -> dict:
        return await self._call("POST", "/mcp/v1/brightness/decrease", amount=amount)

    # ---- Power Control Methods ----

    async def sleep(self) -> dict:
        return await self._call("POST", "/mcp/v1/power/sleep")

    async def shutdown(self, delay: int = 0) -> dict:
        return await self._call("POST", "/mcp/v1/power/shutdown", delay=delay)

    async def restart(self, delay: int = 0) -> dict:
        return await self._call("POST", "/mcp/v1/power/restart", delay=delay)

    async def cancel_shutdown(self) -> dict:
        return await self._call("POST", "/mcp/v1/power/cancel")

    async def get_connected_slaves(self) -> list[SlaveDeviceInfo]:
        """Fetch list of all connected slaves from the master.

        Uses the slave's authenticated pass-through endpoint
        (``GET /mcp/v1/master/slaves`` on the slave's own port). The slave
        calls the master with its own ``X-Device-Id`` / Bearer credentials,
        so this works whether or not the MCP host and the slave host match.
        """
        return [SlaveDeviceInfo(**item) for item in await self._call("GET", "/mcp/v1/master/slaves")]

    async def get_slave_ip(self, target_device_id: str) -> dict:
        """Fetch IP info for a specific slave from the master via the pass-through.

        The slave's own credentials authenticate the request to the master, so
        this works whether or not the MCP host and the slave host match.

        Args:
            target_device_id: The device_id of the slave whose IP to look up
                (may be this same slave or a different one).

        Returns:
            The ``MasterSlaveIp`` dict (``device_id``, ``ip_address``,
            ``last_seen_ip``, ``last_seen``) as reported by the master.
        """
        return await self._call("GET", f"/mcp/v1/master/slaves/{quote(target_device_id, safe='')}/ip")


# ---- MCP Server ----


def _parse_env_devices(raw: str, default_token: str) -> dict[str, SlaveClient]:
    """Parse ``id=host:port`` or ``id=TOKEN@host:port`` entries (comma-separated).

    A malformed entry is skipped with a warning instead of discarding the rest.
    """
    slaves: dict[str, SlaveClient] = {}
    for n, entry in enumerate(raw.split(","), start=1):
        entry = entry.strip()
        if not entry:
            continue
        device_id, sep, target = entry.partition("=")
        device_id = device_id.strip()
        token, _, hostport = target.strip().rpartition("@")
        host, colon, port_s = hostport.rpartition(":")
        if not colon or "]" in port_s:  # no port given (or bare [IPv6])
            host, port_s = hostport, str(DEFAULT_PORT)
        if not sep or not device_id or not host or not port_s.isdigit():
            # Don't echo the entry: it may contain a token.
            log.warning(f"JAVIS_SLAVE_DEVICES: ignoring malformed entry #{n} "
                        "(expected id=host:port or id=TOKEN@host:port)")
            continue
        slaves[device_id] = SlaveClient(device_id=device_id, host=host, port=int(port_s),
                                        token=token or default_token)
    return slaves


def load_slave_config() -> dict[str, SlaveClient]:
    """Load configured slave devices from environment and/or config file.

    - ``JAVIS_SLAVE_DEVICES=id=host:port,id2=TOKEN@host:port``
    - ``JAVIS_SLAVE_MCP_CONFIG`` (default ``~/.javis/slave_mcp_config.json``):
      ``{"slaves": {"id": {"host": "…", "port": 8766, "token": "…"}}}``
    - ``JAVIS_SLAVE_TOKEN``: default control token for entries without one.
    - Fallback: ``localhost`` → 127.0.0.1:8766.
    """
    default_token = os.environ.get("JAVIS_SLAVE_TOKEN", "").strip()
    slaves = _parse_env_devices(os.environ.get("JAVIS_SLAVE_DEVICES", ""), default_token)

    config_file = os.environ.get(
        "JAVIS_SLAVE_MCP_CONFIG",
        os.path.expanduser("~/.javis/slave_mcp_config.json"),
    )
    if os.path.exists(config_file):
        try:
            with open(config_file, encoding="utf-8") as f:
                config = json.load(f)
            entries = config.get("slaves") or {}
            if not isinstance(entries, dict):
                raise ValueError("'slaves' must be an object")
        except (OSError, ValueError, AttributeError) as e:
            log.warning(f"Failed to load slave config from {config_file}: {e}")
            entries = {}
        for device_id, dev in entries.items():
            try:
                slaves[device_id] = SlaveClient(
                    device_id=device_id,
                    host=str(dev.get("host", "127.0.0.1")),
                    port=int(dev.get("port", DEFAULT_PORT)),
                    token=str(dev.get("token") or default_token),
                )
            except (AttributeError, TypeError, ValueError) as e:
                log.warning(f"{config_file}: ignoring slave {device_id!r}: {e}")

    if not slaves:
        log.info(f"No slaves configured; using localhost:{DEFAULT_PORT} as default")
        slaves["localhost"] = SlaveClient(
            device_id="localhost", host="127.0.0.1", port=DEFAULT_PORT, token=default_token)

    return slaves


# Global slave registry
_slaves: dict[str, SlaveClient] = {}


mcp = FastMCP(
    name="javis-slave",
    instructions=(
        "Use these tools to query and control local Javis Slave devices. "
        "Supported operations: status, device info, IP lookup, "
        "mute/unmute/toggle microphone, restart audio engine, "
        "volume get/set/increase/decrease, output mute/unmute, "
        "brightness get/set/increase/decrease, "
        "power sleep/shutdown/restart, and cancel shutdown. "
        "Master-derived tools (slave_get_ip_address, "
        "slave_get_connected_devices) route through the slave's own "
        "credentials, so they work on LAN deployments where the MCP "
        "host and the slave host differ."
    ),
)


# ---- Tools ----

# MCP tool annotations (hints for clients deciding what needs confirmation).
_READ = ToolAnnotations(readOnlyHint=True, openWorldHint=False)
_SET = ToolAnnotations(readOnlyHint=False, destructiveHint=False, idempotentHint=True,
                       openWorldHint=False)
_STEP = ToolAnnotations(readOnlyHint=False, destructiveHint=False, idempotentHint=False,
                        openWorldHint=False)
_POWER = ToolAnnotations(readOnlyHint=False, destructiveHint=True, idempotentHint=False,
                         openWorldHint=False)


def _render(result: Any) -> str:
    if isinstance(result, str):
        return result
    if isinstance(result, BaseModel):
        return result.model_dump_json(indent=2)
    if isinstance(result, list):
        result = [r.model_dump() if isinstance(r, BaseModel) else r for r in result]
    return json.dumps(result, indent=2)


async def _run(device_id: str, op: Callable[[SlaveClient], Awaitable[Any]]) -> str:
    """Run *op* against a configured slave on the server's own event loop.

    Failures are raised as :class:`ToolError`, which FastMCP reports as a
    tool result with ``isError: true`` (visible to the model).
    """
    client = _slaves.get(device_id)
    if client is None:
        raise ToolError(f"device '{device_id}' not found; configured: {', '.join(_slaves) or 'none'}")
    try:
        return _render(await op(client))
    except SlaveAuthError as e:
        msg = str(e)
    except httpx.HTTPStatusError as e:
        msg = f"{device_id}: HTTP {e.response.status_code}: {e.response.text[:200]}"
    except httpx.HTTPError as e:
        msg = f"{device_id} ({client.base_url()}) not reachable: {type(e).__name__}: {e}"
    except ValueError as e:  # bad JSON or unexpected payload shape (pydantic)
        msg = f"{device_id}: unexpected response from slave: {e}"
    log.warning(msg)
    raise ToolError(msg)


@mcp.tool(
    description="Get the current status of a Javis Slave device, including audio state and mic level.",
    annotations=_READ,
)
async def slave_get_status(device_id: str = "localhost") -> str:
    """Get slave device status.
    
    Args:
        device_id: The slave device ID (default: localhost).
    
    Returns:
        JSON status object with device info and current audio state.
    """
    return await _run(device_id, lambda c: c.get_status())


@mcp.tool(
    description="Get device information for a Javis Slave (static config like device name, sample rate, audio devices).",
    annotations=_READ,
)
async def slave_get_info(device_id: str = "localhost") -> str:
    """Get slave device info.
    
    Args:
        device_id: The slave device ID (default: localhost).
    
    Returns:
        JSON device info object.
    """
    return await _run(device_id, lambda c: c.get_device_info())


@mcp.tool(
    description="Get the IP address of a configured Javis Slave device, as last reported by the master.",
    annotations=_READ,
)
async def slave_get_ip_address(device_id: str = "localhost") -> str:
    """Get slave device IP address from master records.

    Routes through the slave's authenticated pass-through
    (``GET /mcp/v1/master/slaves/{id}/ip``), so the slave presents its own
    credentials to the master. This works even when the MCP host and the
    slave host differ.

    Args:
        device_id: The slave device ID to look up (default: localhost).

    Returns:
        IP address string, or error message if not found.
    """
    async def _op(c: SlaveClient) -> str:
        # ``device_id`` is this server's alias (e.g. "localhost"); the master
        # knows the slave by the UUID it reports in its own status.
        real_id = (await c.get_status()).device_id
        data = await c.get_slave_ip(real_id)
        return data.get("ip_address") or data.get("last_seen_ip") or "Unknown"

    return await _run(device_id, _op)


@mcp.tool(
    description="Mute the microphone on a Javis Slave device.",
    annotations=_SET,
)
async def slave_mute_mic(device_id: str = "localhost") -> str:
    """Mute slave microphone.
    
    Args:
        device_id: The slave device ID (default: localhost).
    
    Returns:
        JSON control result.
    """
    return await _run(device_id, lambda c: c.mute_mic())


@mcp.tool(
    description="Unmute the microphone on a Javis Slave device.",
    annotations=_SET,
)
async def slave_unmute_mic(device_id: str = "localhost") -> str:
    """Unmute slave microphone.
    
    Args:
        device_id: The slave device ID (default: localhost).
    
    Returns:
        JSON control result.
    """
    return await _run(device_id, lambda c: c.unmute_mic())


@mcp.tool(
    description="Toggle the microphone on/off on a Javis Slave device.",
    annotations=_STEP,
)
async def slave_toggle_mic(device_id: str = "localhost") -> str:
    """Toggle slave microphone.
    
    Args:
        device_id: The slave device ID (default: localhost).
    
    Returns:
        JSON control result with new mic state.
    """
    return await _run(device_id, lambda c: c.toggle_mic())


@mcp.tool(
    description="Restart the audio engine on a Javis Slave device (useful for recovery from errors).",
    annotations=_SET,
)
async def slave_restart_audio(device_id: str = "localhost") -> str:
    """Restart slave audio engine.
    
    Args:
        device_id: The slave device ID (default: localhost).
    
    Returns:
        JSON control result.
    """
    return await _run(device_id, lambda c: c.restart_audio())


@mcp.tool(
    description="List all configured Javis Slave devices.",
    annotations=_READ,
)
def slave_list_devices() -> str:
    """List all configured slave devices.
    
    Returns:
        JSON array of device IDs.
    """
    devices = list(_slaves.keys())
    return f"Configured slave devices: {devices}"


@mcp.tool(
    description="Get list of all connected slaves with their IP addresses from the master.",
    annotations=_READ,
)
async def slave_get_connected_devices(device_id: str = "localhost") -> str:
    """Get list of connected slaves with IPs from master.

    Routes through the configured slave's authenticated pass-through
    (``GET /mcp/v1/master/slaves``). The slave presents its own credentials
    to the master, so this works on LAN deployments where the MCP host and
    the slave host differ.

    Args:
        device_id: Configured slave to route the master query through
            (any paired device works; default: localhost).

    Returns:
        JSON array of slave device info objects with IP addresses.
    """
    return await _run(device_id, lambda c: c.get_connected_slaves())


# ---- Volume Control Tools ----


@mcp.tool(
    description="Get the current master volume level of a slave device (0.0-1.0).",
    annotations=_READ,
)
async def slave_get_volume(device_id: str = "localhost") -> str:
    """Get slave device volume level.
    
    Args:
        device_id: The slave device ID (default: localhost).
    
    Returns:
        Volume level as JSON response.
    """
    return await _run(device_id, lambda c: c.get_volume())


@mcp.tool(
    description="Set the master volume level on a slave device (0.0-1.0).",
    annotations=_SET,
)
async def slave_set_volume(device_id: str = "localhost", level: float = 0.5) -> str:
    """Set slave device volume level.
    
    Args:
        device_id: The slave device ID (default: localhost).
        level: Volume level 0.0-1.0 (default 0.5).
    
    Returns:
        Volume control result as JSON.
    """
    return await _run(device_id, lambda c: c.set_volume(level))


@mcp.tool(
    description="Increase the master volume on a slave device.",
    annotations=_STEP,
)
async def slave_increase_volume(device_id: str = "localhost", amount: float = 0.05) -> str:
    """Increase slave device volume.
    
    Args:
        device_id: The slave device ID (default: localhost).
        amount: Amount to increase 0.0-1.0 (default 0.05 = 5%).
    
    Returns:
        Volume control result as JSON.
    """
    return await _run(device_id, lambda c: c.increase_volume(amount))


@mcp.tool(
    description="Decrease the master volume on a slave device.",
    annotations=_STEP,
)
async def slave_decrease_volume(device_id: str = "localhost", amount: float = 0.05) -> str:
    """Decrease slave device volume.
    
    Args:
        device_id: The slave device ID (default: localhost).
        amount: Amount to decrease 0.0-1.0 (default 0.05 = 5%).
    
    Returns:
        Volume control result as JSON.
    """
    return await _run(device_id, lambda c: c.decrease_volume(amount))


@mcp.tool(
    description="Mute the master audio output on a slave device.",
    annotations=_SET,
)
async def slave_mute_output(device_id: str = "localhost") -> str:
    """Mute slave device audio output.
    
    Args:
        device_id: The slave device ID (default: localhost).
    
    Returns:
        Volume control result as JSON.
    """
    return await _run(device_id, lambda c: c.mute_output())


@mcp.tool(
    description="Unmute the master audio output on a slave device.",
    annotations=_SET,
)
async def slave_unmute_output(device_id: str = "localhost") -> str:
    """Unmute slave device audio output.
    
    Args:
        device_id: The slave device ID (default: localhost).
    
    Returns:
        Volume control result as JSON.
    """
    return await _run(device_id, lambda c: c.unmute_output())


# ---- Brightness Control Tools ----


@mcp.tool(
    description="Get the current display brightness on a slave device (0-100).",
    annotations=_READ,
)
async def slave_get_brightness(device_id: str = "localhost") -> str:
    """Get slave device brightness level.
    
    Args:
        device_id: The slave device ID (default: localhost).
    
    Returns:
        Brightness level as JSON response.
    """
    return await _run(device_id, lambda c: c.get_brightness())


@mcp.tool(
    description="Set the display brightness on a slave device (0-100).",
    annotations=_SET,
)
async def slave_set_brightness(device_id: str = "localhost", level: int = 50) -> str:
    """Set slave device brightness level.
    
    Args:
        device_id: The slave device ID (default: localhost).
        level: Brightness level 0-100 (default 50).
    
    Returns:
        Brightness control result as JSON.
    """
    return await _run(device_id, lambda c: c.set_brightness(level))


@mcp.tool(
    description="Increase the display brightness on a slave device.",
    annotations=_STEP,
)
async def slave_increase_brightness(device_id: str = "localhost", amount: int = 5) -> str:
    """Increase slave device brightness.
    
    Args:
        device_id: The slave device ID (default: localhost).
        amount: Amount to increase in % points (default 5).
    
    Returns:
        Brightness control result as JSON.
    """
    return await _run(device_id, lambda c: c.increase_brightness(amount))


@mcp.tool(
    description="Decrease the display brightness on a slave device.",
    annotations=_STEP,
)
async def slave_decrease_brightness(device_id: str = "localhost", amount: int = 5) -> str:
    """Decrease slave device brightness.
    
    Args:
        device_id: The slave device ID (default: localhost).
        amount: Amount to decrease in % points (default 5).
    
    Returns:
        Brightness control result as JSON.
    """
    return await _run(device_id, lambda c: c.decrease_brightness(amount))


# ---- Power Control Tools ----


@mcp.tool(
    description="Put a slave device to sleep.",
    annotations=_POWER,
)
async def slave_sleep(device_id: str = "localhost") -> str:
    """Put slave device to sleep.
    
    Args:
        device_id: The slave device ID (default: localhost).
    
    Returns:
        Power control result as JSON.
    """
    return await _run(device_id, lambda c: c.sleep())


@mcp.tool(
    description="Shutdown a slave device.",
    annotations=_POWER,
)
async def slave_shutdown(device_id: str = "localhost", delay: int = 0) -> str:
    """Shutdown slave device.
    
    Args:
        device_id: The slave device ID (default: localhost).
        delay: Delay in seconds before shutdown (default 0 = immediate).
    
    Returns:
        Power control result as JSON.
    """
    return await _run(device_id, lambda c: c.shutdown(delay))


@mcp.tool(
    description="Restart a slave device.",
    annotations=_POWER,
)
async def slave_restart_system(device_id: str = "localhost", delay: int = 0) -> str:
    """Restart slave device.
    
    Args:
        device_id: The slave device ID (default: localhost).
        delay: Delay in seconds before restart (default 0 = immediate).
    
    Returns:
        Power control result as JSON.
    """
    return await _run(device_id, lambda c: c.restart(delay))


@mcp.tool(
    description="Cancel a pending shutdown or restart on a slave device.",
    annotations=_SET,
)
async def slave_cancel_shutdown(device_id: str = "localhost") -> str:
    """Cancel pending shutdown/restart on slave device.
    
    Args:
        device_id: The slave device ID (default: localhost).
    
    Returns:
        Power control result as JSON.
    """
    return await _run(device_id, lambda c: c.cancel_shutdown())


# ---- Entry Point ----


def main() -> None:
    global _slaves
    logging.basicConfig(level=logging.INFO)
    _slaves = load_slave_config()
    if _slaves:
        log.info(f"Loaded {len(_slaves)} slave device(s): {list(_slaves.keys())}")
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()