#!/usr/bin/env python3
"""
claude-sessions — Session lifecycle manager for Claude Code tmux sessions.

Commands:
  (no args)               Interactive menu
  list                    Show live Claude tmux sessions and resumable registry entries
  recover [--days N]      Detect sessions active before a restart
  spawn <project>         Spawn a new session (wraps claude-tmux)
  resume <session-id>     Resume a past session by conversation ID
  register                Record a new session in the registry (called by claude-tmux)
  gc [--days N]           Prune stale entries and backfill session IDs
"""

import argparse
import json
import os
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

# === Terminal colours ===
RED = "\033[0;31m"
GREEN = "\033[0;32m"
YELLOW = "\033[0;33m"
CYAN = "\033[0;36m"
BOLD = "\033[1m"
DIM = "\033[2m"
NC = "\033[0m"


# === Paths ===
HOME = Path.home()
CLAUDE_ROOT = HOME / "dev" / "claude"
CONFIG_FILE = HOME / ".config" / "claude" / "profiles.yaml"
REGISTRY_FILE = HOME / ".local" / "share" / "claude-sessions" / "registry.json"
HOSTNAME = subprocess.run(["hostname", "-s"], capture_output=True, text=True).stdout.strip()

PATH_SCAN_ROOTS = [
    (HOME / "dev" / "claude", 2),
    (HOME / "dev" / "claude" / "octopus", 1),
    (HOME / "dev" / "claude" / "gamedev", 1),
    (HOME / "dev" / "claude" / "worktrees", 1),
    (HOME, 1),
]

DEFAULT_PROFILES = {
    "oreillyit-anthropic": {
        "display_name": "O'Reilly IT",
        "claude_dir": str(HOME / ".claude-oreillyit-anthropic"),
        "default_mode": "quick",
    },
    "default": {
        "display_name": "Personal",
        "claude_dir": str(HOME / ".claude"),
        "default_mode": "quick",
    },
}


# ── Helpers ──────────────────────────────────────────────────────────────────

def warn(msg: str) -> None:
    print(f"{YELLOW}Warning:{NC} {msg}", file=sys.stderr)

def die(msg: str) -> None:
    print(f"{RED}Error:{NC} {msg}", file=sys.stderr)
    sys.exit(1)


# ── Profiles ────────────────────────────────────────────────────────────────

def load_profiles() -> dict:
    if CONFIG_FILE.exists():
        try:
            import yaml
            with CONFIG_FILE.open() as f:
                config = yaml.safe_load(f)
            if config and "profiles" in config:
                profiles = config["profiles"]
                for p in profiles.values():
                    if "claude_dir" in p:
                        p["claude_dir"] = str(Path(p["claude_dir"]).expanduser())
                return profiles
        except Exception as e:
            warn(f"Could not load {CONFIG_FILE}: {e}")
    return DEFAULT_PROFILES


def all_profile_dirs(profiles: dict) -> list[tuple[str, Path]]:
    result = []
    for name, cfg in profiles.items():
        d = Path(cfg.get("claude_dir", "")).expanduser()
        if d.is_dir():
            result.append((name, d))
    for d in HOME.iterdir():
        if d.is_dir() and d.name.startswith(".claude-") and not d.name.endswith((".json", ".backup")):
            pname = d.name[len(".claude-"):]
            if not any(n == pname for n, _ in result):
                result.append((pname, d))
    return result


# ── Registry ────────────────────────────────────────────────────────────────

def load_registry() -> dict:
    if REGISTRY_FILE.exists():
        try:
            return json.loads(REGISTRY_FILE.read_text())
        except Exception:
            pass
    return {"sessions": []}


def save_registry(registry: dict) -> None:
    REGISTRY_FILE.parent.mkdir(parents=True, exist_ok=True)
    REGISTRY_FILE.write_text(json.dumps(registry, indent=2))


# ── Path hash utilities ──────────────────────────────────────────────────────

def path_to_hash(path: str | Path) -> str:
    return str(path).replace("/", "-")


def build_path_map() -> dict[str, str]:
    result: dict[str, str] = {}

    def scan(p: Path, depth: int) -> None:
        if not p.is_dir() or depth < 0:
            return
        result[path_to_hash(p)] = str(p)
        if depth > 0:
            try:
                for child in p.iterdir():
                    if child.is_dir() and not child.name.startswith("."):
                        scan(child, depth - 1)
            except PermissionError:
                pass

    for root, depth in PATH_SCAN_ROOTS:
        scan(root, depth)
    return result


