#!/usr/bin/python3
"""Manage lavtomate MCP server connection for Claude Code and Codex.

Usage:
    lavtomate-mcp on [--local]      - connect MCP server (user scope by default)
    lavtomate-mcp off [--local]     - disconnect MCP server
    lavtomate-mcp update [--local]  - refresh expired token and re-register
    lavtomate-mcp status            - show connection status
    lavtomate-mcp test              - test MCP connectivity and show config
    lavtomate-mcp stdio             - MCP over stdio for clients that start a command
                                      (the Agent Plugin, Cursor, VS Code): the key is
                                      minted with the Kerberos ticket, not in a config

Targets:
    By default acts on every installed CLI (claude and/or codex).
    --claude   Act only on Claude Code
    --codex    Act only on Codex
    --zclaude  Act only on the isolated zclaude profile (~/.claude-z, GLM/Z.AI).
               Implies claude; propagates the internal CA into the profile's
               env.conf and uses an independent token (suffix 'zclaude').

Options:
    --local  Add/remove to current project instead of user config (claude only)

Environment:
    LAVTOMATE_URL  - server URL (default: https://lavtomate.office.etersoft.ru)
    CODEX_HOME     - codex config dir (default: ~/.codex)
"""

import json
import os
import re
import shutil
import subprocess
import sys

from lavtomate_client import auth as lavtomate_auth
from lavtomate_client import transport as lavtomate_transport
from lavtomate_client.errors import LavtomateError, LoginRequired

LAVTOMATE_URL = os.environ.get("LAVTOMATE_URL", "https://lavtomate.office.etersoft.ru")
MCP_NAME = "lavtomate"
CLAUDE_ENV_CONF = "/etc/opt/claude.ai/env.conf"
# Claude Code binary shipped by the ALT package. claude expects to find itself
# at ~/.local/bin/claude (the conventional user-install location) and prints a
# warning when it is missing — symlink it so the warning stays silent.
CLAUDE_REAL_BIN = os.environ.get("CLAUDE_BIN", "/opt/claude.ai/claude")
CODEX_HOME = os.environ.get("CODEX_HOME") or os.path.expanduser("~/.codex")

# When targeting --zclaude, claude subprocesses run with this CLAUDE_CONFIG_DIR
# so 'claude mcp add/remove/list' operate on the isolated zclaude profile rather
# than the default ~/.claude. None means the normal (default profile) behaviour.
CLAUDE_CONFIG_DIR = None
ZCLAUDE_DEFAULT = os.environ.get("CLAUDE_CONFIG_DIR") or os.path.expanduser("~/.claude-z")


# --- CLI detection -----------------------------------------------------------

def has_claude():
    return shutil.which("claude") is not None


def has_codex():
    return shutil.which("codex") is not None


def resolve_targets(args):
    """Return list of target CLIs based on flags and what is installed."""
    want = [t for t in ("claude", "codex") if f"--{t}" in args]
    if want:
        return want
    targets = []
    if has_claude():
        targets.append("claude")
    if has_codex():
        targets.append("codex")
    return targets


# --- token -------------------------------------------------------------------

def check_kerberos():
    """Check if Kerberos ticket is available and valid (python gssapi)."""
    return lavtomate_auth.have_ticket()


def get_token(scope, suffix=""):
    """Get a new 30-day API token with the Kerberos ticket (lavtomate_client).

    Uses different description per scope/client to avoid conflicts and to keep
    claude and codex tokens independent (server deduplicates by description).
    Always mints a new token (no reuse of a cached one); the previous token of
    the same description is revoked.

    A token supplied via LAVTOMATE_API_KEY (e.g. forwarded by lavtomate-ssh's
    SendEnv) is used as-is — no Kerberos needed on the remote host.
    """
    env_token = os.environ.get("LAVTOMATE_API_KEY")
    if env_token and env_token.strip():
        return env_token.strip()

    if not check_kerberos():
        print("Error: No Kerberos ticket. Run 'kinit' first "
              "(or forward a token via LAVTOMATE_API_KEY / lavtomate-ssh).",
              file=sys.stderr)
        return None

    hostname = subprocess.run(
        ["hostname", "-s"], capture_output=True, text=True
    ).stdout.strip()
    # Different token name per scope for clarity
    if scope == "user":
        desc = f"mcp-{hostname}"
    else:
        project = os.path.basename(os.getcwd())
        desc = f"mcp-{hostname}-{project}"
    if suffix:
        desc += f"-{suffix}"
    try:
        return lavtomate_auth.mint("30d", desc, LAVTOMATE_URL).api_key
    except LoginRequired:
        print("Error: Kerberos ticket expired or invalid. Run 'kinit'.", file=sys.stderr)
    except LavtomateError as e:
        print(f"Error: {e}", file=sys.stderr)
    return None


