#!/usr/bin/env python3
"""Shared library for mux — import for Python scripts, CLI for bash.

Python scripts:  from muxlib import load_json, save_json, relative_time
Bash scripts:    python3 ~/.config/mux/muxlib.py project-path "$name"
"""

import json
import os
import sys
import subprocess
from datetime import datetime, timezone

# --- Paths ---

MUX_CONFIG = os.path.expanduser("~/.config/mux")
MUX_DATA = os.path.expanduser("~/.local/share/mux")
MUX_PROJECTS = os.path.join(MUX_CONFIG, "projects.json")
MUX_NARRATIONS = os.path.join(MUX_CONFIG, "seen-narrations.json")
MUX_WORKTREES = os.path.join(MUX_CONFIG, "worktrees.json")
MUX_WORKTREES_DIR = os.path.expanduser("~/.mux/worktrees")
MUX_SNAPSHOT = os.path.join(MUX_DATA, "snapshots", "latest.json")
CLAUDE_PROJECTS_DIR = os.path.expanduser("~/.claude/projects")


# --- ANSI ---

B = "\033[1m"
D = "\033[2m"
R = "\033[0m"
RED = "\033[31m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
CYAN = "\033[36m"

# --- Data I/O ---


def load_json(path):
    # Shape-coercing by design: every caller treats the result as a dict
    # (.get / key access), so a file whose top level drifted to a list or
    # scalar must degrade to {} here — not crash 20 call sites downstream
    # (the cc-registry bare-array incident took out F2/F4/F5 and the
    # picker at once). Callers that can rescue a non-dict shape should
    # use load_json_any and normalize themselves.
    data = load_json_any(path)
    return data if isinstance(data, dict) else {}


def load_json_any(path):
    path = os.path.expanduser(path)
    if not os.path.exists(path):
        return {}
    try:
        with open(path) as f:
            return json.load(f)
    except (json.JSONDecodeError, IOError):
        return {}


def save_json(path, data):
    path = os.path.expanduser(path)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w") as f:
        json.dump(data, f, indent=2)
        f.write("\n")
    os.rename(tmp, path)


def load_jsonl(path):
    path = os.path.expanduser(path)
    entries = []
    if not os.path.exists(path):
        return entries
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    entries.append(json.loads(line))
                except json.JSONDecodeError:
                    pass
    return entries


def save_jsonl(path, entries):
    path = os.path.expanduser(path)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w") as f:
        for e in entries:
            f.write(json.dumps(e) + "\n")
    os.rename(tmp, path)


# --- Time ---


def relative_time(ts_str):
    try:
        ts = datetime.fromisoformat(ts_str)
        delta = datetime.now() - ts
        mins = int(delta.total_seconds() / 60)
        if mins < 1:
            return "just now"
        if mins < 60:
            return f"{mins}m ago"
        hours = mins // 60
        if hours < 24:
            return f"{hours}h ago"
        return f"{delta.days}d ago"
    except Exception:
        return ""


def now_iso():
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


# --- Terminal ---


def getch():
    import tty
    import termios

    fd = sys.stdin.fileno()
    old = termios.tcgetattr(fd)
    try:
        tty.setraw(fd)
        return sys.stdin.read(1)
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old)


def readline_prompt(prompt):
    sys.stdout.write(prompt)
    sys.stdout.flush()
    buf = ""
    while True:
        ch = getch()
        if ch in ("\r", "\n"):
            sys.stdout.write("\n")
            return buf.strip()
        if ch in ("\x7f", "\x08"):
            if buf:
                buf = buf[:-1]
                sys.stdout.write("\b \b")
                sys.stdout.flush()
            continue
        if ch == "\x03":
            return ""
        if ch == "\x1b":
            return ""
        if 32 <= ord(ch) < 127:
            buf += ch
            sys.stdout.write(ch)
            sys.stdout.flush()


# --- Subprocess ---


def run(cmd):
    r = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    return r.stdout.strip(), r.stderr.strip(), r.returncode


def run_stdout(cmd):
    r = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    return r.stdout.strip()


# --- Tmux ---


def current_session():
    out, _, rc = run("tmux display-message -p '#{session_name}'")
    return out if rc == 0 and out else None


# --- Projects ---


CC_REGISTRY = os.path.expanduser("~/.claude/cc-registry.json")