# ── tmux utilities ───────────────────────────────────────────────────────────

def list_tmux_sessions() -> set[str]:
    try:
        r = subprocess.run(["tmux", "ls", "-F", "#{session_name}"], capture_output=True, text=True)
        if r.returncode == 0:
            return set(r.stdout.strip().splitlines())
    except Exception:
        pass
    return set()


def title_case(s: str) -> str:
    return " ".join(w.capitalize() for w in s.replace("-", " ").replace("_", " ").split())


# ── Session file parsing ─────────────────────────────────────────────────────

def read_jsonl_header(path: Path, max_lines: int = 30) -> dict:
    session_id = None
    start_time = None
    try:
        with path.open(errors="replace") as f:
            for i, line in enumerate(f):
                if i >= max_lines:
                    break
                try:
                    obj = json.loads(line.strip())
                    if not session_id and "sessionId" in obj:
                        session_id = obj["sessionId"]
                    if not start_time and "timestamp" in obj:
                        start_time = obj["timestamp"]
                    if session_id and start_time:
                        break
                except json.JSONDecodeError:
                    pass
    except Exception:
        pass
    return {"session_id": session_id or path.stem, "start_time": start_time}


def scan_profile_sessions(profile_name: str, claude_dir: Path,
                           path_map: dict[str, str], cutoff: float) -> list[dict]:
    projects_dir = claude_dir / "projects"
    if not projects_dir.is_dir():
        return []

    records = []
    for proj_dir in projects_dir.iterdir():
        if not proj_dir.is_dir():
            continue
        project_hash = proj_dir.name
        real_path = path_map.get(project_hash)

        jsonl_files = sorted(
            proj_dir.glob("*.jsonl"),
            key=lambda f: f.stat().st_mtime,
            reverse=True,
        )
        for jf in jsonl_files:
            mtime = jf.stat().st_mtime
            if mtime < cutoff:
                break
            meta = read_jsonl_header(jf)
            records.append({
                "session_id": meta["session_id"],
                "profile": profile_name,
                "project_hash": project_hash,
                "project_dir": real_path,
                "start_time": meta["start_time"],
                "last_activity": datetime.fromtimestamp(mtime).strftime("%b %d %H:%M"),
                "mtime": mtime,
            })
    return records


# ── Registry backfill ────────────────────────────────────────────────────────

def backfill_session_ids(registry: dict, profiles: dict) -> int:
    """For registry entries with no session_id, find the matching .jsonl by scan."""
    updated = 0
    for entry in registry.get("sessions", []):
        if entry.get("session_id"):
            continue
        profile_name = entry.get("profile", "")
        profile_cfg = profiles.get(profile_name, {})
        claude_dir = Path(profile_cfg.get("claude_dir", "")).expanduser()
        if not claude_dir.is_dir():
            continue
        project_dir = entry.get("project_dir", "")
        if not project_dir:
            continue
        proj_hash = path_to_hash(project_dir)
        proj_path = claude_dir / "projects" / proj_hash
        if not proj_path.is_dir():
            continue

        spawned_at = entry.get("spawned_at", "")
        spawn_ts = 0.0
        if spawned_at:
            try:
                spawn_ts = datetime.fromisoformat(spawned_at).timestamp()
            except ValueError:
                pass

        best: Optional[Path] = None
        best_mtime = 0.0
        for jf in proj_path.glob("*.jsonl"):
            mtime = jf.stat().st_mtime
            if mtime > spawn_ts and mtime > best_mtime:
                best, best_mtime = jf, mtime

        if best:
            meta = read_jsonl_header(best)
            entry["session_id"] = meta["session_id"]
            entry["last_activity"] = datetime.fromtimestamp(best_mtime).isoformat()
            updated += 1
    return updated


# ── Spawning ─────────────────────────────────────────────────────────────────

def resolve_claude_profile_bin() -> str:
    r = subprocess.run(["which", "claude-profile"], capture_output=True, text=True)
    if r.returncode == 0 and r.stdout.strip():
        return r.stdout.strip()
    die("claude-profile not found in PATH")


