#!/usr/bin/env python3
"""muse_resume_pool: per-repo, profile-aware Muse Code resume pool.

Why this exists
---------------
``muse resume`` scopes its picker by workspace but is blind to muse-auth
profiles, and the TUI ``/resume`` reads session logs itself and lumps every
workspace into one heap. A session resumed under a different credential
than the one that created it fails server-side: the continuation is
cryptographically bound to the creating account, so the server rejects the
resume. This tool lists only the sessions that can actually resume here
and now, and guards ``resume`` calls before they fail.

Data sources (read-only, no secrets)
------------------------------------
- ``~/.local/share/muse/session-index.db`` ``sessions`` table (``mode=ro``).
  Fallback when the index is missing: scan ``sessions/*/*/*/*/session.jsonl``
  for ``runtime.session.metadata`` (workspace_root) and
  ``session.name.changed`` (session_name) records, mirroring ``/resume``.
- ``~/.config/muse/active_profile``, ``session_profiles.json``,
  ``switch_history.jsonl``: profile *names* only. ``auth.json`` token bytes
  are never read, logged, or compared.

Profile resolution mirrors ``muse-auth``: cached ``session_profiles.json``
mapping wins; otherwise the latest switch at or before session start; else
the earliest switch; else the active profile as fallback. When the auth
dir is unreadable (e.g. inside the Muse sandbox, which masks it), the
profile is ``unknown`` and only proven mismatches are hidden/blocked.
"""

import argparse
import glob
import json
import os
import sqlite3
import subprocess
import sys

INDEX_COLUMNS = (
    "session_id",
    "session_name",
    "workspace_root",
    "workspace_key",
    "provider_id",
    "model_id",
    "git_branch",
    "title",
    "first_user_prompt",
    "created_at_us",
    "updated_at_us",
    "prompt_count",
    "status",
)


def default_paths():
    home = os.path.expanduser("~")
    data_home = os.environ.get("XDG_DATA_HOME", os.path.join(home, ".local", "share"))
    return {
        "index_db": os.path.join(data_home, "muse", "session-index.db"),
        "sessions_dir": os.path.join(data_home, "muse", "sessions"),
        "config_dir": os.path.expanduser("~/.config/muse"),
    }


def canonical_workspace(cwd=None):
    """Repo root for the pool: git top-level, else real cwd."""
    cwd = cwd or os.getcwd()
    try:
        out = subprocess.run(
            ["git", "-C", cwd, "rev-parse", "--show-toplevel"],
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            text=True,
            timeout=10,
        )
        if out.returncode == 0 and out.stdout.strip():
            return os.path.realpath(out.stdout.strip())
    except Exception:
        pass
    return os.path.realpath(cwd)


def load_index_rows(index_db):
    """Read session rows from the index (read-only). None if unavailable."""
    if not os.path.exists(index_db):
        return None
    cols = ", ".join(INDEX_COLUMNS)
    try:
        uri = "file:{}?mode=ro".format(index_db.replace("?", "%3F"))
        conn = sqlite3.connect(uri, uri=True, timeout=5)
        try:
            conn.row_factory = sqlite3.Row
            cur = conn.execute(
                "SELECT {} FROM sessions ORDER BY "
                "updated_at_us DESC, created_at_us DESC, session_id ASC".format(cols)
            )
            return [dict(r) for r in cur.fetchall()]
        finally:
            conn.close()
    except sqlite3.Error:
        return None


def _scan_log_for_session(session_log):
    """Extract workspace/name/title from one session.jsonl (bounded read)."""
    workspace = None
    name = None
    title = None
    first_prompt = None
    created_at_us = None
    updated_at_us = None
    prompt_count = 0
    try:
        with open(session_log, "r", errors="replace") as fh:
            for line in fh:
                line = line.strip()
                if not line:
                    continue
                try:
                    rec = json.loads(line)
                except ValueError:
                    continue
                if "children" in rec:  # retained permission frame wrapper
                    continue
                rec_at = rec.get("recorded_at")
                if isinstance(rec_at, int):
                    if created_at_us is None:
                        created_at_us = rec_at
                    updated_at_us = rec_at
                ptype = rec.get("payload_type", "")
                payload = rec.get("payload", {}) if isinstance(rec.get("payload"), dict) else {}
                if ptype == "runtime.session.metadata":
                    record = payload.get("record", {})
                    workspace = workspace or record.get("workspace_root")
                elif ptype == "session.name.changed":
                    if payload.get("new_name"):
                        name = payload["new_name"]
                elif ptype == "runtime.session":
                    event = payload.get("event", {})
                    if event.get("kind") == "started" and not first_prompt:
                        prompt = event.get("prompt") or ""
                        first_prompt = prompt[:200]
                        title = title or prompt[:80]
                        prompt_count += 1
    except OSError:
        return None
    if workspace is None and name is None and created_at_us is None:
        return None
    session_id = os.path.basename(os.path.dirname(session_log))
    return {
        "session_id": session_id,
        "session_name": name,
        "workspace_root": workspace,
        "workspace_key": workspace,
        "provider_id": None,
        "model_id": None,
        "git_branch": None,
        "title": title or "New session",
        "first_user_prompt": first_prompt,
        "created_at_us": created_at_us,
        "updated_at_us": updated_at_us or created_at_us,
        "prompt_count": prompt_count,
        "status": "valid",
    }