AUTO_COLORS = [
    # #2a1a0a (muddy brown) replaced with #1c2128 (clean slate) — the
    # brown read as "dirty" on the terminal; see DX capture 2026-06-07.
    "#0d1b2a", "#0a2a1a", "#1c2128", "#1a0a2a",
    "#1a1a2e", "#2a0a0a", "#0a2a2a", "#1a2a0a",
    "#2a0a1a", "#0a1a2a", "#1a2a1a", "#2a1a1a",
]


def load_projects():
    manual = load_json(MUX_PROJECTS).get("projects", {})
    if not isinstance(manual, dict):
        manual = {}
    registry = load_json_any(CC_REGISTRY)

    # Tolerate registry shape drift: canonical is {"projects": [...]}, but a
    # bare top-level list has been observed in the wild (ad-hoc cleanup
    # rewrote the file without the wrapper, crashing every load_projects
    # caller — F2 dashboard, picker, sidebars). Degrade, never crash.
    if isinstance(registry, list):
        registry = {"projects": registry}
    if not isinstance(registry, dict) or not isinstance(registry.get("projects"), list):
        return manual

    known_paths = {p.get("path") for p in manual.values() if isinstance(p, dict)}
    used_colors = {p.get("color") for p in manual.values() if isinstance(p, dict)}
    merged = dict(manual)

    for entry in registry["projects"]:
        path = entry.get("path", "")
        if not path or path in known_paths:
            continue
        name = os.path.basename(path)
        if name in merged:
            continue
        color = next((c for c in AUTO_COLORS if c not in used_colors), "#1a1a2e")
        used_colors.add(color)
        merged[name] = {"path": path, "color": color, "auto": True}

    return merged


def project_path(name):
    p = load_projects().get(name)
    return p["path"] if p else None


def project_names():
    return sorted(load_projects().keys())


def project_color(name):
    p = load_projects().get(name, {})
    return p.get("color", "")


def project_setting(key):
    d = load_json(MUX_PROJECTS)
    return d.get("settings", {}).get(key)


def project_setting_set(key, value):
    d = load_json(MUX_PROJECTS)
    d.setdefault("settings", {})[key] = value
    save_json(MUX_PROJECTS, d)


# --- Narrations ---


def narrate_check(key):
    d = load_json(MUX_NARRATIONS)
    return "yes" if d.get(key) else "no"


def narrate_mark(key):
    d = load_json(MUX_NARRATIONS)
    if not d:
        d = {}
    d[key] = True
    save_json(MUX_NARRATIONS, d)


# --- Notes ---


def notes_path(project):
    return os.path.join(MUX_CONFIG, "notes", f"{project}.json")


def notes_active(project):
    d = load_json(notes_path(project))
    return [n for n in d.get("notes", []) if not n.get("archived")]


def notes_has(project):
    return "yes" if notes_active(project) else "no"


def notes_count(project):
    return str(len(notes_active(project)))


# --- Worktrees ---


def _worktree_init():
    if not os.path.exists(MUX_WORKTREES):
        save_json(MUX_WORKTREES, {"active": []})


def worktree_add(project, slug, branch, wt_path):
    _worktree_init()
    d = load_json(MUX_WORKTREES)
    d["active"].append({
        "project": project,
        "task_slug": slug,
        "branch": branch,
        "worktree_path": wt_path,
        "created_at": now_iso(),
        "tmux_window": slug,
    })
    save_json(MUX_WORKTREES, d)


def worktree_remove(project, slug):
    _worktree_init()
    d = load_json(MUX_WORKTREES)
    d["active"] = [
        e for e in d["active"]
        if not (e["project"] == project and e["task_slug"] == slug)
    ]
    save_json(MUX_WORKTREES, d)


def worktree_is_active(project, slug):
    _worktree_init()
    d = load_json(MUX_WORKTREES)
    matches = [
        e for e in d["active"]
        if e["project"] == project and e["task_slug"] == slug
    ]
    return "yes" if matches else "no"


def worktree_list():
    _worktree_init()
    d = load_json(MUX_WORKTREES)
    if not d["active"]:
        return "No active worktrees."
    lines = []
    for e in d["active"]:
        lines.append(
            f"  {e['project']}/{e['task_slug']}"
            f"  branch: {e['branch']}"
            f"  path: {e['worktree_path']}"
        )
    return "\n".join(lines)