def spawn_session(project_dir: str, profile: str, mode: str,
                  resume_id: Optional[str] = None,
                  name_override: Optional[str] = None,
                  attach: bool = False) -> bool:
    proj_path = Path(project_dir)
    if not proj_path.is_dir():
        warn(f"Project directory not found: {project_dir}")
        return False

    project_basename = proj_path.name
    base_name = name_override or title_case(project_basename)
    if resume_id:
        base_name = f"{base_name} (resume)"

    live = list_tmux_sessions()
    name = base_name
    n = 2
    while name in live:
        name = f"{base_name} #{n}"
        n += 1

    rc_flag = f'--remote-control "{name}"'
    resume_flag = f"--resume {resume_id}" if resume_id else ""
    extra = " ".join(filter(None, [rc_flag, resume_flag]))

    claude_profile_bin = resolve_claude_profile_bin()
    launch_cmd = (
        f"{claude_profile_bin} {profile} --mode {mode} "
        f"--project {project_basename} -- {extra}"
    )

    r = subprocess.run(
        ["tmux", "new-session", "-d", "-s", name, "-c", str(proj_path)],
        capture_output=True,
    )
    if r.returncode != 0:
        warn(f"tmux new-session failed: {r.stderr.decode().strip()}")
        return False

    subprocess.run(["tmux", "send-keys", "-t", name, "-l", launch_cmd])
    subprocess.run(["tmux", "send-keys", "-t", name, "Enter"])

    print(f"{GREEN}Spawned{NC} {BOLD}{name}{NC}")
    print(f"  dir:     {project_dir}")
    print(f"  profile: {profile} / mode: {mode}")
    if resume_id:
        print(f"  resume:  {resume_id[:8]}…")
    print(f"  attach:  tmux attach -t \"{name}\"")

    if attach:
        if os.environ.get("TMUX"):
            subprocess.run(["tmux", "switch-client", "-t", f"={name}"])
        else:
            os.execvp("tmux", ["tmux", "attach", "-t", f"={name}"])

    return True


def find_claude_md_dir(start: Path) -> Optional[Path]:
    """Nearest directory at or above start (stopping at $HOME) containing a CLAUDE.md."""
    for d in [start, *start.parents]:
        if d == HOME.parent or d == Path("/"):
            break
        if (d / "CLAUDE.md").is_file():
            return d
        if d == HOME:
            break
    return None


# ── Commands ─────────────────────────────────────────────────────────────────

def cmd_list(args) -> None:
    live = list_tmux_sessions()

    if live:
        print(f"\n{BOLD}Live tmux sessions:{NC}")
        for s in sorted(live):
            print(f"  {GREEN}●{NC}  {s}")
    else:
        print(f"\n{DIM}No active tmux sessions.{NC}")

    registry = load_registry()
    resumable = [s for s in registry.get("sessions", []) if s.get("tmux_name") not in live]

    if resumable:
        print(f"\n{BOLD}Resumable (registry, not currently live):{NC}")
        fmt = "  {:<4} {:<32} {:<24} {:<10} {}"
        print(fmt.format("", "Project", "Profile", "Session ID", "Last active"))
        print(fmt.format("", "-" * 32, "-" * 24, "-" * 10, "-" * 15))
        for i, s in enumerate(resumable, 1):
            proj = Path(s.get("project_dir", "")).name or s.get("project", "?")
            sid = (s.get("session_id") or "")[:8] or "pending…"
            ts = (s.get("last_activity") or s.get("spawned_at") or "?")[:16]
            profile = s.get("profile", "?")
            print(fmt.format(f"{i})", proj, profile, sid, ts))
    else:
        print(f"\n{DIM}Registry is empty — run 'spawn' to start tracking sessions.{NC}")

    print()