def scan_session_logs(sessions_dir):
    """Fallback pool source: scan top-level session.jsonl files directly."""
    pattern = os.path.join(sessions_dir, "*", "*", "*", "*", "session.jsonl")
    rows = []
    for path in glob.glob(pattern):
        row = _scan_log_for_session(path)
        if row:
            rows.append(row)
    rows.sort(
        key=lambda r: (
            r.get("updated_at_us") or 0,
            r.get("created_at_us") or 0,
            r.get("session_id") or "",
        ),
        reverse=True,
    )
    return rows


def load_auth_state(config_dir):
    """Load profile names only. Never touches auth.json token bytes."""
    state = {
        "readable": False,
        "active": None,
        "session_profiles": {},
        "switch_history": [],
    }
    if not os.path.isdir(config_dir):
        return state
    if not (os.access(config_dir, os.R_OK) and os.access(config_dir, os.X_OK)):
        return state
    state["readable"] = True
    try:
        with open(os.path.join(config_dir, "active_profile"), "r") as fh:
            state["active"] = fh.read().strip() or None
    except OSError:
        pass
    try:
        with open(os.path.join(config_dir, "session_profiles.json"), "r") as fh:
            data = json.load(fh)
            if isinstance(data, dict):
                state["session_profiles"] = {
                    str(k): str(v) for k, v in data.items()
                }
    except (OSError, ValueError):
        pass
    try:
        history = []
        with open(os.path.join(config_dir, "switch_history.jsonl"), "r") as fh:
            for line in fh:
                line = line.strip()
                if not line:
                    continue
                try:
                    entry = json.loads(line)
                except ValueError:
                    continue
                if entry.get("profile"):
                    history.append(entry)
        history.sort(key=lambda e: e.get("epoch", 0))
        state["switch_history"] = history
    except OSError:
        pass
    return state


def resolve_profile(session_id, created_at_us, auth_state):
    """Return (profile_or_None, source). Mirrors muse-auth resolution order."""
    if not auth_state.get("readable"):
        return None, "unknown"
    cached = auth_state.get("session_profiles", {})
    if session_id in cached:
        return cached[session_id], "cached"
    history = auth_state.get("switch_history", [])
    start_epoch = (created_at_us / 1e6) if created_at_us else None
    if history and start_epoch:
        for switch in reversed(history):
            if switch.get("epoch", 0) <= start_epoch:
                return switch.get("profile"), "history"
        return history[0].get("profile"), "history"
    if auth_state.get("active"):
        return auth_state["active"], "fallback"
    return None, "unknown"


def annotate_rows(rows, auth_state):
    """Attach profile + source to each row (mutates and returns rows)."""
    for row in rows:
        profile, source = resolve_profile(
            row.get("session_id"), row.get("created_at_us"), auth_state
        )
        row["auth_profile"] = profile
        row["auth_source"] = source
    return rows


def pool_for_workspace(rows, workspace):
    """Exact workspace_key match (workspace_root fallback), index order kept."""
    pool = []
    for row in rows:
        key = row.get("workspace_key") or row.get("workspace_root")
        if key and os.path.realpath(key) == workspace:
            pool.append(row)
    return pool


def split_resumable(pool, active_profile):
    """(resumable, blocked): only proven profile mismatches are blocked.

    Unknown profiles (sandboxed auth dir, no mapping/history) stay resumable
    but flagged, since blocking them would hide possibly valid sessions.
    """
    resumable, blocked = [], []
    for row in pool:
        profile = row.get("auth_profile")
        if active_profile and profile and profile != active_profile:
            blocked.append(row)
        else:
            resumable.append(row)
    return resumable, blocked