def worktree_count():
    _worktree_init()
    d = load_json(MUX_WORKTREES)
    return str(len(d["active"]))


def worktree_slugs(project):
    d = load_json(MUX_WORKTREES)
    return [e["task_slug"] for e in d.get("active", []) if e["project"] == project]


def worktree_entry_by_path(wt_path):
    """The registry entry (task_slug, branch, ...) for a live worktree path,
    or None. Used by snapshot to tell a worktree window from a main-station
    one without re-deriving the slug from the path — the registry is the
    source of truth create_worktree() itself writes to."""
    _worktree_init()
    d = load_json(MUX_WORKTREES)
    for e in d.get("active", []):
        if e.get("worktree_path") == wt_path:
            return e
    return None


# --- Snapshot / restore (mux snapshot, mux restore) ---
#
# The problem a snapshot has to solve: a tmux window only knows its own cwd
# and pane state. Reopening the SAME Claude conversation after a reboot needs
# the Claude Code session id, which mux never captured at launch time (`mux
# new`/`mux resume` hand a prompt to `claude`, not the other way around) and
# has no live API to ask for. Claude Code's own transcript directory is the
# only place that id is recorded, keyed by a dashified cwd — so resolution
# works backward from cwd to id, not forward.


def _dashify_path(path):
    """The EXACT slugification Claude Code's own project-dir naming uses —
    matches worktree-cleanup.sh's `sed 's|[/.]|-|g'` byte for byte. Verified
    empirically against real ~/.claude/projects/ entries (both '/' and '.'
    become '-'; a worktree path's leading '/Users/x/.mux/...' therefore
    produces a double dash where the '.' sits next to a path separator)."""
    out = []
    for ch in path:
        out.append('-' if ch in ('/', '.') else ch)
    return ''.join(out)


def resolve_session_id(cwd):
    """The most-recently-modified transcript file under this cwd's project
    dir whose OWN recorded cwd exactly matches — not just the newest file in
    the directory, because mux worktrees share their main checkout's project
    dir via a Claude Code identity symlink (the workaround for upstream issue
    #34437), so several worktrees' transcripts can live side by side there.
    Matching on the embedded cwd field disambiguates; picking the newest
    match among worktrees started at the same address chooses the current
    session over a leftover from a since-removed one. Returns None if the
    directory doesn't exist or nothing matches — never guesses."""
    slug = _dashify_path(cwd)
    proj_dir = os.path.join(CLAUDE_PROJECTS_DIR, slug)
    if not os.path.isdir(proj_dir):
        return None
    try:
        files = [f for f in os.listdir(proj_dir) if f.endswith('.jsonl')]
    except OSError:
        return None
    files.sort(key=lambda f: os.path.getmtime(os.path.join(proj_dir, f)), reverse=True)
    for fname in files:
        fpath = os.path.join(proj_dir, fname)
        try:
            with open(fpath, 'r') as fh:
                for line in fh:
                    if '"cwd"' not in line:
                        continue
                    try:
                        rec = json.loads(line)
                    except (json.JSONDecodeError, ValueError):
                        continue
                    if rec.get('cwd') == cwd:
                        return fname[:-len('.jsonl')]
        except OSError:
            continue
    return None


def snapshot_build(rows):
    """rows: iterable of (desk, window_name, window_index, cwd) tuples for
    windows already confirmed live-Claude by the bash caller (pane_is_live_claude
    — snapshot has no business capturing a dead shell). Returns the list of
    entry dicts snapshot-build prints as JSON."""
    entries = []
    for desk, window_name, window_index, cwd in rows:
        entry = {
            'desk': desk,
            'window_name': window_name,
            'window_index': window_index,
            'cwd': cwd,
            'is_worktree': False,
            'task_slug': None,
            'branch': None,
        }
        if cwd.startswith(MUX_WORKTREES_DIR + os.sep):
            wt = worktree_entry_by_path(cwd)
            if wt:
                entry['is_worktree'] = True
                entry['task_slug'] = wt.get('task_slug')
                entry['branch'] = wt.get('branch')
            else:
                # A worktree-shaped path with no registry entry is a stale
                # or hand-made worktree — still snapshot the cwd (restore can
                # still cd there and resume) but don't claim registry-backed
                # worktree status we can't actually verify.
                entry['is_worktree'] = True
        entry['session_id'] = resolve_session_id(cwd)
        entries.append(entry)
    return entries