# --- Claude Code -------------------------------------------------------------

def claude(*args):
    """Run claude CLI command, return (returncode, stdout).

    When CLAUDE_CONFIG_DIR is set (--zclaude), it is passed through so the
    command operates on the isolated zclaude profile. The claude wrapper still
    sources /etc/opt/claude.ai/env.conf, so the internal CA is available there.
    """
    env = os.environ.copy()
    if CLAUDE_CONFIG_DIR:
        env["CLAUDE_CONFIG_DIR"] = CLAUDE_CONFIG_DIR
    result = subprocess.run(
        ["claude", *args],
        capture_output=True, text=True, env=env,
    )
    return result.returncode, result.stdout


def claude_is_configured():
    rc, output = claude("mcp", "list")
    for line in output.splitlines():
        if line.startswith(f"{MCP_NAME}:"):
            return True
    return False


def ensure_claude_symlink():
    """Ensure ~/.local/bin/claude points at the real binary.

    claude looks for itself at ~/.local/bin/claude (the standard user-install
    location) and warns when it is absent. On hosts where claude is installed
    under /opt/claude.ai/claude, create the symlink so the warning does not
    appear. Idempotent; silent when the target binary is missing or the link
    already exists.
    """
    link = os.path.join(os.path.expanduser("~/.local/bin"), "claude")
    if os.path.lexists(link):
        return
    if not os.path.exists(CLAUDE_REAL_BIN):
        return
    os.makedirs(os.path.dirname(link), exist_ok=True)
    try:
        os.symlink(CLAUDE_REAL_BIN, link)
        print(f"claude: created symlink {link} -> {CLAUDE_REAL_BIN}")
    except OSError as e:
        print(f"claude: could not create symlink {link}: {e}", file=sys.stderr)


def check_extra_ca():
    """Check that NODE_EXTRA_CA_CERTS is set in env.conf for internal TLS."""
    try:
        with open(CLAUDE_ENV_CONF) as f:
            for line in f:
                line = line.strip()
                if line.startswith("#") or "=" not in line:
                    continue
                key, _, value = line.partition("=")
                if key.strip() == "NODE_EXTRA_CA_CERTS" and value.strip():
                    return
    except FileNotFoundError:
        pass
    print(f"Error: NODE_EXTRA_CA_CERTS not found in {CLAUDE_ENV_CONF}", file=sys.stderr)
    print("Run as root: echo 'NODE_EXTRA_CA_CERTS=/path/to/ca.pem' >> "
          f"{CLAUDE_ENV_CONF}", file=sys.stderr)
    sys.exit(1)


def claude_register(scope):
    """Issue token and add claude MCP entry. Caller is responsible for removing
    any pre-existing entry to avoid 'already configured' errors."""
    label = "zclaude" if CLAUDE_CONFIG_DIR else "claude"
    print(f"Getting lavtomate API token ({label})...")
    # Distinct token description for zclaude so it does not collide with the
    # normal claude token (server deduplicates by description).
    token = get_token(scope, suffix="zclaude" if CLAUDE_CONFIG_DIR else "")
    if not token:
        print("Error: Failed to get token", file=sys.stderr)
        sys.exit(1)

    rc, output = claude(
        "mcp", "add",
        "--transport", "http",
        "--scope", scope,
        MCP_NAME, f"{LAVTOMATE_URL}/mcp",
        "--header", f"Authorization: Bearer {token}",
    )
    if rc != 0:
        print(f"Error: {output.strip()}", file=sys.stderr)
        sys.exit(1)


def claude_remove_all():
    """Remove from all scopes to avoid ghost configs."""
    for s in ("user", "local", "project"):
        claude("mcp", "remove", "--scope", s, MCP_NAME)


# --- Codex -------------------------------------------------------------------

def codex(*args):
    """Run codex mcp command, return (returncode, combined output)."""
    result = subprocess.run(
        ["codex", "mcp", *args],
        capture_output=True, text=True
    )
    return result.returncode, (result.stdout or "") + (result.stderr or "")