def resolve_ref(rows, ref):
    """Resolve UUID / UUID prefix / session name. Returns (matches, kind)."""
    ref = (ref or "").strip()
    if not ref:
        return [], "empty"
    exact = [r for r in rows if r.get("session_id") == ref]
    if exact:
        return exact, "uuid"
    named = [r for r in rows if (r.get("session_name") or "") == ref]
    if named:
        return named, "name"
    if len(ref) >= 8:
        prefixed = [
            r for r in rows if (r.get("session_id") or "").startswith(ref)
        ]
        if prefixed:
            return prefixed, "prefix"
    return [], "none"


def check_resume(rows, ref, workspace, auth_state):
    """Guard decision for resuming ``ref`` from ``workspace``.

    Returns dict(ok=bool, reason=str, detail=str, fix=str, row=row|None).
    """
    matches, kind = resolve_ref(rows, ref)
    if kind == "empty" or not matches:
        return {
            "ok": False,
            "reason": "unknown-session",
            "detail": "No session matches '{}'.".format(ref),
            "fix": "List this repo's pool: muse_resume_pool.py pool",
            "row": None,
        }
    if len(matches) > 1:
        ids = ", ".join(m["session_id"][:12] for m in matches[:5])
        return {
            "ok": False,
            "reason": "ambiguous",
            "detail": "'{}' matches {} sessions: {}".format(ref, len(matches), ids),
            "fix": "Use a longer UUID prefix or the full session id.",
            "row": None,
        }
    row = matches[0]
    key = row.get("workspace_key") or row.get("workspace_root")
    if not key or os.path.realpath(key) != workspace:
        return {
            "ok": False,
            "reason": "wrong-workspace",
            "detail": "Session '{}' belongs to workspace '{}', not '{}'.".format(
                row.get("session_name") or row["session_id"][:12], key, workspace
            ),
            "fix": "cd '{}' first, or pick a session from this repo's pool.".format(key or "?"),
            "row": row,
        }
    if row.get("status") and row["status"] != "valid":
        return {
            "ok": False,
            "reason": "bad-status",
            "detail": "Session '{}' has status '{}'.".format(
                row.get("session_name") or row["session_id"][:12], row["status"]
            ),
            "fix": "Pick a session with status 'valid' from this repo's pool.",
            "row": row,
        }
    active = auth_state.get("active")
    profile = row.get("auth_profile")
    if active and profile and profile != active:
        return {
            "ok": False,
            "reason": "wrong-profile",
            "detail": "Session '{}' was created under muse-auth profile '{}' "
            "but the active profile is '{}'; the server would reject the "
            "resume (continuation is bound to the creating credential).".format(
                row.get("session_name") or row["session_id"][:12], profile, active
            ),
            "fix": "Run `muse-auth use {}` (outside the sandbox), then resume.".format(profile),
            "row": row,
        }
    detail = "Session '{}' is resumable here.".format(
        row.get("session_name") or row["session_id"][:12]
    )
    if not auth_state.get("readable"):
        detail += " (Profile unverified: auth dir unreadable from this shell.)"
    elif not profile:
        detail += " (Profile unknown: no mapping or switch history.)"
    return {
        "ok": True,
        "reason": "ok",
        "detail": detail,
        "fix": "",
        "row": row,
    }


def load_rows(paths):
    """Index first, log-scan fallback. Returns (rows, source)."""
    rows = load_index_rows(paths["index_db"])
    if rows is not None:
        return rows, "index"
    return scan_session_logs(paths["sessions_dir"]), "log-scan"


def format_pool_table(resumable, blocked, show_all):
    lines = []
    header = "{:<18} {:<12} {:<10} {:<5} {}".format(
        "NAME", "SESSION", "PROFILE", "MSGS", "TITLE"
    )
    lines.append(header)
    for row in resumable:
        flag = "?" if not row.get("auth_profile") else " "
        lines.append(
            "{:<18} {:<12} {:<10} {:<5} {}{}".format(
                (row.get("session_name") or "-")[:18],
                (row.get("session_id") or "")[:12],
                (row.get("auth_profile") or "?")[:10],
                row.get("prompt_count", 0),
                flag,
                (row.get("title") or "")[:60],
            )
        )
    if show_all:
        for row in blocked:
            lines.append(
                "{:<18} {:<12} {:<10} {:<5} {} [BLOCKED: profile mismatch]".format(
                    (row.get("session_name") or "-")[:18],
                    (row.get("session_id") or "")[:12],
                    (row.get("auth_profile") or "?")[:10],
                    row.get("prompt_count", 0),
                    (row.get("title") or "")[:60],
                )
            )
    elif blocked:
        lines.append(
            "({} session(s) hidden: wrong muse-auth profile; use --all to show)".format(
                len(blocked)
            )
        )
    return "\n".join(lines)


