#!/usr/bin/env python3
"""
pi-resume (aka `pir`) - fuzzy-find and resume old pi sessions for the current dir.

Instead of `pi -c` (blindly grabs the newest) or digging through
~/.pi/agent/sessions/--home-user-dev-foo--/ for a UUID by hand, this lists every
session for the cwd in an fzf picker with a live transcript preview, then execs
`pi --session <file>` on the one you pick.

Usage:
  pir                 pick a session for the current dir, resume it
  pir -a/--all        pick across every project dir
  pir -q QUERY        start fzf with an initial query
  pir <query...>      same, bare words are treated as the query
  pir -g TEXT         only sessions whose full transcript contains TEXT
                      (the fzf query alone matches title/id, not deep history)

In fzf:
  enter     resume the session
  ctrl-f    fork it instead (leaves the original untouched)
  ctrl-o    print the session file path and exit (no launch)
  ctrl-a    toggle current-dir <-> all-dirs
  ctrl-s    toggle sort: recent <-> longest
  ctrl-/    toggle preview
  ?         toggle preview window size

Internal modes (used by fzf itself): --list, --preview <file>
"""

from __future__ import annotations

import json
import os
import re
import shlex
import shutil
import subprocess
import sys
import time
from datetime import datetime, timezone

# ---------------------------------------------------------------- paths

APP = os.environ.get("PI_APP_NAME", "pi")


def agent_dir() -> str:
    env = os.environ.get(f"{APP.upper()}_CODING_AGENT_DIR")
    if env:
        return os.path.expanduser(env)
    return os.path.join(os.path.expanduser("~"), f".{APP}", "agent")


def sessions_root() -> str:
    return os.path.join(agent_dir(), "sessions")


def encode_cwd(cwd: str) -> str:
    """Mirror pi's getDefaultSessionDirPath() encoding exactly."""
    resolved = os.path.realpath(os.path.expanduser(cwd))
    safe = re.sub(r"[/\\:]", "-", re.sub(r"^[/\\]", "", resolved))
    return f"--{safe}--"


def session_dir_for(cwd: str) -> str:
    env = os.environ.get(f"{APP.upper()}_CODING_AGENT_SESSION_DIR")
    if env:
        return os.path.expanduser(env)
    return os.path.join(sessions_root(), encode_cwd(cwd))


CACHE = os.path.join(
    os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache")),
    "pi-resume",
    "index.json",
)

# ---------------------------------------------------------------- scanning

MSG_MARK = b'"type":"message"'
INFO_MARK = b'"type":"session_info"'
HEAD_BYTES = 512 * 1024


def _text_of(msg: dict) -> str:
    """Flatten a message's content into plain text (ignore thinking/tool noise)."""
    content = msg.get("content")
    if isinstance(content, str):
        return content
    if not isinstance(content, list):
        return ""
    out = []
    for part in content:
        if isinstance(part, dict) and part.get("type") == "text":
            out.append(part.get("text") or "")
    return " ".join(out)


def _clean(s: str, limit: int = 400) -> str:
    s = re.sub(r"\s+", " ", (s or "")).strip()
    return s[:limit]


def scan(path: str) -> dict | None:
    """Cheap metadata scan: header, msg count, name, first user message."""
    try:
        st = os.stat(path)
    except OSError:
        return None

    msgs = 0
    has_info = False
    head = bytearray()
    tail = b""
    keep = max(len(MSG_MARK), len(INFO_MARK)) - 1

    try:
        with open(path, "rb") as f:
            while True:
                chunk = f.read(1 << 20)
                if not chunk:
                    break
                if len(head) < HEAD_BYTES:
                    head += chunk[: HEAD_BYTES - len(head)]
                buf = tail + chunk
                msgs += buf.count(MSG_MARK)
                if not has_info and INFO_MARK in buf:
                    has_info = True
                tail = buf[-keep:] if keep else b""
    except OSError:
        return None

    header, first_user, model = {}, "", ""
    for line in head.split(b"\n"):
        if not line.strip():
            continue
        try:
            rec = json.loads(line)
        except Exception:
            continue  # truncated final line of the head window
        t = rec.get("type")
        if t == "session" and not header:
            header = rec
        elif t == "model_change" and not model:
            model = rec.get("modelId") or ""
        elif t == "message" and not first_user:
            m = rec.get("message") or {}
            if m.get("role") == "user":
                txt = _clean(_text_of(m))
                if txt:
                    first_user = txt
        if header and first_user and model:
            break

    name = ""
    if has_info:
        # Rare; a full line pass is fine here. Last one wins.
        try:
            with open(path, "rb") as f:
                for line in f:
                    if INFO_MARK not in line:
                        continue
                    try:
                        rec = json.loads(line)
                    except Exception:
                        continue
                    if rec.get("type") == "session_info" and rec.get("name"):
                        name = _clean(rec["name"], 120)
        except OSError:
            pass

    return {
        "path": path,
        "id": header.get("id") or os.path.basename(path).split("_", 1)[-1][:-6],
        "cwd": header.get("cwd") or "",
        "forked": bool(header.get("parentSession")),
        "created": header.get("timestamp") or "",
        "mtime": st.st_mtime,
        "size": st.st_size,
        "msgs": msgs,
        "name": name,
        "model": model,
        "first": first_user,
    }


