from __future__ import annotations

import json
import os
import shlex
import subprocess
import sys
import tempfile
import time
import webbrowser
from pathlib import Path
from typing import Any

import httpx
from InquirerPy import inquirer


DEFAULT_SERVER_URL = "https://alsoml.com"
CLAUDE_CONFIG_PATH = Path.home() / ".claude.json"
WORKOS_CLIENT_ID = "client_01KHFBZZ3JQF0NT8NVJ07GG932"


def _read_claude_config() -> dict[str, Any]:
    if not CLAUDE_CONFIG_PATH.exists():
        return {}
    try:
        return json.loads(CLAUDE_CONFIG_PATH.read_text())
    except (json.JSONDecodeError, OSError):
        return {}


def _write_claude_config(config: dict[str, Any]) -> None:
    fd, tmp_path = tempfile.mkstemp(
        dir=CLAUDE_CONFIG_PATH.parent, suffix=".tmp"
    )
    try:
        with os.fdopen(fd, "w") as f:
            json.dump(config, f, indent=2)
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp_path, str(CLAUDE_CONFIG_PATH))
    except BaseException:
        try:
            os.unlink(tmp_path)
        except OSError:
            pass
        raise


def _get_api_key() -> str:
    config = _read_claude_config()
    entry = config.get("mcpServers", {}).get("also-mcp", {})
    return entry.get("headers", {}).get("X-Api-Key", "")


def _get_server_url() -> str:
    config = _read_claude_config()
    entry = config.get("mcpServers", {}).get("also-mcp", {})
    url = entry.get("url", "")
    if url.endswith("/mcp"):
        return url[:-4]
    return DEFAULT_SERVER_URL


def _build_mcp_entry(server_url: str, api_key: str) -> dict[str, Any]:
    return {
        "type": "http",
        "url": f"{server_url}/mcp",
        "headers": {"X-Api-Key": api_key},
    }


def _sync_mcp_config(server_url: str, api_key: str) -> bool:
    if not api_key:
        return _remove_mcp_config()

    config = _read_claude_config()
    mcp_servers = config.setdefault("mcpServers", {})
    entry = _build_mcp_entry(server_url, api_key)
    mcp_servers["also-mcp"] = entry
    _write_claude_config(config)
    print(f"  Wrote MCP config to {CLAUDE_CONFIG_PATH}")
    print(f"  Entry: {json.dumps(entry, indent=2)}")
    return True


def _remove_mcp_config() -> bool:
    if not CLAUDE_CONFIG_PATH.exists():
        return False

    config = _read_claude_config()
    mcp_servers = config.get("mcpServers", {})
    if "also-mcp" not in mcp_servers:
        return False

    del mcp_servers["also-mcp"]
    if not mcp_servers:
        del config["mcpServers"]

    _write_claude_config(config)
    return True


def _api_headers(api_key: str) -> dict[str, str]:
    return {"X-Api-Key": api_key}


def _has_claude_descendant(pid: str) -> bool:
    try:
        ps_result = subprocess.run(
            ["ps", "-eo", "pid,ppid,command"],
            capture_output=True,
            text=True,
        )
    except FileNotFoundError:
        return False

    children_of: dict[str, list[tuple[str, str]]] = {}
    for line in ps_result.stdout.strip().splitlines()[1:]:
        parts = line.strip().split(None, 2)
        if len(parts) < 3:
            continue
        child_pid, parent_pid, cmd = parts
        children_of.setdefault(parent_pid, []).append((child_pid, cmd))

    stack = [pid]
    while stack:
        current = stack.pop()
        for child_pid, cmd in children_of.get(current, []):
            if "claude" in cmd.lower():
                return True
            stack.append(child_pid)
    return False


