#!/usr/bin/env python3
import argparse
import os
import re
import shlex
import subprocess
import sys
import textwrap
import uuid
from pathlib import Path
import time as _time

USER_NAME = "claude"
HOME_DIR = Path("/home/claude")
CLAUDE_HOME = HOME_DIR / ".claude"
PROMPTS_DIR = HOME_DIR / ".prompts"
SESSIONS_DIR = HOME_DIR / ".delegate-claude-code" / "sessions"
DEFAULT_REPO = Path("/home/workspace")
TMUX_PREFIX = "claude-delegate-"
TMUX_TMPDIR = "/home/claude/.tmux-tmp"
UUID_RE = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$", re.I)


def run(cmd, check=True, capture=False, user=None):
    env = os.environ.copy()
    env["TMUX_TMPDIR"] = TMUX_TMPDIR
    if user:
        cmd = ["runuser", "-u", user, "--", *cmd]
    return subprocess.run(cmd, check=check, text=True, capture_output=capture, env=env)


def ensure_user():
    if subprocess.run(["id", USER_NAME], capture_output=True, text=True).returncode != 0:
        raise SystemExit("claude user does not exist. Run: python3 Skills/claude-code-delegate/scripts/setup.py")


def ensure_auth():
    creds = CLAUDE_HOME / ".credentials.json"
    if not creds.exists():
        raise SystemExit("Claude credentials missing for claude user. Ask the user to sign in via Zo terminal, then rerun setup.")
    status = run(["claude", "auth", "status"], check=False, capture=True, user=USER_NAME)
    if status.returncode != 0:
        raise SystemExit("Claude auth status failed for claude user. Ask the user to sign in via Zo terminal.")
    if '"loggedIn": true' not in status.stdout:
        raise SystemExit("claude user is not signed in. Ask the user to sign in via Zo terminal.")


def ensure_dirs():
    for path in [PROMPTS_DIR, SESSIONS_DIR, CLAUDE_HOME]:
        path.mkdir(parents=True, exist_ok=True)
    subprocess.run(["chown", "-R", f"{USER_NAME}:{USER_NAME}", str(HOME_DIR)], check=True)


def valid_uuid(value):
    if not UUID_RE.match(value):
        raise argparse.ArgumentTypeError(f"invalid uuid: {value}")
    return value


def session_tmux_name(session_id):
    return f"{TMUX_PREFIX}{session_id}"


def session_meta_path(session_id):
    return SESSIONS_DIR / session_id / "meta.env"


def prompt_path(session_id, prompt_id):
    return PROMPTS_DIR / session_id / f"{prompt_id}.md"


def session_exists(session_id):
    name = session_tmux_name(session_id)
    return run(["tmux", "has-session", "-t", name], check=False, capture=True, user=USER_NAME).returncode == 0


def write_session_meta(session_id, repo_dir):
    session_dir = SESSIONS_DIR / session_id
    session_dir.mkdir(parents=True, exist_ok=True)
    meta = session_meta_path(session_id)
    meta.write_text(f"SESSION_ID={session_id}\nREPO_DIR={repo_dir}\nTMUX_SESSION={session_tmux_name(session_id)}\n")
    subprocess.run(["chown", "-R", f"{USER_NAME}:{USER_NAME}", str(session_dir)], check=True)
    return meta


def read_session_meta(session_id):
    meta = session_meta_path(session_id)
    if not meta.exists():
        raise SystemExit(f"session metadata missing for {session_id}")
    data = {}
    for line in meta.read_text().splitlines():
        if "=" in line:
            k, v = line.split("=", 1)
            data[k] = v
    return data


def create_prompt(args):
    ensure_user()
    ensure_dirs()
    session_id = args.session_id or str(uuid.uuid4())
    prompt_id = str(uuid.uuid4())
    session_dir = PROMPTS_DIR / session_id
    session_dir.mkdir(parents=True, exist_ok=True)
    content = args.text if args.text is not None else sys.stdin.read()
    if not content.strip():
        raise SystemExit("prompt text is empty")
    target = prompt_path(session_id, prompt_id)
    target.write_text(content)
    subprocess.run(["chown", "-R", f"{USER_NAME}:{USER_NAME}", str(session_dir)], check=True)
    print(f"SESSION_ID={session_id}")
    print(f"PROMPT_ID={prompt_id}")
    print(f"PROMPT_FILE={target}")