def cmd_recover(args) -> None:
    profiles = load_profiles()
    path_map = build_path_map()
    cutoff = time.time() - args.days * 86400

    all_sessions: list[dict] = []
    for profile_name, claude_dir in all_profile_dirs(profiles):
        all_sessions.extend(scan_profile_sessions(profile_name, claude_dir, path_map, cutoff))

    all_sessions.sort(key=lambda s: s["mtime"], reverse=True)

    if not all_sessions:
        print(f"No sessions found in the last {args.days} days.")
        return

    print(f"\n{BOLD}Sessions active in the last {args.days} day(s):{NC}\n")
    print(f"  {'#':<4} {'Project':<32} {'Profile':<24} {'Session ID':<10} {'Last Active':<16} Note")
    print(f"  {'-'*4} {'-'*32} {'-'*24} {'-'*10} {'-'*16} {'-'*20}")

    for i, s in enumerate(all_sessions, 1):
        if s["project_dir"]:
            proj_raw = Path(s["project_dir"]).name
            proj = proj_raw[:30] + ("…" if len(proj_raw) > 30 else "")
            dim_proj = ""
        else:
            h = s["project_hash"]
            for prefix in ("-home-paul-", "-root-", "-home-"):
                if h.startswith(prefix):
                    h = h[len(prefix):]
                    break
            parts = h.rsplit("-", 4)
            hint = "/".join(parts[-2:]) if len(parts) >= 2 else h
            proj = hint[:30]
            dim_proj = DIM

        sid = s["session_id"][:8]
        ts = s["last_activity"]
        profile_raw = s["profile"]
        profile = profile_raw[:22] + ("…" if len(profile_raw) > 22 else "")
        note = "" if s["project_dir"] else f"{YELLOW}path?{NC}"
        proj_padded = f"{proj:<32}"
        print(f"  {i:<4} {dim_proj}{proj_padded}{NC if dim_proj else ''} {profile:<24} {sid:<10} {ts:<16} {note}")

    print()

    if not sys.stdin.isatty():
        return

    try:
        raw = input("Resume which sessions? (e.g. 1,3 or 'all', Enter to skip): ").strip()
    except (KeyboardInterrupt, EOFError):
        print()
        return

    if not raw:
        return

    if raw.lower() == "all":
        selected = [s for s in all_sessions if s["project_dir"]]
    else:
        try:
            indices = [int(x.strip()) - 1 for x in raw.replace(" ", "").split(",")]
            selected = [all_sessions[i] for i in indices if 0 <= i < len(all_sessions)]
        except ValueError:
            die("Invalid selection — expected numbers like 1,3,5 or 'all'")
            return

    if not selected:
        return

    profiles_cfg = load_profiles()
    print()
    # Land in the session only when resuming exactly one; several stay detached.
    attach = len(selected) == 1
    for s in selected:
        if not s["project_dir"]:
            warn(f"Skipping: project directory not found for {s['project_hash']}")
            continue
        profile_cfg = profiles_cfg.get(s["profile"], {})
        mode = getattr(args, "mode", None) or profile_cfg.get("default_mode", "quick")
        spawn_session(project_dir=s["project_dir"], profile=s["profile"],
                      mode=mode, resume_id=s["session_id"], attach=attach)
        print()


def cmd_spawn(args) -> None:
    claude_tmux = subprocess.run(["which", "claude-tmux"], capture_output=True, text=True).stdout.strip()
    if not claude_tmux:
        die("claude-tmux not found in PATH")
    cmd = [claude_tmux, args.project]
    if args.mode:
        cmd += ["--mode", args.mode]
    if args.profile:
        cmd += ["--profile", args.profile]
    if args.name:
        cmd += ["--name", args.name]
    if args.attach:
        cmd += ["--attach"]
    os.execvp(claude_tmux, cmd)


def cmd_resume(args) -> None:
    session_id = args.session_id
    registry = load_registry()
    entry = next(
        (s for s in registry.get("sessions", [])
         if (s.get("session_id") or "").startswith(session_id)),
        None,
    )

    if entry:
        project_dir = entry.get("project_dir", "")
        profile = args.profile or entry.get("profile", "oreillyit-anthropic")
        mode = args.mode or entry.get("mode", "quick")
    else:
        if not args.project:
            die(f"Session {session_id!r} not in registry.\nProvide --project (and optionally --profile, --mode).")
        profile = args.profile or "oreillyit-anthropic"
        profiles_cfg = load_profiles()
        mode = args.mode or profiles_cfg.get(profile, {}).get("default_mode", "quick")
        path_map = build_path_map()
        project_dir = next((p for p in path_map.values() if Path(p).name == args.project), None)
        if not project_dir:
            die(f"Project directory not found for '{args.project}'")

    spawn_session(project_dir=project_dir, profile=profile, mode=mode,
                  resume_id=session_id, attach=args.attach)