def load_cache() -> dict:
    try:
        with open(CACHE) as f:
            return json.load(f)
    except Exception:
        return {}


def save_cache(cache: dict) -> None:
    try:
        os.makedirs(os.path.dirname(CACHE), exist_ok=True)
        tmp = CACHE + f".{os.getpid()}.tmp"
        with open(tmp, "w") as f:
            json.dump(cache, f)
        os.replace(tmp, CACHE)
    except Exception:
        pass


def grep_filter(files: list[str], pattern: str) -> set[str]:
    """Full-text prefilter over raw session files. Returns matching paths."""
    if not files:
        return set()
    tool = "rg" if shutil.which("rg") else "grep"
    cmd = [tool, "-l", "-i", "-F", "--", pattern] if tool == "rg" else [
        "grep", "-l", "-i", "-F", "-e", pattern, "--"
    ]
    hits: set[str] = set()
    # chunk to stay under ARG_MAX
    for i in range(0, len(files), 400):
        batch = files[i : i + 400]
        try:
            p = subprocess.run(cmd + batch, capture_output=True, text=True)
        except OSError:
            return set(files)
        hits.update(l.strip() for l in p.stdout.splitlines() if l.strip())
    return hits


def collect(all_dirs: bool, grep: str = "") -> list[dict]:
    if all_dirs:
        root = sessions_root()
        dirs = [os.path.join(root, d) for d in os.listdir(root)] if os.path.isdir(root) else []
    else:
        dirs = [session_dir_for(os.getcwd())]

    files = []
    for d in dirs:
        if not os.path.isdir(d):
            continue
        for fn in os.listdir(d):
            if fn.endswith(".jsonl"):
                files.append(os.path.join(d, fn))

    cache = load_cache()
    out, dirty = [], False
    for p in files:
        try:
            st = os.stat(p)
        except OSError:
            continue
        hit = cache.get(p)
        if hit and hit.get("mtime") == st.st_mtime and hit.get("size") == st.st_size:
            out.append(hit)
            continue
        rec = scan(p)
        if rec:
            cache[p] = rec
            out.append(rec)
            dirty = True

    live = set(files)
    for stale in [k for k in cache if k not in live]:
        del cache[stale]
        dirty = True
    if dirty:
        save_cache(cache)

    out = [r for r in out if r.get("msgs")]
    if grep:
        keep = grep_filter([r["path"] for r in out], grep)
        out = [r for r in out if r["path"] in keep]
    return out


# ---------------------------------------------------------------- display

C_DIM = "\x1b[38;5;242m"
C_AGE = "\x1b[38;5;110m"
C_NUM = "\x1b[38;5;108m"
C_NAME = "\x1b[38;5;215m"
C_TXT = "\x1b[0m"
C_HDR = "\x1b[1;38;5;215m"
C_USER = "\x1b[1;38;5;110m"
C_ASST = "\x1b[38;5;250m"
C_TOOL = "\x1b[38;5;242m"
R = "\x1b[0m"