def codex_is_configured():
    rc, _ = codex("get", MCP_NAME)
    return rc == 0


def codex_config_path():
    return os.path.join(CODEX_HOME, "config.toml")


def codex_register(scope):
    """Issue token and write codex MCP entry with the token baked into
    http_headers (codex CLI exposes only --bearer-token-env-var, but the
    config supports http_headers directly)."""
    print("Getting lavtomate API token (codex)...")
    token = get_token(scope, suffix="codex")
    if not token:
        print("Error: Failed to get token", file=sys.stderr)
        sys.exit(1)

    # Clean any existing entry so we never produce a duplicate TOML table.
    codex("remove", MCP_NAME)

    cfg = codex_config_path()
    os.makedirs(os.path.dirname(cfg), exist_ok=True)
    existing = ""
    if os.path.isfile(cfg):
        with open(cfg) as f:
            existing = f.read()
    if existing and not existing.endswith("\n"):
        existing += "\n"

    block = (
        f"\n[mcp_servers.{MCP_NAME}]\n"
        f'url = "{LAVTOMATE_URL}/mcp"\n'
        f'http_headers = {{ Authorization = "Bearer {token}" }}\n'
    )
    with open(cfg, "w") as f:
        f.write(existing + block)
    os.chmod(cfg, 0o600)


def codex_remove():
    codex("remove", MCP_NAME)


# --- commands ----------------------------------------------------------------

def cmd_on(scope, targets):
    if "claude" in targets:
        ensure_claude_symlink()
        check_extra_ca()
        if claude_is_configured():
            print("claude: lavtomate MCP already configured (use 'update' to refresh token)")
        else:
            claude_register(scope)
            print(f"claude: lavtomate MCP connected ({scope})")
    if "codex" in targets:
        if codex_is_configured():
            print("codex: lavtomate MCP already configured (use 'update' to refresh token)")
        else:
            codex_register(scope)
            print("codex: lavtomate MCP connected (global)")


def cmd_off(scope, targets):
    if "claude" in targets:
        if claude_is_configured():
            claude_remove_all()
            print("claude: lavtomate MCP disconnected")
        else:
            print("claude: lavtomate MCP is not configured")
    if "codex" in targets:
        if codex_is_configured():
            codex_remove()
            print("codex: lavtomate MCP disconnected")
        else:
            print("codex: lavtomate MCP is not configured")


def cmd_update(scope, targets):
    """Re-issue token and re-register MCP (use after token expiry)."""
    if "claude" in targets:
        check_extra_ca()
        claude_remove_all()
        claude_register(scope)
        print(f"claude: lavtomate MCP token refreshed ({scope})")
    if "codex" in targets:
        codex_register(scope)
        print("codex: lavtomate MCP token refreshed (global)")


def cmd_status(scope, targets):
    if "claude" in targets:
        if claude_is_configured():
            _, output = claude("mcp", "list")
            for line in output.splitlines():
                if MCP_NAME in line:
                    print(f"claude: {line}")
        else:
            print("claude: lavtomate MCP not configured")
    if "codex" in targets:
        if codex_is_configured():
            _, output = codex("list")
            for line in output.splitlines():
                if MCP_NAME in line:
                    print(f"codex: {line.strip()}")
        else:
            print("codex: lavtomate MCP not configured")


def find_mcp_config():
    """Find MCP config file containing lavtomate."""
    import json

    cwd = os.getcwd()
    # Project-level .mcp.json
    mcp_json = os.path.join(cwd, ".mcp.json")
    if os.path.isfile(mcp_json):
        try:
            with open(mcp_json) as f:
                data = json.load(f)
            servers = data.get("mcpServers", {})
            if any("lavtomate" in k for k in servers):
                return mcp_json
        except (json.JSONDecodeError, OSError):
            pass

    # Claude project settings
    home = os.path.expanduser("~")
    project_key = "-" + cwd.replace("/", "-")
    project_settings = os.path.join(home, ".claude", "projects", project_key, "settings.json")
    if os.path.isfile(project_settings):
        try:
            with open(project_settings) as f:
                if "lavtomate" in f.read():
                    return project_settings
        except OSError:
            pass

    # Claude global settings
    global_settings = os.path.join(home, ".claude", "settings.json")
    if os.path.isfile(global_settings):
        try:
            with open(global_settings) as f:
                if "lavtomate" in f.read():
                    return global_settings
        except OSError:
            pass

    return None