def snapshot_write(rows, path=MUX_SNAPSHOT):
    entries = snapshot_build(rows)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    save_json(path, {'saved_at': now_iso(), 'windows': entries})
    return entries


def snapshot_read(path=MUX_SNAPSHOT):
    d = load_json(path)
    return d.get('windows', [])


# --- DX ---


def dx_path(project):
    return os.path.join(MUX_CONFIG, "dx", f"{project}.json")


def dx_active(project):
    d = load_json(dx_path(project))
    return [item for item in d.get("items", []) if not item.get("done")]


# --- CLI dispatcher (for bash callers) ---


def _cli_main():
    if len(sys.argv) < 2:
        print("Usage: muxlib.py <command> [args...]", file=sys.stderr)
        sys.exit(1)

    cmd = sys.argv[1]
    args = sys.argv[2:]

    try:
        if cmd == "project-path":
            result = project_path(args[0])
            if result:
                print(result)
            else:
                sys.exit(1)

        elif cmd == "project-names":
            print("\n".join(project_names()))

        elif cmd == "project-color":
            print(project_color(args[0]))

        elif cmd == "project-setting":
            val = project_setting(args[0])
            if val is not None:
                print(val)

        elif cmd == "project-setting-set":
            value = args[1]
            if value == "true":
                value = True
            elif value == "false":
                value = False
            project_setting_set(args[0], value)

        elif cmd == "narrate-check":
            print(narrate_check(args[0]))

        elif cmd == "narrate-mark":
            narrate_mark(args[0])

        elif cmd == "notes-has":
            print(notes_has(args[0]))

        elif cmd == "notes-count":
            print(notes_count(args[0]))

        elif cmd == "worktree-add":
            worktree_add(args[0], args[1], args[2], args[3])

        elif cmd == "worktree-remove":
            worktree_remove(args[0], args[1])

        elif cmd == "worktree-is-active":
            print(worktree_is_active(args[0], args[1]))

        elif cmd == "worktree-list":
            print(worktree_list())

        elif cmd == "worktree-count":
            print(worktree_count())

        elif cmd == "worktree-slugs":
            print("\n".join(worktree_slugs(args[0])))

        elif cmd == "snapshot-build":
            # stdin: one "desk|window_name|window_index|cwd" row per line
            # (bash has already filtered to live-Claude panes). Writes the
            # snapshot file and prints entries as pipe rows for the caller's
            # summary line — desk and is_worktree only, cheap to count.
            rows = []
            for line in sys.stdin:
                line = line.rstrip('\n')
                if not line:
                    continue
                parts = line.split('|', 3)
                if len(parts) != 4:
                    continue
                desk, win_name, win_idx, cwd = parts
                rows.append((desk, win_name, win_idx, cwd))
            path = args[0] if args else MUX_SNAPSHOT
            entries = snapshot_write(rows, path)
            for e in entries:
                print(f"{e['desk']}|{e['window_name']}|{'1' if e['is_worktree'] else '0'}|{'1' if e['session_id'] else '0'}")

        elif cmd == "snapshot-read":
            # Emits pipe rows for `mux restore`'s bash loop: one per snapshot
            # entry. session_id/task_slug/branch print as the literal string
            # "None" when absent (never blank — blank is indistinguishable
            # from a field that was never emitted, given a plain `cut -d|`).
            path = args[0] if args else MUX_SNAPSHOT
            for e in snapshot_read(path):
                print('|'.join(str(e.get(k)) for k in (
                    'desk', 'window_name', 'window_index', 'cwd',
                    'is_worktree', 'task_slug', 'branch', 'session_id',
                )))

        elif cmd == "snapshot-meta":
            # saved_at + count, for `mux restore`'s pre-flight summary line.
            path = args[0] if args else MUX_SNAPSHOT
            d = load_json(path)
            print(f"{d.get('saved_at', '')}|{len(d.get('windows', []))}")

        else:
            print(f"Unknown command: {cmd}", file=sys.stderr)
            sys.exit(1)

    except IndexError:
        print(f"Error: missing argument for '{cmd}'", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    _cli_main()