#!/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

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 os
import shutil
import subprocess
import sys

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."""
    result = subprocess.run(
        ["klist", "-s"],
        capture_output=True, text=True,
    )
    return result.returncode == 0


def get_token(scope, suffix=""):
    """Get API token via lavtomate-get-token.

    Uses different description per scope/client to avoid conflicts and to keep
    claude and codex tokens independent (server deduplicates by description).
    No client-side caching.

    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}"
    result = subprocess.run(
        ["lavtomate-get-token", "30d", desc],
        capture_output=True, text=True,
        env={**os.environ, "LAVTOMATE_NO_CACHE": "1"},
    )
    if result.returncode != 0:
        stderr = result.stderr.strip()
        if "401" in stderr or "Authorization" in stderr:
            print("Error: Kerberos ticket expired or invalid. Run 'kinit'.", file=sys.stderr)
        elif "jq" in stderr:
            print("Error: Server returned invalid response (check Kerberos: run 'kinit').", file=sys.stderr)
        else:
            print(f"Error: {stderr}", file=sys.stderr)
        return None
    return result.stdout.strip()


# --- 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 curl_get(url, headers=None):
    """Run curl and return (http_code, body)."""
    cmd = ["curl", "-s", "-k", "-w", "\n%{http_code}", "--max-time", "10"]
    if headers:
        for k, v in headers.items():
            cmd += ["-H", f"{k}: {v}"]
    cmd.append(url)
    result = subprocess.run(cmd, capture_output=True, text=True)
    lines = result.stdout.rsplit("\n", 1)
    body = lines[0] if len(lines) > 1 else ""
    code = lines[-1].strip() if lines else "000"
    return code, 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():
        klist = subprocess.run(["klist"], capture_output=True, text=True)
        for line in klist.stdout.splitlines():
            if "principal" in line.lower() or "expires" in line.lower():
                print(f"  {line.strip()}")
                break
        else:
            print("  OK (ticket valid)")
    else:
        print("  NO TICKET — run 'kinit' to authenticate")
    print()

    # Server connectivity (shared)
    print(f"Server: {LAVTOMATE_URL}")
    code, body = curl_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)
    print("Authentication:")
    token = get_token("user")
    if token:
        print(f"  Token: ***{token[-4:]}")
        code, body = curl_get(
            f"{LAVTOMATE_URL}/api/v1/bugzilla/instances",
            headers={"Authorization": f"Bearer {token}"},
        )
        if code == "200":
            try:
                data = json.loads(body)
                instances = data.get("instances", [])
                print(f"  API: OK ({len(instances)} Bugzilla instances)")
                for inst in instances:
                    print(f"    - {inst['name']} ({inst['url']})")
            except json.JSONDecodeError:
                print(f"  API: OK (HTTP {code})")
        else:
            print(f"  API: HTTP {code}")
    else:
        print("  Token: FAIL (cannot get token)")


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 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()
        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()