def configured_tokens(targets):
    """[(where, token)] of the lavtomate MCP tokens already written into the
    claude (user and this project) and codex configs; nothing is minted."""
    found = []
    if "claude" in targets:
        base = CLAUDE_CONFIG_DIR or os.path.expanduser("~")
        path = os.path.join(base, ".claude.json")
        try:
            with open(path) as f:
                data = json.load(f)
        except (OSError, ValueError):
            data = {}
        scopes = [("claude user", data.get("mcpServers") or {}),
                  ("claude project", ((data.get("projects") or {}).get(os.getcwd()) or {}).get("mcpServers") or {})]
        for where, servers in scopes:
            for name, server in servers.items():
                auth = ((server or {}).get("headers") or {}).get("Authorization", "")
                if MCP_NAME in name and auth.startswith("Bearer "):
                    found.append((where, auth[len("Bearer "):]))
    if "codex" in targets:
        try:
            with open(codex_config_path()) as f:
                text = f.read()
        except OSError:
            text = ""
        for match in re.finditer(r'Authorization\s*=\s*"Bearer ([^"]+)"', text):
            found.append(("codex", match[1]))
    return found


def http_get(url, headers=None):
    """GET in process (lavtomate_client transport); returns (http_code, body text).
    http_code "000" when the server could not be reached."""
    try:
        status, data, _ = lavtomate_transport.request("GET", url, headers=headers, timeout=10)
    except LavtomateError as e:
        print(f"  ({e})")
        return "000", ""
    body = json.dumps(data) if isinstance(data, (dict, list)) else (data or "")
    return str(status), body


def cmd_test(scope, targets):
    """Test MCP connectivity and show config."""
    import json

    print("=== Lavtomate MCP Test ===")
    print()

    print(f"Working directory: {os.getcwd()}")
    print(f"Targets: {', '.join(targets) if targets else 'none (no CLI found)'}")
    if CLAUDE_CONFIG_DIR:
        print(f"zclaude profile: {CLAUDE_CONFIG_DIR}")
    print()

    # Claude-specific config and status
    if "claude" in targets:
        config_path = find_mcp_config()
        if config_path:
            print(f"Claude MCP config: {config_path}")
        else:
            print("Claude MCP config: NOT FOUND (run 'lavtomate-mcp on' first)")

        print("Claude MCP status:")
        if claude_is_configured():
            _, output = claude("mcp", "list")
            has_failed = False
            for line in output.splitlines():
                if MCP_NAME in line:
                    if "Failed to connect" in line:
                        has_failed = True
                    print(f"  {line}")
            if has_failed:
                print("  HINT: Try running 'claude' in this directory first to approve workspace trust")
        else:
            print("  NOT configured in claude")

        print("Claude TLS config:")
        try:
            with open(CLAUDE_ENV_CONF) as f:
                for line in f:
                    if "NODE_EXTRA_CA_CERTS" in line and not line.strip().startswith("#"):
                        print(f"  {line.strip()}")
                        break
                else:
                    print(f"  WARNING: NODE_EXTRA_CA_CERTS not in {CLAUDE_ENV_CONF}")
        except FileNotFoundError:
            print(f"  WARNING: {CLAUDE_ENV_CONF} not found")
        print()

    # Codex-specific status
    if "codex" in targets:
        print(f"Codex config: {codex_config_path()}")
        print("Codex MCP status:")
        if codex_is_configured():
            _, output = codex("list")
            for line in output.splitlines():
                if MCP_NAME in line or line.lower().startswith("name"):
                    print(f"  {line.rstrip()}")
        else:
            print("  NOT configured in codex")
        print()

    # Kerberos (shared)
    print("Kerberos:")
    if check_kerberos():
        info = lavtomate_auth.ticket_info()
        if info:
            print(f"  Principal: {info[0]} (valid {info[1] // 3600}h {info[1] % 3600 // 60}m more)")
        else:
            print("  OK (ticket valid)")
    else:
        print("  NO TICKET — run 'kinit' to authenticate")
    print()

    # Server connectivity (shared)
    print(f"Server: {LAVTOMATE_URL}")
    code, body = http_get(f"{LAVTOMATE_URL}/mcp-status")
    if code == "200":
        try:
            data = json.loads(body)
            print(f"  Status: {data.get('status', 'unknown')}")
            for name, info in data.get("transports", {}).items():
                print(f"  Transport {name}: {info.get('endpoint', '?')}")
        except json.JSONDecodeError:
            print(f"  OK (HTTP {code}) but invalid JSON")
    elif code == "000":
        print("  FAIL: Connection refused")
    else:
        print(f"  FAIL: HTTP {code}")
    print()

    # Authenticated API (shared): the tokens already configured, never a new one
    # (minting a token of the same description would replace the configured one)
    print("Authentication:")
    tokens = configured_tokens(targets)
    if not tokens:
        print("  No configured token (run 'lavtomate-mcp on'); checking Kerberos login only")
        try:
            code, _ = lavtomate_transport.negotiate_request(
                "GET", f"{LAVTOMATE_URL}/api/auth/api-keys", timeout=10)
            print(f"  Kerberos login: {'OK' if code == 200 else f'HTTP {code}'}")
        except LavtomateError as e:
            print(f"  Kerberos login: FAIL ({e})")
    for where, token in tokens:
        print(f"  Token ({where}): ***{token[-4:]}")
        code, body = http_get(
            f"{LAVTOMATE_URL}/api/v1/bugzilla/instances",
            headers={"Authorization": f"Bearer {token}"},
        )
        if code == "200":
            try:
                instances = json.loads(body).get("instances", [])
                print(f"    API: OK ({len(instances)} Bugzilla instances)")
            except (json.JSONDecodeError, AttributeError):
                print(f"    API: OK (HTTP {code})")
        else:
            print(f"    API: HTTP {code} (expired or replaced: run 'lavtomate-mcp update')")


