#!/usr/bin/env python3
"""MCP server for mux — lets Claude manage desks, windows, and notes."""

import json
import subprocess
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from muxlib import load_json, load_projects as mux_load_projects, run

MUX_INSTRUCTIONS = """
The user manages multiple projects using `mux`, a tmux wrapper. Here's how it works:

## Core concepts
- A "desk" is a tmux session — one per project. Each desk has windows.
- The user thinks in desks, not tmux sessions. Never say "tmux session" — say "desk."
- tmux is the invisible engine. The user doesn't know tmux commands.

## Key commands (user types these in a shell, not in Claude)
- `mux` — fuzzy picker to switch desks
- `mux <project>` — open/switch to a project desk
- `mux <project> "prompt"` — open a desk and start Claude with a prompt
- `mux new "prompt"` — new window in current desk (starts Claude)
- `mux ls` — list all desks and windows
- `mux note "text"` — add a sticky note (max 3 per project)
- `mux note` — show notes on current desk
- `mux note rm <n>` — remove a note
- `mux resume <session-id>` — resume a Claude session in a new worktree window
- `mux close` / `mux done` — park a desk (keeps running)
- `mux kill <project>` — destroy a desk (asks first)
- `mux split` — split pane
- `mux help` — contextual help
- `mux portal on/off` — toggle spoken project name on switch
- `mux worktree ls/cleanup/remove` — manage git worktrees

## tmux keys (the user MUST learn these)
- Ctrl-Space is the prefix key. Nothing visible happens — it tells tmux "next key is a command."
- Ctrl-Space then d — detach (leave tmux, desk keeps running)
- Ctrl-Space then c — new window
- Ctrl-Space then 1/2/3 — switch to window by number
- Ctrl-Space then s — session picker (switch desks)
- F1 — help popup
- The status bar at the bottom shows: [session-name] 1:window-name 2:window-name

## Portal
Each project has a distinct background color. Switching projects changes the terminal color
and speaks the project name. Toggle voice with `mux portal on/off`.

## Worktrees
When a second Claude session opens in the same project, mux auto-creates a git worktree
for file isolation while sharing project identity (skills, memory, hooks, settings)
via symlinks. Windows in worktrees show (wt) in `mux ls`. Clean worktrees are
auto-removed when the window closes; dirty ones leave a sticky note.

## Sticky notes
Max 3 per project. Auto-dim after 2 sessions, auto-archive after 5. Show as popup on arrival.

## When the user is confused about tmux
Explain using desk/window metaphors. The status bar is their map. Ctrl-Space is like
pressing Shift — it just tells tmux the next key is special. If they're stuck in a
split pane, tell them to type `exit`. If they need a quick shell, tell them to press
F3 (popup shell — appears, use it, press exit, it's gone). F3 is the escape hatch.
Don't tell them to open new windows for shell access — that causes window proliferation.
Windows are for Claude sessions. Shells are F3 popups.

## Available projects
Use mux_list to see what's currently open, or read ~/.config/mux/projects.json for
all configured projects.
"""

TOOLS = [
    {
        "name": "mux_switch",
        "description": "Switch to a different project desk. The terminal will switch to that project's tmux session. The current Claude session keeps running in the background.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "project": {
                    "type": "string",
                    "description": "Project name (e.g., 'maginnis', 'cabinet', 'flow')"
                }
            },
            "required": ["project"]
        }
    },
    {
        "name": "mux_list",
        "description": "List all open desks (tmux sessions) and their windows with activity status.",
        "inputSchema": {"type": "object", "properties": {}}
    },
    {
        "name": "mux_new",
        "description": "Open a new window in the current desk, optionally starting Claude with a prompt. Auto-creates an isolated worktree (with shared project identity) if another Claude session is already running.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "prompt": {
                    "type": "string",
                    "description": "Optional prompt for the new Claude session. Also becomes the window name."
                }
            }
        }
    },
    {
        "name": "mux_note_add",
        "description": "Add a sticky note to a project desk. Notes surface as a popup when you switch to that project. Max 3 per project.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "text": {
                    "type": "string",
                    "description": "The note text"
                },
                "project": {
                    "type": "string",
                    "description": "Project name. If omitted, uses the current desk."
                }
            },
            "required": ["text"]
        }
    },
    {
        "name": "mux_note_list",
        "description": "List sticky notes on the current or specified project desk.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "project": {
                    "type": "string",
                    "description": "Project name. If omitted, uses the current desk."
                }
            }
        }
    },
    {
        "name": "mux_note_remove",
        "description": "Remove (archive) a sticky note by number.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "number": {
                    "type": "integer",
                    "description": "Note number (from mux_note_list)"
                }
            },
            "required": ["number"]
        }
    },
    {
        "name": "mux_status",
        "description": "Show status of a specific desk or portfolio summary of all desks.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "project": {
                    "type": "string",
                    "description": "Project name. If omitted, shows portfolio summary."
                }
            }
        }
    }
]