def ensure_session_started(session_id, repo_dir):
    name = session_tmux_name(session_id)
    if session_exists(session_id):
        return name
    repo_dir = str(Path(repo_dir).resolve())
    quoted_repo = shlex.quote(repo_dir)
    cmd = (
        f"cd {quoted_repo} && "
        f"exec claude --dangerously-skip-permissions"
    )
    run(["tmux", "new-session", "-d", "-s", name, cmd], user=USER_NAME)
    subprocess.run(["sleep", "3"], check=True)
    run(["tmux", "send-keys", "-t", name, "\x1b[B"], user=USER_NAME)
    subprocess.run(["sleep", "0.5"], check=True)
    run(["tmux", "send-keys", "-t", name, "Enter"], user=USER_NAME)
    wait_for_claude_prompt(name, timeout=30)
    return name


def wait_for_claude_prompt(name, timeout=30):
    for _ in range(timeout * 2):
        out = run(["tmux", "capture-pane", "-t", name, "-p", "-S", "-80"], check=False, capture=True, user=USER_NAME)
        text = out.stdout
        if "❯" in text and ("bypass permissions" in text or "Send a message" in text or "Claude Code" in text):
            return True
        if out.returncode != 0:
            break
        subprocess.run(["sleep", "0.5"], check=True)
    return False


def send_prompt(session_id, prompt_id):
    ensure_user()
    ensure_auth()
    ensure_dirs()
    target = prompt_path(session_id, prompt_id)
    if not target.exists():
        raise SystemExit(f"prompt file not found: {target}")
    meta = read_session_meta(session_id)
    repo_dir = meta.get("REPO_DIR")
    if not repo_dir:
        raise SystemExit(f"session repo dir missing for {session_id}")
    name = ensure_session_started(session_id, repo_dir)
    wait_for_claude_prompt(name)
    text = target.read_text().rstrip()
    if not text:
        raise SystemExit("prompt file is empty")
    run(["tmux", "send-keys", "-t", name, "-l", text], user=USER_NAME)
    run(["tmux", "send-keys", "-t", name, "Enter"], user=USER_NAME)
    print(f"DISPATCHED session={session_id} prompt={prompt_id} tmux={name}")


def create_session(args):
    ensure_user()
    ensure_auth()
    ensure_dirs()
    session_id = args.session_id or str(uuid.uuid4())
    repo_dir = str(Path(args.repo).resolve())
    if not Path(repo_dir).exists():
        raise SystemExit(f"repo dir does not exist: {repo_dir}")
    write_session_meta(session_id, repo_dir)
    name = ensure_session_started(session_id, repo_dir)
    print(f"SESSION_ID={session_id}")
    print(f"TMUX_SESSION={name}")
    print(f"REPO_DIR={repo_dir}")


def close_session(args):
    session_id = args.session_id
    name = session_tmux_name(session_id)
    run(["tmux", "kill-session", "-t", name], check=False, user=USER_NAME)
    print(f"CLOSED {session_id}")


def list_sessions(_args):
    ensure_dirs()
    found = False
    for meta in sorted(SESSIONS_DIR.glob("*/meta.env")):
        data = {}
        for line in meta.read_text().splitlines():
            if "=" in line:
                k, v = line.split("=", 1)
                data[k] = v
        session_id = data.get("SESSION_ID", meta.parent.name)
        tmux_name = data.get("TMUX_SESSION", session_tmux_name(session_id))
        active = session_exists(session_id)
        found = True
        print(f"{session_id}\t{'active' if active else 'stopped'}\t{data.get('REPO_DIR','') }\t{tmux_name}")
    if not found:
        print("NO_SESSIONS")


def read_output(args):
    """Capture tmux pane content and report idle/busy status."""
    session_id = args.session_id
    if not session_exists(session_id):
        raise SystemExit(f"session {session_id} is not active")
    name = session_tmux_name(session_id)
    lines = str(args.lines)
    out = run(["tmux", "capture-pane", "-t", name, "-p", "-S", f"-{lines}"],
              check=True, capture=True, user=USER_NAME)
    text = out.stdout
    print(text, end="")
    stripped_lines = [l for l in text.splitlines() if l.strip()]
    tail = stripped_lines[-5:] if len(stripped_lines) >= 5 else stripped_lines
    if any("❯" in l for l in tail):
        print("\nSTATUS=idle")
    else:
        print("\nSTATUS=busy")