def main():
    args = sys.argv[1:]
    # Default scope is user; --local overrides to local (claude only)
    scope = "local" if "--local" in args else "user"

    if "--zclaude" in args:
        # Target the isolated zclaude profile (GLM/Z.AI). Implies claude only.
        global CLAUDE_CONFIG_DIR
        CLAUDE_CONFIG_DIR = ZCLAUDE_DEFAULT
        if not os.path.isdir(CLAUDE_CONFIG_DIR):
            print(f"Error: zclaude profile {CLAUDE_CONFIG_DIR} does not exist yet.",
                  file=sys.stderr)
            print("Run 'zclaude' once first (it seeds the GLM/Z.AI config), then "
                  "re-run this command.", file=sys.stderr)
            sys.exit(1)
        targets = ["claude"]
    else:
        targets = resolve_targets(args)
    args = [a for a in args if a not in ("--local", "--user", "--claude", "--codex", "--zclaude")]

    cmd = args[0] if args else "status"

    if cmd == "stdio":
        # stdout carries the protocol: nothing else may be printed there
        from lavtomate_client import mcp_bridge
        mcp_bridge.main(LAVTOMATE_URL)
        return

    if cmd in ("-h", "--help", "help"):
        print("Usage: lavtomate-mcp [on|off|update|status|test] [--claude|--codex|--zclaude] [--local]")
        print()
        print("Commands:")
        print("  on      Connect lavtomate MCP server (user scope by default)")
        print("  off     Disconnect lavtomate MCP server")
        print("  update  Refresh expired token and re-register MCP")
        print("  status  Show connection status (default)")
        print("  test    Test MCP connectivity and show config")
        print("  stdio   MCP over stdio (for the Agent Plugin, Cursor, VS Code)")
        print()
        print("Targets (default: every installed CLI):")
        print("  --claude   Act only on Claude Code")
        print("  --codex    Act only on Codex")
        print("  --zclaude  Act only on the isolated zclaude profile (~/.claude-z, GLM/Z.AI)")
        print()
        print("Options:")
        print("  --local  Use project config instead of user (claude only)")
        return

    commands = {
        "on": cmd_on, "add": cmd_on,
        "off": cmd_off, "remove": cmd_off, "rm": cmd_off,
        "update": cmd_update, "refresh": cmd_update, "renew": cmd_update,
        "status": cmd_status,
        "test": cmd_test,
    }

    if cmd not in commands:
        print("Usage: lavtomate-mcp [on|off|update|status|test] [--claude|--codex] [--local]",
              file=sys.stderr)
        sys.exit(1)

    if not targets:
        print("Error: neither 'claude' nor 'codex' CLI found in PATH", file=sys.stderr)
        sys.exit(1)

    commands[cmd](scope, targets)


if __name__ == "__main__":
    main()