def cmd_register(args) -> None:
    """Called by claude-tmux after spawning a session."""
    registry = load_registry()
    entry = {
        "tmux_name": args.tmux_name,
        "project": Path(args.project_dir).name,
        "project_dir": args.project_dir,
        "profile": args.profile,
        "mode": args.mode,
        "machine": HOSTNAME,
        "spawned_at": datetime.now(timezone.utc).isoformat(),
        "session_id": None,
        "last_activity": None,
    }
    # Remove any prior entry for this tmux session name
    registry["sessions"] = [
        s for s in registry.get("sessions", [])
        if s.get("tmux_name") != args.tmux_name
    ]
    registry["sessions"].append(entry)
    save_registry(registry)


def cmd_gc(args) -> None:
    registry = load_registry()
    profiles = load_profiles()
    live = list_tmux_sessions()
    cutoff = time.time() - args.days * 86400

    # Backfill session IDs first
    backfilled = backfill_session_ids(registry, profiles)
    if backfilled:
        print(f"Backfilled {backfilled} session ID(s).")

    before = len(registry.get("sessions", []))

    def keep(s: dict) -> bool:
        if s.get("tmux_name") in live:
            return True
        spawned = s.get("spawned_at", "")
        if spawned:
            try:
                ts = datetime.fromisoformat(spawned).timestamp()
                return ts > cutoff
            except ValueError:
                pass
        return False

    registry["sessions"] = [s for s in registry.get("sessions", []) if keep(s)]
    after = len(registry["sessions"])
    save_registry(registry)
    print(f"Pruned {before - after} stale entries ({after} remaining).")


# ── Interactive menu ─────────────────────────────────────────────────────────

def cmd_menu(_args=None) -> None:
    """Interactive session picker. Shown when claude-sessions is invoked with no args."""
    profiles_cfg = load_profiles()

    while True:
        live = sorted(list_tmux_sessions())
        registry = load_registry()
        resumable = [
            s for s in registry.get("sessions", [])
            if s.get("tmux_name") not in live and s.get("session_id")
        ]

        # Build numbered options
        options: list[tuple[str, dict]] = []  # (kind, data)
        for name in live:
            options.append(("live", {"tmux_name": name}))
        for entry in resumable:
            options.append(("resumable", entry))

        # ── Print menu ──
        os.system("clear 2>/dev/null || true")
        print(f"\n{BOLD}Claude Sessions{NC}")
        print("─" * 50)

        if not options:
            print(f"\n  {DIM}No sessions.{NC}")
        else:
            if live:
                print(f"\n  {BOLD}Live:{NC}")
                for i, (kind, data) in enumerate(options, 1):
                    if kind != "live":
                        continue
                    idx = options.index(("live", data)) + 1
                    print(f"    {CYAN}{idx:>2}){NC}  {GREEN}●{NC}  {data['tmux_name']}")

            if resumable:
                print(f"\n  {BOLD}Resumable:{NC}")
                for i, (kind, data) in enumerate(options, 1):
                    if kind != "resumable":
                        continue
                    proj = Path(data.get("project_dir", "")).name or data.get("project", "?")
                    sid = (data.get("session_id") or "")[:8]
                    ts = (data.get("last_activity") or data.get("spawned_at") or "?")[:16]
                    profile = data.get("profile", "?")
                    print(f"    {CYAN}{i:>2}){NC}  {DIM}○{NC}  {proj:<28} {DIM}{profile} · {sid} · {ts}{NC}")

        print(f"\n  {CYAN} r){NC}  Recover sessions from restart  (--days 7)")
        print(f"  {CYAN} n){NC}  New session")
        print(f"  {CYAN} q){NC}  Quit")
        print()

        try:
            raw = input("  Choice: ").strip().lower()
        except (KeyboardInterrupt, EOFError):
            print()
            break

        if raw == "q":
            break

        elif raw == "r":
            import types
            a = types.SimpleNamespace(days=7, mode=None, attach=False)
            cmd_recover(a)
            input("\nPress Enter to return to menu…")

        elif raw == "n":
            default_dir = find_claude_md_dir(Path.cwd())
            try:
                if default_dir:
                    print(f"  Found CLAUDE.md in {BOLD}{default_dir}{NC}")
                    project = input(
                        f"  Project [{BOLD}{default_dir.name}{NC}] (Enter = use this, or type another): "
                    ).strip()
                    if not project:
                        project = str(default_dir)
                else:
                    project = input("  Project name: ").strip()
            except (KeyboardInterrupt, EOFError):
                print()
                continue
            if not project:
                continue
            import types
            a = types.SimpleNamespace(project=project, mode=None, profile=None, name=None, attach=True)
            print()
            cmd_spawn(a)  # exec's claude-tmux --attach; replaces this process

        elif raw.isdigit():
            idx = int(raw) - 1
            if idx < 0 or idx >= len(options):
                print(f"  {YELLOW}Invalid choice{NC}")
                time.sleep(1)
                continue
            kind, data = options[idx]

            if kind == "live":
                name = data["tmux_name"]
                print(f"\n  Attaching to {BOLD}{name}{NC}…")
                if os.environ.get("TMUX"):
                    subprocess.run(["tmux", "switch-client", "-t", f"={name}"])
                else:
                    os.execvp("tmux", ["tmux", "attach", "-t", f"={name}"])
                break  # exec'd — this line only reached on switch-client

            elif kind == "resumable":
                project_dir = data.get("project_dir", "")
                profile = data.get("profile", "oreillyit-anthropic")
                mode = profiles_cfg.get(profile, {}).get("default_mode", "quick")
                session_id = data.get("session_id", "")
                if not project_dir:
                    print(f"  {YELLOW}Project directory unknown — cannot resume{NC}")
                    time.sleep(1)
                    continue
                print()
                spawn_session(project_dir=project_dir, profile=profile,
                              mode=mode, resume_id=session_id, attach=True)
        else:
            print(f"  {YELLOW}Unknown choice: {raw!r}{NC}")
            time.sleep(1)