def send_cmd(args):
    """Create a prompt and dispatch it in one step."""
    ensure_user()
    ensure_auth()
    ensure_dirs()
    session_id = args.session_id
    meta = read_session_meta(session_id)
    content = args.text if args.text is not None else sys.stdin.read()
    if not content.strip():
        raise SystemExit("prompt text is empty")
    prompt_id = str(uuid.uuid4())
    session_dir = PROMPTS_DIR / session_id
    session_dir.mkdir(parents=True, exist_ok=True)
    target = prompt_path(session_id, prompt_id)
    target.write_text(content)
    subprocess.run(["chown", "-R", f"{USER_NAME}:{USER_NAME}", str(session_dir)], check=True)
    send_prompt(session_id, prompt_id)


def watch_cmd(args):
    """Poll tmux output until Claude is idle, then print the result."""
    session_id = args.session_id
    if not session_exists(session_id):
        raise SystemExit(f"session {session_id} is not active")
    name = session_tmux_name(session_id)
    timeout = args.timeout
    interval = args.interval
    quiet = args.quiet
    start = _time.time()
    initial_snapshot = _capture_pane(name, 500)
    _time.sleep(interval)
    while _time.time() - start < timeout:
        text = _capture_pane(name, 500)
        stripped = [l for l in text.splitlines() if l.strip()]
        tail = stripped[-5:] if len(stripped) >= 5 else stripped
        is_idle = bool(stripped) and any("❯" in l for l in tail)
        has_new_content = text != initial_snapshot
        if is_idle and has_new_content:
            if not quiet:
                print(text, end="")
            print("\nSTATUS=idle")
            return
        if not quiet:
            sys.stderr.write(".")
            sys.stderr.flush()
        _time.sleep(interval)
    print(_capture_pane(name, 500), end="")
    print("\nSTATUS=timeout", file=sys.stderr)
    sys.exit(1)


def _capture_pane(name, lines):
    out = run(["tmux", "capture-pane", "-t", name, "-p", "-S", f"-{lines}"],
              check=True, capture=True, user=USER_NAME)
    return out.stdout


def install_stub(target_path):
    script = textwrap.dedent(
        f"""\
        #!/usr/bin/env bash
        set -euo pipefail
        exec python3 /home/workspace/Skills/claude-code-delegate/scripts/claude_delegate.py dispatch "$@"
        """
    )
    Path(target_path).write_text(script)
    os.chmod(target_path, 0o755)
    print(target_path)


def dispatch_compat(args):
    parts = args.prompt_ref.split("/")
    if len(parts) != 2:
        raise SystemExit("-p must be sessionid/promptid")
    session_id, prompt_id = parts
    valid_uuid(session_id)
    valid_uuid(prompt_id)
    send_prompt(session_id, prompt_id)


def main():
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="cmd", required=True)

    p_setup = sub.add_parser("install-stub")
    p_setup.add_argument("--target", default="/usr/local/bin/delegate-claude-code")
    p_setup.set_defaults(func=lambda a: install_stub(a.target))

    p_create_session = sub.add_parser("create-session")
    p_create_session.add_argument("--session-id")
    p_create_session.add_argument("--repo", default=str(DEFAULT_REPO))
    p_create_session.set_defaults(func=create_session)

    p_create_prompt = sub.add_parser("create-prompt")
    p_create_prompt.add_argument("--session-id")
    p_create_prompt.add_argument("--text")
    p_create_prompt.set_defaults(func=create_prompt)

    p_dispatch = sub.add_parser("dispatch")
    p_dispatch.add_argument("-p", "--prompt-ref", required=True)
    p_dispatch.set_defaults(func=dispatch_compat)

    p_close = sub.add_parser("close-session")
    p_close.add_argument("session_id", type=valid_uuid)
    p_close.set_defaults(func=close_session)

    p_list = sub.add_parser("list-sessions")
    p_list.set_defaults(func=list_sessions)

    p_read = sub.add_parser("read-output")
    p_read.add_argument("session_id", type=valid_uuid)
    p_read.add_argument("--lines", type=int, default=500)
    p_read.set_defaults(func=read_output)

    p_send = sub.add_parser("send")
    p_send.add_argument("--session-id", required=True, type=valid_uuid)
    p_send.add_argument("--text")
    p_send.set_defaults(func=send_cmd)

    p_watch = sub.add_parser("watch")
    p_watch.add_argument("session_id", type=valid_uuid)
    p_watch.add_argument("--timeout", type=int, default=300)
    p_watch.add_argument("--interval", type=float, default=2.0)
    p_watch.add_argument("--quiet", action="store_true")
    p_watch.set_defaults(func=watch_cmd)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()