def age(ts: float) -> str:
    d = max(0.0, time.time() - ts)
    if d < 3600:
        return f"{int(d // 60)}m"
    if d < 86400:
        return f"{int(d // 3600)}h"
    if d < 86400 * 7:
        return f"{int(d // 86400)}d"
    return datetime.fromtimestamp(ts).strftime("%b %d")


def age_phrase(ts: float) -> str:
    """Human phrasing: '3h ago' for recent, absolute date for old."""
    if time.time() - ts < 86400 * 7:
        return f"{age(ts)} ago"
    return datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M")


def fmt_line(r: dict, show_cwd: bool) -> str:
    title = r["name"] or r["first"] or "(no messages)"
    label = f"{C_NAME}{title}{R}" if r["name"] else f"{C_TXT}{title}{R}"
    bits = [
        f"{C_AGE}{age(r['mtime']):>6}{R}",
        f"{C_NUM}{r['msgs']:>4} msg{R}",
    ]
    if show_cwd:
        home = os.path.expanduser("~")
        cwd = r["cwd"] or "?"
        cwd = "~" + cwd[len(home):] if cwd.startswith(home) else cwd
        bits.append(f"{C_DIM}{cwd:<28}{R}")
    if r["forked"]:
        bits.append(f"{C_DIM}fork{R}")
    return "  ".join(bits) + "  " + label


def emit_list(all_dirs: bool, sort: str, grep: str = "") -> None:
    rows = collect(all_dirs, grep)
    if sort == "longest":
        rows.sort(key=lambda r: r["msgs"], reverse=True)
    else:
        rows.sort(key=lambda r: r["mtime"], reverse=True)
    buf = []
    for r in rows:
        # field 1 = path (hidden from display via --with-nth), field 2 = pretty
        # searchable haystack is appended dim so fzf can match on it too
        hay = f"{r['id']} {r['first']}"
        buf.append(f"{r['path']}\t{fmt_line(r, all_dirs)}\t{C_DIM}{_clean(hay, 300)}{R}")
    sys.stdout.write("\n".join(buf) + ("\n" if buf else ""))


# ---------------------------------------------------------------- preview

TAIL_BYTES = 1_500_000


def preview(path: str) -> None:
    try:
        st = os.stat(path)
    except OSError:
        print("gone")
        return

    truncated = False
    with open(path, "rb") as f:
        if st.st_size <= 2_000_000:
            data = f.read()
        else:
            head = f.read(256 * 1024)
            f.seek(-TAIL_BYTES, os.SEEK_END)
            tailbuf = f.read()
            # drop partial lines at both seams
            head = head[: head.rfind(b"\n") + 1]
            tailbuf = tailbuf[tailbuf.find(b"\n") + 1 :]
            data = head + b"\n@@TRUNC@@\n" + tailbuf
            truncated = True

    meta, out = {}, []
    for raw in data.split(b"\n"):
        if not raw.strip():
            continue
        if raw == b"@@TRUNC@@":
            out.append(f"{C_TOOL}{'-' * 20} [ ... middle of session omitted ... ] {'-' * 20}{R}")
            continue
        try:
            rec = json.loads(raw)
        except Exception:
            continue
        t = rec.get("type")
        if t == "session":
            meta = rec
        elif t == "session_info" and rec.get("name"):
            meta["name"] = rec["name"]
        elif t == "compaction":
            out.append(f"{C_TOOL}--- compacted ---{R}")
        elif t == "message":
            m = rec.get("message") or {}
            role = m.get("role")
            txt = _clean(_text_of(m), 100000)
            if role == "user":
                if not txt:
                    continue
                out.append("")
                out.append(f"{C_USER}> {txt[:1500]}{R}")
            elif role == "assistant":
                tools = [
                    p.get("name")
                    for p in (m.get("content") or [])
                    if isinstance(p, dict) and p.get("type") == "toolCall"
                ]
                if txt:
                    out.append(f"{C_ASST}  {txt[:600]}{R}")
                if tools:
                    seen = ", ".join(dict.fromkeys(t for t in tools if t))
                    out.append(f"{C_TOOL}  [{seen}]{R}")

    created = meta.get("timestamp", "")
    if created:
        try:
            created = (
                datetime.fromisoformat(created.replace("Z", "+00:00"))
                .astimezone()
                .strftime("%Y-%m-%d %H:%M")
            )
        except Exception:
            pass

    print(f"{C_HDR}{meta.get('name') or os.path.basename(path)}{R}")
    print(f"{C_DIM}id      {meta.get('id', '?')}{R}")
    print(f"{C_DIM}cwd     {meta.get('cwd', '?')}{R}")
    print(
        f"{C_DIM}started {created}   modified {age_phrase(st.st_mtime)}"
        f"   {st.st_size // 1024} KB{R}"
    )
    if truncated:
        print(f"{C_DIM}(large session: showing head + tail){R}")
    print()
    print("\n".join(out))