def _discover_claude_sessions() -> list[dict[str, str]]:
    result = subprocess.run(
        [
            "tmux", "list-panes", "-a",
            "-F", "#{session_name}:#{window_index}.#{pane_index} #{pane_pid} #{pane_current_command}",
        ],
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        return []

    sessions: list[dict[str, str]] = []
    for line in result.stdout.strip().splitlines():
        parts = line.split(" ", 2)
        if len(parts) < 3:
            continue
        pane_id, pid, command = parts
        if "claude" not in command.lower() and not _has_claude_descendant(pid):
            continue
        session_name, rest = pane_id.split(":", 1)
        sessions.append({
            "pane_id": pane_id,
            "session_name": session_name,
            "pid": pid,
            "command": command,
        })
    return sessions


def _launch_new_claude_session() -> str | None:
    directory = inquirer.filepath(
        message="Directory to start Claude Code in:",
        default=str(Path.cwd()),
        only_directories=True,
    ).execute()
    if not directory:
        return None

    directory = str(Path(directory).resolve())

    has_tmux = subprocess.run(
        ["tmux", "list-sessions"],
        capture_output=True,
    ).returncode == 0

    if has_tmux:
        result = subprocess.run(
            [
                "tmux", "new-window",
                "-P", "-F", "#{session_name}:#{window_index}.#{pane_index}",
            ],
            capture_output=True,
            text=True,
        )
    else:
        result = subprocess.run(
            [
                "tmux", "new-session", "-d", "-s", "also",
                "-P", "-F", "#{session_name}:#{window_index}.#{pane_index}",
            ],
            capture_output=True,
            text=True,
        )

    if result.returncode != 0:
        print(f"  Failed to create tmux session: {result.stderr.strip()}", file=sys.stderr)
        return None

    pane_id = result.stdout.strip()
    subprocess.run(
        ["tmux", "send-keys", "-t", pane_id, f"cd {shlex.quote(directory)} && claude", "Enter"],
        check=True,
    )
    if not has_tmux:
        print("  Created tmux session 'also'. Run 'tmux attach -t also' to view it.")
    print(f"  Launched Claude Code in {directory} (pane {pane_id})")
    return pane_id


def _fetch_groups(server_url: str) -> list[str]:
    try:
        resp = httpx.get(f"{server_url}/auth/groups", timeout=10)
        resp.raise_for_status()
        return resp.json()
    except (httpx.HTTPError, json.JSONDecodeError):
        return []


def _create_group(server_url: str, group_id: str) -> str | None:
    try:
        resp = httpx.post(
            f"{server_url}/auth/groups",
            json={"group_id": group_id},
            timeout=10,
        )
        if resp.status_code == 409:
            print(f"  Group '{group_id}' already exists.", file=sys.stderr)
            return None
        resp.raise_for_status()
        return resp.json()["password"]
    except httpx.HTTPError as e:
        print(f"  Failed to create group: {e}", file=sys.stderr)
        return None


def _get_version() -> str:
    from importlib.metadata import version
    return version("alsogc")


def _authenticate_workos(server_url: str) -> tuple[str, str]:
    resp = httpx.post(
        "https://api.workos.com/user_management/authorize/device",
        data={"client_id": WORKOS_CLIENT_ID},
        timeout=30,
    )
    resp.raise_for_status()
    device_data = resp.json()

    device_code = device_data["device_code"]
    user_code = device_data["user_code"]
    verification_uri_complete = device_data["verification_uri_complete"]
    interval = device_data["interval"]

    print(f"\n  To sign in, visit: {verification_uri_complete}")
    print(f"  Enter code: {user_code}\n")

    try:
        webbrowser.open(verification_uri_complete)
    except Exception:
        pass

    while True:
        time.sleep(interval)
        print(f"  Polling... (interval={interval}s)")
        try:
            poll_resp = httpx.post(
                "https://api.workos.com/user_management/authenticate",
                data={
                    "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
                    "device_code": device_code,
                    "client_id": WORKOS_CLIENT_ID,
                },
                timeout=30,
            )
        except httpx.ReadTimeout:
            print("  Poll timed out, retrying...")
            continue

        print(f"  Response: {poll_resp.status_code} {poll_resp.text[:200]}")

        if poll_resp.status_code == 200:
            token_data = poll_resp.json()
            email = token_data["user"]["email"]
            workos_user_id = token_data["user"]["id"]
            break

        error_body = poll_resp.json()
        error_code = error_body.get("error", "")

        if error_code == "authorization_pending":
            continue
        elif error_code == "slow_down":
            interval += 1
            continue
        else:
            print(f"  Unexpected error: {error_code} — {poll_resp.text[:300]}")
            poll_resp.raise_for_status()

    print(f"  WorkOS auth success: {email}")
    print(f"  Calling {server_url}/auth/login ...")
    login_resp = httpx.post(
        f"{server_url}/auth/login",
        json={"email": email, "workos_user_id": workos_user_id},
        timeout=30,
    )
    print(f"  Login response: {login_resp.status_code} {login_resp.text[:200]}")
    login_resp.raise_for_status()
    api_key = login_resp.json()["api_key"]

    return (api_key, email)


def _action_manage_users(server_url: str, api_key: str) -> None:
    headers = _api_headers(api_key)

    while True:
        action = inquirer.select(
            message="User management:",
            choices=["Create user", "Delete user", "Manage user groups", "Back"],
        ).execute()

        if action == "Back":
            break

        if action == "Create user":
            display_name = inquirer.text(message="Display name:").execute()
            if not display_name:
                continue
            resp = httpx.post(
                f"{server_url}/auth/users",
                headers=headers,
                json={"display_name": display_name},
                timeout=10,
            )
            resp.raise_for_status()
            user = resp.json()
            print(f"\n  Created user: {user}\n")

        elif action == "Delete user":
            resp = httpx.get(f"{server_url}/auth/users", headers=headers, timeout=10)
            resp.raise_for_status()
            users = resp.json()
            if not users:
                print("\n  No users found.\n")
                continue
            choices = [
                {"name": u.get("display_name", u["user_id"]), "value": u["user_id"]}
                for u in users
            ]
            user_id = inquirer.select(message="Select user to delete:", choices=choices).execute()
            httpx.delete(
                f"{server_url}/auth/users/{user_id}", headers=headers, timeout=10
            ).raise_for_status()
            print(f"\n  Deleted user {user_id}.\n")

        elif action == "Manage user groups":
            resp = httpx.get(f"{server_url}/auth/users", headers=headers, timeout=10)
            resp.raise_for_status()
            users = resp.json()
            if not users:
                print("\n  No users found.\n")
                continue
            choices = [
                {"name": u.get("display_name", u["user_id"]), "value": u["user_id"]}
                for u in users
            ]
            user_id = inquirer.select(message="Select user:", choices=choices).execute()

            group_action = inquirer.select(
                message="Group action:",
                choices=["Add to group", "Remove from group", "Back"],
            ).execute()

            if group_action == "Add to group":
                remote_groups = _fetch_groups(server_url)
                group_choices = [{"name": g, "value": g} for g in remote_groups]
                group_choices.append({"name": "-- Create new group --", "value": "__create__"})

                selected = inquirer.select(
                    message="Select group:", choices=group_choices
                ).execute()

                if selected == "__create__":
                    group_name = inquirer.text(message="New group name:").execute()
                    if not group_name:
                        continue
                    password = _create_group(server_url, group_name)
                    if password is None:
                        continue
                    selected = group_name
                    print(f"\n  Created group '{group_name}' with password: {password}\n")
                else:
                    password = inquirer.secret(
                        message=f"Password for '{selected}':"
                    ).execute()

                httpx.post(
                    f"{server_url}/auth/users/{user_id}/groups",
                    headers=headers,
                    json={"group_id": selected, "password": password},
                    timeout=10,
                ).raise_for_status()
                print(f"\n  Added user to group '{selected}'.\n")

            elif group_action == "Remove from group":
                ug_resp = httpx.get(
                    f"{server_url}/auth/users/{user_id}/groups",
                    headers=headers,
                    timeout=10,
                )
                ug_resp.raise_for_status()
                user_groups = ug_resp.json().get("groups", [])
                if not user_groups:
                    print("\n  User is not in any groups.\n")
                    continue
                group_to_remove = inquirer.select(
                    message="Select group to remove from:",
                    choices=user_groups,
                ).execute()
                httpx.delete(
                    f"{server_url}/auth/users/{user_id}/groups/{group_to_remove}",
                    headers=headers,
                    timeout=10,
                ).raise_for_status()
                print(f"\n  Removed user from group '{group_to_remove}'.\n")


def _setup_tui_sidebar(
    target_pane: str,
    server_url: str,
    user_id: str,
    api_key: str,
) -> str | None:
    split_result = subprocess.run(
        [
            "tmux", "split-window", "-h",
            "-t", target_pane,
            "-l", "40%",
            "-P", "-F", "#{session_name}:#{window_index}.#{pane_index}",
        ],
        capture_output=True,
        text=True,
    )
    if split_result.returncode != 0:
        print(f"  Failed to split tmux window: {split_result.stderr}", file=sys.stderr)
        return None

    new_pane = split_result.stdout.strip()
    cmd = (
        f"relay-tui "
        f"--server-url {server_url} "
        f"--user-id {user_id} "
        f"--api-key {api_key} "
        f"--target-pane {target_pane}"
    )
    subprocess.run(["tmux", "send-keys", "-t", new_pane, cmd, "Enter"], check=True)
    subprocess.run(["tmux", "select-pane", "-t", target_pane], check=True)
    return new_pane


def _action_setup_sessions(server_url: str, api_key: str) -> None:
    headers = _api_headers(api_key)

    resp = httpx.get(f"{server_url}/auth/users", headers=headers, timeout=10)
    resp.raise_for_status()
    users = resp.json()

    if not users:
        print("\n  No users found. Create users first.\n")
        return

    sessions = _discover_claude_sessions()
    if not sessions:
        launch = inquirer.confirm(
            message="No Claude sessions found. Launch a new one?", default=True
        ).execute()
        if launch:
            new_pane = _launch_new_claude_session()
            if new_pane:
                sessions = _discover_claude_sessions()
        if not sessions:
            return

    used_users: set[str] = set()
    used_sessions: set[str] = set()
    mappings: list[tuple[dict[str, Any], dict[str, str]]] = []

    while True:
        available_users = [u for u in users if u["user_id"] not in used_users]
        available_sessions = [s for s in sessions if s["pane_id"] not in used_sessions]

        if not available_users or not available_sessions:
            break

        user_choices = [
            {"name": u.get("display_name", u["user_id"]), "value": u["user_id"]}
            for u in available_users
        ]
        user_choices.append({"name": "-- Done mapping --", "value": "__done__"})

        selected_user_id = inquirer.select(
            message="Select a user to map:", choices=user_choices
        ).execute()
        if selected_user_id == "__done__":
            break

        session_choices = [
            {"name": f"{s['pane_id']}  (pid {s['pid']})", "value": s["pane_id"]}
            for s in available_sessions
        ]
        selected_pane = inquirer.select(
            message="Select a session:", choices=session_choices
        ).execute()

        selected_user = next(u for u in users if u["user_id"] == selected_user_id)
        selected_session = next(s for s in sessions if s["pane_id"] == selected_pane)

        used_users.add(selected_user_id)
        used_sessions.add(selected_pane)
        mappings.append((selected_user, selected_session))

    for user, session in mappings:
        pane_id = session["pane_id"]
        user_id = user["user_id"]
        display_name = user.get("display_name", user_id)

        open_tui = inquirer.confirm(
            message=f"Open TUI sidebar for {display_name}?",
            default=True,
        ).execute()

        if open_tui:
            tui_pane = _setup_tui_sidebar(pane_id, server_url, user_id, api_key)
            if tui_pane:
                print(f"  TUI running in pane: {tui_pane}")

    for user, session in mappings:
        pane_id = session["pane_id"]
        user_id = user["user_id"]
        display_name = user.get("display_name", user_id)

        text = (
            f'[ALSO SETUP] You are connected to the Also relay as "{display_name}". '
            f'Your user_id is "{user_id}". Pass this user_id as the first argument '
            f'to every relay tool call (send_message, read_new_messages, etc).'
        )
        subprocess.run(["tmux", "send-keys", "-t", pane_id, "-l", text])
        subprocess.run(["tmux", "send-keys", "-t", pane_id, "Enter"])
        print(f"  Injected setup message into {pane_id} for {display_name}")


def _action_show_status(server_url: str, api_key: str) -> None:
    masked_key = api_key[:8] + "..." if len(api_key) > 8 else api_key
    print(f"\n  Server:   {server_url}")
    print(f"  API Key:  {masked_key}")
    headers = _api_headers(api_key)
    try:
        resp = httpx.get(f"{server_url}/auth/users", headers=headers, timeout=10)
        resp.raise_for_status()
        users = resp.json()
        if users:
            print("  Users:")
            for u in users:
                groups_resp = httpx.get(
                    f"{server_url}/auth/users/{u['user_id']}/groups",
                    headers=headers,
                    timeout=10,
                )
                groups = groups_resp.json().get("groups", []) if groups_resp.status_code == 200 else []
                groups_str = ", ".join(groups) if groups else "(no groups)"
                print(f"    {u.get('display_name', u['user_id'])} [{u['user_id']}] -> {groups_str}")
        else:
            print("  Users:    (none)")
    except httpx.HTTPError:
        print("  Users:    (could not fetch)")
    print()


def main() -> None:
    if len(sys.argv) > 1 and sys.argv[1] in ("--version", "-v"):
        print(f"alsogc {_get_version()}")
        return

    api_key = _get_api_key()
    server_url = _get_server_url()

    if not api_key:
        print("  No API key found. Authenticating...")
        api_key, email = _authenticate_workos(server_url)
        _sync_mcp_config(server_url, api_key)
        print(f"  Authenticated as {email}.\n")
    else:
        _sync_mcp_config(server_url, api_key)
        print(f"  Synced also-mcp server in {CLAUDE_CONFIG_PATH}\n")

    while True:
        action = inquirer.select(
            message="What would you like to do?",
            choices=[
                "Manage users",
                "Setup sessions",
                "Status",
                "Exit",
            ],
        ).execute()

        if action == "Manage users":
            _action_manage_users(server_url, api_key)
        elif action == "Setup sessions":
            _action_setup_sessions(server_url, api_key)
        elif action == "Status":
            _action_show_status(server_url, api_key)
        else:
            break


if __name__ == "__main__":
    main()