def cmd_pool(args, paths):
    workspace = os.path.realpath(args.workspace or canonical_workspace())
    rows, source = load_rows(paths)
    auth_state = load_auth_state(paths["config_dir"])
    annotate_rows(rows, auth_state)
    pool = pool_for_workspace(rows, workspace)
    resumable, blocked = split_resumable(pool, auth_state.get("active"))
    if args.json:
        print(json.dumps({
            "workspace": workspace,
            "source": source,
            "active_profile": auth_state.get("active"),
            "auth_readable": auth_state.get("readable"),
            "resumable": resumable,
            "blocked": blocked if args.all else [],
            "blocked_count": len(blocked),
        }, indent=2, default=str))
        return 0
    print("workspace: {}  (source: {})".format(workspace, source))
    if auth_state.get("readable"):
        print("active profile: {}".format(auth_state.get("active") or "(none)"))
    else:
        print("active profile: ? (auth dir unreadable from this shell)")
    if not pool:
        print("No sessions for this workspace.")
        return 0
    print(format_pool_table(resumable, blocked, args.all))
    return 0


def cmd_check(args, paths):
    workspace = os.path.realpath(args.workspace or canonical_workspace())
    rows, _ = load_rows(paths)
    auth_state = load_auth_state(paths["config_dir"])
    annotate_rows(rows, auth_state)
    decision = check_resume(rows, args.ref, workspace, auth_state)
    if args.json:
        row = dict(decision["row"]) if decision["row"] else None
        print(json.dumps({
            "ok": decision["ok"],
            "reason": decision["reason"],
            "detail": decision["detail"],
            "fix": decision["fix"],
            "row": row,
        }, indent=2, default=str))
    else:
        status = "OK" if decision["ok"] else "BLOCKED ({})".format(decision["reason"])
        print("{}: {}".format(status, decision["detail"]))
        if decision["fix"]:
            print("fix: {}".format(decision["fix"]))
    return 0 if decision["ok"] else 1


def cmd_resume(args, paths):
    workspace = os.path.realpath(args.workspace or canonical_workspace())
    rows, _ = load_rows(paths)
    auth_state = load_auth_state(paths["config_dir"])
    annotate_rows(rows, auth_state)
    if args.ref == "--last":
        pool = pool_for_workspace(rows, workspace)
        resumable, _ = split_resumable(pool, auth_state.get("active"))
        if not resumable:
            print("No resumable sessions for this workspace.", file=sys.stderr)
            return 1
        session_id = resumable[0]["session_id"]
    else:
        decision = check_resume(rows, args.ref, workspace, auth_state)
        if not decision["ok"]:
            print("refusing to resume: {}".format(decision["detail"]), file=sys.stderr)
            if decision["fix"]:
                print("fix: {}".format(decision["fix"]), file=sys.stderr)
            return 1
        session_id = decision["row"]["session_id"]
    cmd = ["muse-code", "resume", session_id]
    if args.dry_run:
        print("would exec: {}".format(" ".join(cmd)))
        return 0
    os.execvp(cmd[0], cmd)
    return 0  # unreachable


def main(argv=None):
    parser = argparse.ArgumentParser(
        prog="muse_resume_pool",
        description="Per-repo, profile-aware Muse Code resume pool and guard.",
    )
    parser.add_argument(
        "--workspace",
        help="Workspace root to scope to (default: git top-level or cwd).",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    p_pool = sub.add_parser("pool", help="List sessions resumable here and now.")
    p_pool.add_argument("--all", action="store_true",
                        help="Also show profile-blocked sessions.")
    p_pool.add_argument("--json", action="store_true", help="Machine-readable output.")
    p_pool.set_defaults(func=cmd_pool)

    p_check = sub.add_parser("check", help="Explain whether a resume would succeed.")
    p_check.add_argument("ref", help="Session UUID, UUID prefix, or session name.")
    p_check.add_argument("--json", action="store_true", help="Machine-readable output.")
    p_check.set_defaults(func=cmd_check)

    p_resume = sub.add_parser("resume", help="Guard then exec muse-code resume.")
    p_resume.add_argument("ref", help="Session UUID, prefix, name, or --last.")
    p_resume.add_argument("--dry-run", action="store_true",
                          help="Print the resume command instead of exec'ing.")
    p_resume.set_defaults(func=cmd_resume)

    args = parser.parse_args(argv)
    return args.func(args, default_paths())


if __name__ == "__main__":
    sys.exit(main())