# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Tomasz Głuc
"""A complete MCP server for Omni in about 110 lines.

Every method of the running Omni becomes an MCP tool, built from the catalog
(`api.methods`) at start-up, so a newer Omni brings its new tools along.

    pip install mcp
    python omni_mcp_server.py            # speaks MCP over stdio

Register it in an MCP client like any stdio server, for example in
Claude Desktop's config:

    {"mcpServers": {"omni": {"command": "python",
                             "args": ["C:/path/to/omni_mcp_server.py"]}}}

Start Omni first, with the API on (Settings, AI Create): the server reads the
catalog from it and waits up to 30 seconds for Omni to appear. After that, if
Omni is closed or restarted, tool calls answer with an error and the next call
connects again.
"""

from __future__ import annotations

import asyncio
import json
import sys
import time

from mcp import types
from mcp.server.lowlevel import Server
from mcp.server.stdio import stdio_server

from omnidaw import Omni

INSTRUCTIONS = (
    "Tools drive the project open in Omni, a DAW. Start with session_info and "
    "tracks_list. Time is in beats, levels in dB, pitch as MIDI numbers. Every "
    "edit is one undo step: edit_undo takes back the newest. Long work "
    "(analysis_capture, export_render) returns a job: call jobs_wait with its id."
)

WAIT_FOR_OMNI_SECONDS = 30

_omni: Omni | None = None
#: MCP tool name -> catalog method, filled once at start-up.
_catalog: dict[str, dict] = {}


def connect() -> Omni:
    """The connection to Omni, opened on first use and again after it drops."""
    global _omni
    if _omni is None:
        _omni = Omni("omni-mcp-example")
    return _omni


def tool_of(method: dict) -> types.Tool:
    """One catalog method as an MCP tool."""
    effect = method["effect"]
    return types.Tool(
        name=method["name"].replace(".", "_"),
        description=method["summary"],
        input_schema=method["params"],
        annotations=types.ToolAnnotations(
            title=method["title"],
            read_only_hint=effect == "read",
            destructive_hint=effect == "destructive",
            idempotent_hint=method["idempotent"],
            open_world_hint=False,
        ),
    )


def run_call(method: str, arguments: dict) -> str:
    """The blocking part of a tool call; runs in a thread."""
    global _omni
    try:
        return json.dumps(connect().call(method, arguments), indent=2, ensure_ascii=False)
    except (ConnectionError, OSError):
        _omni = None  # Omni closed or restarted: the next call connects again
        raise


def error(text: str) -> types.CallToolResult:
    """A tool error the model can read and correct itself from."""
    return types.CallToolResult(is_error=True, content=[types.TextContent(type="text", text=text)])


async def list_tools(ctx, params) -> types.ListToolsResult:
    return types.ListToolsResult(tools=[tool_of(m) for m in _catalog.values()])


async def call_tool(ctx, params) -> types.CallToolResult:
    method = _catalog.get(params.name)
    if method is None:
        return error(f"unknown tool {params.name}")
    try:
        text = await asyncio.to_thread(run_call, method["name"], params.arguments or {})
        return types.CallToolResult(content=[types.TextContent(type="text", text=text)])
    except (ConnectionError, OSError) as e:
        return error(f"Omni is not reachable ({e}). Start Omni with the API on and try again.")
    except Exception as e:  # an OmniError: bad parameters, unknown track, ...
        return error(str(e))


def load_catalog() -> None:
    """Read the tool list from Omni, waiting for it to start if needed."""
    global _omni
    deadline = time.monotonic() + WAIT_FOR_OMNI_SECONDS
    while True:
        try:
            _catalog.update({m["name"].replace(".", "_"): m for m in connect().methods.values()})
            return
        except (ConnectionError, OSError) as e:
            _omni = None
            if time.monotonic() > deadline:
                sys.exit(f"Omni is not running with the API on ({e}). Start it and try again.")
            time.sleep(1)


async def main() -> None:
    server = Server(
        "omni-example",
        instructions=INSTRUCTIONS,
        on_list_tools=list_tools,
        on_call_tool=call_tool,
    )
    async with stdio_server() as (read, write):
        await server.run(read, write, server.create_initialization_options())


if __name__ == "__main__":
    load_catalog()
    asyncio.run(main())