def current_session():
    out, _, rc = run("tmux display-message -p '#{session_name}' 2>/dev/null")
    return out if rc == 0 else None

def handle_tool(name, args):
    projects = mux_load_projects()
    helper = os.path.expanduser("~/.config/mux/manage-notes.py")

    if name == "mux_switch":
        project = args["project"]
        if project not in projects:
            return f"No project called '{project}'. Available: {', '.join(sorted(projects.keys()))}"
        out, err, rc = run(f"tmux switch-client -t '={project}' 2>&1")
        if rc == 0:
            return f"Switched to {project}."
        # Session might not exist yet — create it
        path = projects[project]["path"]
        run(f"tmux new-session -d -s '{project}' -c '{path}'")
        out, err, rc = run(f"tmux switch-client -t '={project}' 2>&1")
        if rc == 0:
            return f"Opened and switched to {project}."
        return f"Failed to switch: {err or out}"

    elif name == "mux_list":
        out, _, rc = run(os.path.expanduser("~/.local/bin/mux") + " ls 2>&1")
        return out if out else "No desks open."

    elif name == "mux_new":
        prompt = args.get("prompt", "")
        if prompt:
            out, _, _ = run(f'{os.path.expanduser("~/.local/bin/mux")} new "{prompt}" 2>&1')
            return out or f"New window '{prompt}' created."
        else:
            out, _, _ = run(f'{os.path.expanduser("~/.local/bin/mux")} new 2>&1')
            return out or "New window created."

    elif name == "mux_note_add":
        text = args["text"]
        project = args.get("project") or current_session()
        if not project:
            return "Not in a desk and no project specified."
        result = subprocess.run(
            ["python3", helper, "add", project, text],
            capture_output=True, text=True
        )
        out = result.stdout.strip()
        return out

    elif name == "mux_note_list":
        project = args.get("project") or current_session()
        if not project:
            return "Not in a desk and no project specified."
        result = subprocess.run(
            ["python3", helper, "list", project],
            capture_output=True, text=True
        )
        out = result.stdout.strip()
        return out or "No notes."

    elif name == "mux_note_remove":
        num = args["number"]
        project = current_session()
        if not project:
            return "Not in a desk."
        result = subprocess.run(
            ["python3", helper, "rm", project, str(num)],
            capture_output=True, text=True
        )
        out = result.stdout.strip()
        return out

    elif name == "mux_status":
        project = args.get("project", "")
        if project:
            out, _, _ = run(f'{os.path.expanduser("~/.local/bin/mux")} status "{project}" 2>&1')
        else:
            out, _, _ = run(f'{os.path.expanduser("~/.local/bin/mux")} status 2>&1')
        return out or "No desks open."

    return f"Unknown tool: {name}"

def send(msg):
    sys.stdout.write(json.dumps(msg) + "\n")
    sys.stdout.flush()

def main():
    buffer = ""
    for line in sys.stdin:
        buffer += line.strip()
        if not buffer:
            continue
        try:
            req = json.loads(buffer)
            buffer = ""
        except json.JSONDecodeError:
            continue

        if req is None:
            break

        rid = req.get("id")
        method = req.get("method", "")
        params = req.get("params", {})

        if method == "initialize":
            send({"jsonrpc": "2.0", "id": rid, "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {"tools": {}},
                "serverInfo": {"name": "mux", "version": "1.0.0"},
                "instructions": MUX_INSTRUCTIONS
            }})
        elif method == "notifications/initialized":
            pass
        elif method == "tools/list":
            send({"jsonrpc": "2.0", "id": rid, "result": {"tools": TOOLS}})
        elif method == "tools/call":
            tool_name = params.get("name", "")
            tool_args = params.get("arguments", {})
            try:
                result = handle_tool(tool_name, tool_args)
            except Exception as e:
                result = f"Error: {e}"
            send({"jsonrpc": "2.0", "id": rid, "result": {
                "content": [{"type": "text", "text": result}]
            }})
        elif method == "ping":
            send({"jsonrpc": "2.0", "id": rid, "result": {}})
        else:
            send({"jsonrpc": "2.0", "id": rid, "error": {
                "code": -32601, "message": f"Unknown method: {method}"
            }})

if __name__ == "__main__":
    main()