# ---------------------------------------------------------------- main

def die(msg: str, code: int = 1):
    print(f"pir: {msg}", file=sys.stderr)
    sys.exit(code)


def main() -> None:
    argv = sys.argv[1:]
    self_path = os.path.abspath(__file__)

    # internal modes invoked by fzf
    if argv and argv[0] == "--preview":
        preview(argv[1])
        return
    if argv and argv[0] == "--list":
        g = ""
        if "--grep" in argv:
            gi = argv.index("--grep")
            g = argv[gi + 1] if gi + 1 < len(argv) else ""
        emit_list("--all" in argv, "longest" if "--longest" in argv else "recent", g)
        return

    all_dirs, query, grep, words = False, "", "", []
    i = 0
    while i < len(argv):
        a = argv[i]
        if a in ("-a", "--all"):
            all_dirs = True
        elif a in ("-q", "--query"):
            i += 1
            query = argv[i] if i < len(argv) else ""
        elif a in ("-g", "--grep"):
            i += 1
            grep = argv[i] if i < len(argv) else ""
        elif a in ("-h", "--help"):
            print(__doc__.strip())
            return
        else:
            words.append(a)
        i += 1
    if words and not query:
        query = " ".join(words)

    if not shutil.which("fzf"):
        die("fzf not found in PATH")
    pi_bin = os.environ.get("PI_BIN", APP)
    if not shutil.which(pi_bin):
        die(f"'{pi_bin}' not found in PATH")

    rows = collect(all_dirs, grep)
    if not rows:
        where = "any project" if all_dirs else os.getcwd()
        extra = f" matching {grep!r}" if grep else ""
        die(f"no sessions found for {where}{extra}", 1)

    gargs = ["--grep", grep] if grep else []
    listing = subprocess.run(
        [sys.executable, self_path, "--list"] + (["--all"] if all_dirs else []) + gargs,
        capture_output=True,
        text=True,
    ).stdout

    gsuffix = f" --grep {shlex.quote(grep)}" if grep else ""
    base = f"{shlex.quote(sys.executable)} {shlex.quote(self_path)} --list"
    fzf_args = [
        "fzf",
        "--ansi",
        "--delimiter=\t",
        "--with-nth=2..",
        "--layout=reverse",
        "--info=inline",
        "--border",
        "--height=100%",
        f"--query={query}",
        "--prompt=session> ",
        "--header="
        "enter resume | ctrl-f fork | ctrl-o print path | ctrl-a all dirs | ctrl-s sort",
        f"--preview={shlex.quote(sys.executable)} {shlex.quote(self_path)} --preview {{1}}",
        "--preview-window=right,55%,wrap",
        "--bind=ctrl-/:toggle-preview",
        "--bind=?:change-preview-window(right,80%|right,55%)",
        f"--bind=ctrl-a:reload({base} --all{gsuffix})",
        f"--bind=ctrl-s:reload({base}{' --all' if all_dirs else ''}{gsuffix} --longest)",
        "--expect=ctrl-f,ctrl-o",
    ]

    proc = subprocess.run(fzf_args, input=listing, capture_output=True, text=True)
    if proc.returncode not in (0, 1) or not proc.stdout.strip():
        sys.exit(130)  # cancelled

    out = proc.stdout.split("\n")
    key = out[0].strip()
    pick = next((l for l in out[1:] if l.strip()), "")
    if not pick:
        sys.exit(130)
    path = pick.split("\t", 1)[0]

    if key == "ctrl-o":
        print(path)
        return

    flag = "--fork" if key == "ctrl-f" else "--session"
    os.execvp(pi_bin, [pi_bin, flag, path])


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(130)
    except BrokenPipeError:
        sys.exit(0)