# ── Entrypoint ───────────────────────────────────────────────────────────────

def main() -> None:
    p = argparse.ArgumentParser(
        prog="claude-sessions",
        description="Session lifecycle manager for Claude Code tmux sessions",
        add_help=True,
    )
    sub = p.add_subparsers(dest="command", metavar="COMMAND")

    sub.add_parser("list", help="Show live and resumable sessions")
    sub.add_parser("menu", help="Interactive session menu (default when no args)")

    rp = sub.add_parser("recover", help="Detect sessions active before a restart")
    rp.add_argument("--days", type=int, default=7, help="Look back N days (default: 7)")
    rp.add_argument("--mode", default=None)
    rp.add_argument("--attach", action="store_true")

    sp = sub.add_parser("spawn", help="Spawn a new session (wraps claude-tmux)")
    sp.add_argument("project", help="Project name or path")
    sp.add_argument("--mode", default=None)
    sp.add_argument("--profile", default=None)
    sp.add_argument("--name", default=None)
    sp.add_argument("--attach", action="store_true")

    rsp = sub.add_parser("resume", help="Resume a past conversation by session ID")
    rsp.add_argument("session_id", help="Session UUID or prefix")
    rsp.add_argument("--project", default=None)
    rsp.add_argument("--profile", default=None)
    rsp.add_argument("--mode", default=None)
    rsp.add_argument("--attach", action="store_true")

    reg = sub.add_parser("register", help="Record a new session (called by claude-tmux)")
    reg.add_argument("--project-dir", required=True)
    reg.add_argument("--profile", required=True)
    reg.add_argument("--mode", required=True)
    reg.add_argument("--tmux-name", required=True)

    gp = sub.add_parser("gc", help="Prune stale entries and backfill session IDs")
    gp.add_argument("--days", type=int, default=30)

    args = p.parse_args()

    # No args + interactive TTY → menu
    if not args.command:
        if sys.stdin.isatty() and sys.stdout.isatty():
            cmd_menu()
        else:
            p.print_help()
        sys.exit(0)

    {
        "list": cmd_list,
        "menu": cmd_menu,
        "recover": cmd_recover,
        "spawn": cmd_spawn,
        "resume": cmd_resume,
        "register": cmd_register,
        "gc": cmd_gc,
    }[args.command](args)


if __name__ == "__main__":
    main()
