#!/usr/bin/env python3
"""fzf-ai-preview <agent> <session_id> <source>

Render a compact, readable preview of an AI coding session for fzf's preview
window. Works for all agents handled by fzf-ai-index.
"""

from __future__ import annotations

from collections import Counter
import json
try:
    import orjson
    _JSON_LOADS = orjson.loads
    _JSON_DUMPS = orjson.dumps
except ImportError:
    _JSON_LOADS = json.loads
    _JSON_DUMPS = json.dumps
import os
import re
import shutil
import sqlite3
import sys
import textwrap
from datetime import datetime, timezone
from pathlib import Path

# ANSI colors (respect NO_COLOR). fzf renders ANSI in the preview window.
NO_COLOR = bool(os.environ.get("NO_COLOR"))

# Preview cache: avoid re-parsing large JSONL files on every preview.
CACHE_DIR = Path(os.environ.get("XDG_CACHE_HOME", str(Path.home() / ".cache"))) / "fzf-ai"
PREVIEW_CACHE_VERSION = 1


def _get_preview_cache_path() -> Path:
    return CACHE_DIR / "preview.json"


def _load_preview_cache() -> dict:
    path = _get_preview_cache_path()
    if not path.is_file():
        return {"v": PREVIEW_CACHE_VERSION, "sessions": {}}
    try:
        data = _JSON_LOADS(path.read_bytes())
        if isinstance(data, dict) and data.get("v") == PREVIEW_CACHE_VERSION:
            return data
    except Exception:
        pass
    return {"v": PREVIEW_CACHE_VERSION, "sessions": {}}


def _save_preview_cache(cache: dict) -> None:
    path = _get_preview_cache_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")
    content = _JSON_DUMPS(cache)
    if isinstance(content, bytes):
        tmp.write_bytes(content)
    else:
        tmp.write_text(content, encoding="utf-8")
    tmp.rename(path)


def _preview_cache_key(agent: str, sid: str) -> str:
    return f"{agent}:{sid}"


def _validate_preview_cache_entry(entry: dict) -> bool:
    """Check if a preview cache entry is still valid."""
    src = entry.get("source", "")
    if src.startswith("sqlite:"):
        db_path = Path(src[len("sqlite:"):])
        return db_path.is_file() and db_path.stat().st_mtime == entry.get("mtime", 0)
    p = Path(src)
    if not p.is_file():
        return False
    try:
        st = p.stat()
        return st.st_mtime == entry.get("mtime", 0) and st.st_size == entry.get("size", 0)
    except OSError:
        return False


def c(code: str, s: str) -> str:
    if NO_COLOR:
        return s
    return f"\033[{code}m{s}\033[0m"


BOLD = lambda s: c("1", s)
DIM = lambda s: c("2", s)
CYAN = lambda s: c("36", s)
GREEN = lambda s: c("32", s)
YELLOW = lambda s: c("33", s)
MAGENTA = lambda s: c("35", s)
RED = lambda s: c("31", s)
BLUE = lambda s: c("34", s)

ROLE_COLOR = {
    "user": CYAN,
    "assistant": GREEN,
    "developer": MAGENTA,
    "system": DIM,
    "tool": YELLOW,
}

NOISE_PREFIXES = (
    "<permissions instructions>",
    "<environment_context>",
    "You are Codex, a coding agent",
    "You are an AI assistant accessed via an API.",
    "Filesystem sandboxing defines",
    "# Escalation Requests",
)

MESSAGE_CLIP = {
    "user": (5, 420),
    "assistant": (8, 900),
    "tool": (6, 700),
    "developer": (4, 260),
    "system": (4, 260),
}


def term_width() -> int:
    try:
        cols = int(os.environ.get("FZF_PREVIEW_COLUMNS", "0")) or shutil.get_terminal_size((100, 40)).columns
    except Exception:
        cols = 100
    return max(40, cols)


def wrap(text: str, indent: str = "  ") -> str:
    w = term_width() - len(indent) - 1
    if w < 20:
        w = 60
    out = []
    for line in text.splitlines() or [""]:
        if not line:
            out.append(indent)
            continue
        out.extend(
            textwrap.wrap(
                line,
                width=w,
                initial_indent=indent,
                subsequent_indent=indent,
                break_long_words=False,
                break_on_hyphens=False,
            )
            or [indent]
        )
    return "\n".join(out)


def home_rel(path: str) -> str:
    home = str(Path.home())
    if path.startswith(home):
        return "~" + path[len(home):]
    return path


def compact_path(path: str, keep: int = 2) -> str:
    if path.startswith("sqlite:"):
        return "sqlite:" + compact_path(path[len("sqlite:"):], keep=keep)
    path = home_rel(path)
    parts = [p for p in path.split("/") if p]
    prefix = "~/" if path.startswith("~/") else "/" if path.startswith("/") else ""
    if len(parts) <= keep:
        return path
    return prefix + "…/" + "/".join(parts[-keep:])


def short_sid(sid: str) -> str:
    return sid if len(sid) <= 8 else sid[:8]


def header(agent: str, sid: str, source: str, cwd: str = "") -> str:
    lines = [
        BOLD(agent) + DIM(f"  sid:{short_sid(sid)}"),
    ]
    if cwd:
        lines.append(DIM("cwd  ") + compact_path(cwd, keep=3))
    lines.append(DIM("src  ") + compact_path(source, keep=2))
    lines.append(DIM("─" * term_width()))
    return "\n".join(lines)


def render_message(role: str, text: str, ts: str = "", terms: list[str] | None = None) -> str:
    color = ROLE_COLOR.get(role, lambda s: s)
    tag = f"{role.upper():<9}"
    head = color(BOLD(tag)) + (DIM(f" {ts}") if ts else "")
    body = wrap(clip_text(clean_text(text), role=role), "  ")
    body = highlight_terms(body, terms or [])
    return head + "\n" + body


def clean_text(text: str) -> str:
    return "\n".join(line.rstrip() for line in (text or "").strip().splitlines()).strip()


def clip_text(text: str, role: str = "assistant", max_lines: int | None = None, max_chars: int | None = None) -> str:
    text = clean_text(text)
    default_lines, default_chars = MESSAGE_CLIP.get(role, (6, 500))
    max_lines = default_lines if max_lines is None else max_lines
    max_chars = default_chars if max_chars is None else max_chars
    if not text:
        return text

    clipped = text[:max_chars].rstrip()
    char_truncated = len(clipped) < len(text)
    lines = clipped.splitlines()
    line_truncated = len(lines) > max_lines
    if line_truncated:
        clipped = "\n".join(lines[:max_lines]).rstrip()
    if char_truncated or line_truncated:
        clipped = clipped.rstrip(" .") + " …"
    return clipped


def query_terms(query: str) -> list[str]:
    terms: list[str] = []
    for raw in query.split():
        token = raw.strip()
        if not token or token == "|":
            continue
        token = token.lstrip("!'^")
        token = token.rstrip("$")
        if token:
            terms.append(token.lower())
    return sorted(set(terms), key=len, reverse=True)


def message_matches(text: str, terms: list[str]) -> bool:
    haystack = clean_text(text).lower()
    return any(term in haystack for term in terms)


def highlight_terms(text: str, terms: list[str]) -> str:
    if NO_COLOR or not terms:
        return text
    out = text
    for term in terms:
        out = re.sub(
            re.escape(term),
            lambda m: c("30;43", m.group(0)),
            out,
            flags=re.IGNORECASE,
        )
    return out


def is_noise_message(role: str, text: str) -> bool:
    text = clean_text(text)
    if not text:
        return True
    if text.startswith("<environment_context>"):
        return True
    if role not in {"system", "developer"}:
        return False
    if any(text.startswith(prefix) for prefix in NOISE_PREFIXES):
        return True
    if "workspace-write" in text and "Filesystem sandboxing" in text:
        return True
    return len(text) > 900


def summarize_messages(total: int, visible: list[tuple[str, str, str]], model: str) -> str:
    counts = Counter(role for role, _text, _ts in visible)
    parts = [f"{total} raw", f"{len(visible)} shown"]
    for role in ("user", "assistant", "tool", "developer", "system"):
        if counts.get(role):
            parts.append(f"{role}={counts[role]}")
    if model:
        parts.append(f"model={model}")
    return "  ·  ".join(parts)


def same_message(a: str, b: str) -> bool:
    return clean_text(a) == clean_text(b)


def select_excerpt(
    messages: list[tuple[str, str, str]],
    max_msgs: int,
) -> tuple[list[tuple[str, str, str]], int]:
    if len(messages) <= max_msgs:
        return messages, 0
    head_count = min(2, len(messages))
    tail_count = max(1, max_msgs - head_count)
    visible = messages[:head_count] + messages[-tail_count:]
    hidden = max(0, len(messages) - len(visible))
    return visible, hidden


def select_match_excerpt(
    messages: list[tuple[str, str, str]],
    terms: list[str],
    max_msgs: int,
) -> tuple[list[tuple[str, str, str]], int, list[int]]:
    matched = [i for i, (_role, text, _ts) in enumerate(messages) if message_matches(text, terms)]
    if not matched:
        visible, hidden = select_excerpt(messages, max_msgs)
        return visible, hidden, matched

    keep: list[int] = []
    seen: set[int] = set()
    for idx in matched:
        for pos in range(max(0, idx - 1), min(len(messages), idx + 2)):
            if pos not in seen:
                seen.add(pos)
                keep.append(pos)
        if len(keep) >= max_msgs:
            break
    keep = keep[:max_msgs]
    visible = [messages[i] for i in keep]
    hidden = max(0, len(messages) - len(visible))
    return visible, hidden, matched


def extract_text(parts):
    """Extract concatenated text from a content list/string in many shapes."""
    if parts is None:
        return ""
    if isinstance(parts, str):
        return parts
    if isinstance(parts, list):
        chunks = []
        for p in parts:
            if not isinstance(p, dict):
                continue
            t = p.get("type")
            if t == "text":
                tx = p.get("text")
                if tx is None:
                    tx = (p.get("data") or {}).get("text")
                if tx:
                    chunks.append(tx)
            elif t == "tool_use":
                name = p.get("name", "tool")
                inp = p.get("input")
                chunks.append(f"⟨tool:{name}⟩ {json.dumps(inp)[:200]}")
            elif t == "tool_result":
                content = p.get("content")
                if isinstance(content, list):
                    for cc in content:
                        if isinstance(cc, dict) and cc.get("type") == "text":
                            chunks.append(f"⟨result⟩ {cc.get('text','')[:300]}")
                else:
                    chunks.append(f"⟨result⟩ {str(content)[:300]}")
            elif t in ("input_text", "output_text"):
                tx = p.get("text", "")
                if tx:
                    chunks.append(tx)
            elif t == "reasoning":
                summaries = p.get("summary") or []
                for s in summaries:
                    if isinstance(s, dict) and s.get("type") == "summary_text":
                        chunks.append("⟨reasoning⟩ " + (s.get("text") or ""))
        return "\n".join(chunks)
    return str(parts)


def fmt_ts(ts) -> str:
    if not ts:
        return ""
    if isinstance(ts, (int, float)):
        try:
            return (
                datetime.fromtimestamp(ts / 1000 if ts > 1e12 else ts, tz=timezone.utc)
                .strftime("%Y-%m-%d %H:%M:%S")
            )
        except Exception:
            return ""
    if isinstance(ts, str):
        try:
            return datetime.fromisoformat(ts.replace("Z", "+00:00")).strftime(
                "%Y-%m-%d %H:%M:%S"
            )
        except Exception:
            return ts
    return ""


# ---------------------------------------------------------------
def preview_jsonl(
    path: Path,
    agent: str,
    sid: str,
    query: str = "",
    max_msgs: int = 20,
) -> str:
    cwd = ""
    total = 0
    model = ""
    messages: list[tuple[str, str, str]] = []
    if not path.is_file():
        return RED(f"file not found: {path}")

    # Check preview cache
    preview_cache = _load_preview_cache()
    cache_key = _preview_cache_key(agent, sid)
    cache_entry = preview_cache.get("sessions", {}).get(cache_key)
    cache_valid = (
        cache_entry
        and cache_entry.get("source") == str(path)
        and _validate_preview_cache_entry(cache_entry)
    )
    if cache_valid:
        cwd = cache_entry.get("cwd", "")
        model = cache_entry.get("model", "")
        total = cache_entry.get("total", 0)
        messages = [tuple(m) for m in cache_entry.get("messages", [])]
    else:
        with path.open("r", encoding="utf-8", errors="replace") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    obj = _JSON_LOADS(line)
                except json.JSONDecodeError:
                    continue
                if agent == "claude":
                    if obj.get("cwd") and not cwd:
                        cwd = obj["cwd"]
                    if obj.get("type") in ("user", "assistant"):
                        if obj.get("isMeta"):
                            continue
                        m = obj.get("message") or {}
                        role = m.get("role") or obj["type"]
                        text = extract_text(m.get("content"))
                        if text:
                            ts = fmt_ts(obj.get("timestamp"))
                            messages.append((role, text, ts))
                            total += 1
                elif agent == "codex":
                    p = obj.get("payload") or {}
                    t = obj.get("type")
                    if t == "session_meta":
                        if p.get("cwd") and not cwd:
                            cwd = p["cwd"]
                    if t == "response_item" and p.get("type") == "message":
                        role = p.get("role")
                        text = extract_text(p.get("content"))
                        if text:
                            messages.append((role, text, fmt_ts(obj.get("timestamp"))))
                            total += 1
                    elif t == "response_item" and p.get("type") == "reasoning":
                        text = extract_text([p])
                        if text:
                            messages.append(("assistant", text, fmt_ts(obj.get("timestamp"))))
                            total += 1
                elif agent == "droid":
                    t = obj.get("type")
                    if t == "session_start":
                        if obj.get("cwd") and not cwd:
                            cwd = obj["cwd"]
                    elif t == "message":
                        m = obj.get("message") or {}
                        role = m.get("role")
                        text = extract_text(m.get("content"))
                        if text:
                            messages.append((role, text, fmt_ts(obj.get("timestamp"))))
                            total += 1
                elif agent == "pi":
                    t = obj.get("type")
                    if t == "session":
                        if obj.get("cwd") and not cwd:
                            cwd = obj["cwd"]
                    elif t == "model_change":
                        model = obj.get("modelId") or model
                    elif t == "message":
                        m = obj.get("message") or {}
                        role = m.get("role")
                        text = extract_text(m.get("content"))
                        if text:
                            messages.append((role, text, fmt_ts(obj.get("timestamp"))))
                            total += 1

        # Save to preview cache
        try:
            st = path.stat()
            preview_cache.setdefault("sessions", {})[cache_key] = {
                "mtime": st.st_mtime,
                "size": st.st_size,
                "source": str(path),
                "cwd": cwd,
                "model": model,
                "total": total,
                "messages": [list(m) for m in messages],
            }
            _save_preview_cache(preview_cache)
        except Exception:
            pass

    visible_messages = [
        (role, text, ts)
        for role, text, ts in messages
        if not is_noise_message(role, text)
    ]
    if not visible_messages:
        visible_messages = messages

    user_prompts = [
        clean_text(text)
        for role, text, _ts in visible_messages
        if role == "user" and clean_text(text)
    ]
    terms = query_terms(query)
    shown_msgs, hidden_count, matched_indices = select_match_excerpt(
        visible_messages,
        terms=terms,
        max_msgs=max_msgs,
    )
    first_prompt = user_prompts[0] if user_prompts else ""
    last_prompt = user_prompts[-1] if user_prompts else ""
    if not matched_indices and shown_msgs and first_prompt and shown_msgs[0][0] == "user" and same_message(shown_msgs[0][1], first_prompt):
        shown_msgs = shown_msgs[1:]
    if not matched_indices and shown_msgs and len(user_prompts) > 1 and shown_msgs[-1][0] == "user" and same_message(shown_msgs[-1][1], last_prompt):
        shown_msgs = shown_msgs[:-1]

    parts = [header(agent, sid, str(path), cwd)]
    parts.append(DIM(summarize_messages(total, visible_messages, model)))
    skipped = total - len(visible_messages)
    if skipped > 0:
        parts.append(DIM(f"{skipped} system/developer messages hidden"))
    parts.append("")

    if terms and matched_indices:
        parts.append(BOLD(f"Query Matches ({len(matched_indices)})"))
        for idx in matched_indices[:2]:
            role, text, _ts = visible_messages[idx]
            parts.append(
                wrap(
                    highlight_terms(clip_text(text, role=role, max_lines=2, max_chars=180), terms),
                    "  ",
                )
            )
            parts.append("")
    elif user_prompts:
        parts.append(BOLD("First Prompt"))
        parts.append(wrap(clip_text(first_prompt, role="user", max_lines=3, max_chars=220), "  "))
        if len(user_prompts) > 1 and last_prompt != first_prompt:
            parts.append("")
            parts.append(BOLD("Last User Prompt"))
            parts.append(wrap(clip_text(last_prompt, role="user", max_lines=3, max_chars=220), "  "))
        parts.append("")

    if hidden_count > 0:
        parts.append(DIM(f"    … ({hidden_count} messages hidden) …"))
        parts.append("")

    for role, text, ts in shown_msgs:
        parts.append(render_message(role, text, ts, terms=terms))
        parts.append("")
    return "\n".join(parts)


# ---------------------------------------------------------------
def preview_opencode(sid: str, db_path: str) -> str:
    db = Path(db_path.replace("sqlite:", ""))
    if not db.is_file():
        return RED(f"db not found: {db}")
    con = sqlite3.connect(f"file:{db}?mode=ro", uri=True)
    try:
        tables = {
            r[0]
            for r in con.execute(
                "SELECT name FROM sqlite_master WHERE type='table'"
            ).fetchall()
        }
        if {"session", "message", "part"}.issubset(tables):
            return _preview_opencode_current(con, db, sid)
        if {"sessions", "messages"}.issubset(tables):
            return _preview_opencode_legacy(con, db, sid)
        return RED(f"{db}: unknown opencode schema")
    finally:
        con.close()


def _preview_opencode_current(con: sqlite3.Connection, db: Path, sid: str) -> str:
    row = con.execute(
        "SELECT id, directory, title, time_updated FROM session WHERE id=?",
        (sid,),
    ).fetchone()
    if not row:
        return RED(f"session not found: {sid}")
    _, directory, title, updated = row
    lines = [header("opencode", sid, f"sqlite:{db}", cwd=directory or "")]
    lines.append(BOLD(title or "(untitled)"))
    lines.append(DIM(f"updated {fmt_ts(updated)}"))
    lines.append("")

    # Collect messages ordered by time; fetch parts grouped by message.
    msgs = con.execute(
        "SELECT id, data, time_created FROM message WHERE session_id=? ORDER BY time_created ASC",
        (sid,),
    ).fetchall()
    total = len(msgs)
    if total == 0:
        return "\n".join(lines)

    # Pre-fetch parts for all messages in one go.
    parts_by_msg: dict[str, list] = {}
    for mid, pdata, _tc in con.execute(
        "SELECT message_id, data, time_created FROM part WHERE session_id=? ORDER BY time_created ASC",
        (sid,),
    ):
        try:
            parts_by_msg.setdefault(mid, []).append(_JSON_LOADS(pdata))
        except Exception:
            pass

    # Pick messages to show: first 3 + last 3
    show_indices = list(range(min(3, total)))
    if total > 6:
        show_indices += list(range(max(3, total - 3), total))
    hidden = total - len(show_indices)
    prev_i = -2
    for i in show_indices:
        mid, mdata, tc = msgs[i]
        if i > prev_i + 1 and hidden > 0:
            lines.append(DIM(f"    … ({hidden} messages hidden) …"))
            lines.append("")
        try:
            d = _JSON_LOADS(mdata)
        except Exception:
            d = {}
        role = d.get("role", "?")
        text = extract_text(parts_by_msg.get(mid) or [])
        if text:
            lines.append(render_message(role, text, fmt_ts(tc)))
            lines.append("")
        prev_i = i
    return "\n".join(lines)


def _preview_opencode_legacy(con: sqlite3.Connection, db: Path, sid: str) -> str:
    cur = con.execute(
        "SELECT id,title,message_count,cost,prompt_tokens,completion_tokens,updated_at "
        "FROM sessions WHERE id=?",
        (sid,),
    )
    row = cur.fetchone()
    if not row:
        return RED(f"session not found: {sid}")
    _, title, mcount, cost, pt, ct, updated = row
    lines = [header("opencode", sid, f"sqlite:{db}")]
    lines.append(BOLD(title or "(untitled)"))
    lines.append(
        DIM(
            f"{mcount} msgs  ·  ${cost:.4f}  ·  {pt} in / {ct} out  ·  updated {fmt_ts(updated)}"
        )
    )
    lines.append("")
    cur2 = con.execute(
        "SELECT role,parts,created_at FROM messages WHERE session_id=? ORDER BY created_at ASC",
        (sid,),
    )
    rows = cur2.fetchall()
    total = len(rows)
    show_indices = list(range(min(3, total)))
    if total > 6:
        show_indices += list(range(max(3, total - 3), total))
    hidden = total - len(show_indices)
    prev_i = -2
    for i in show_indices:
        role, parts_json, created = rows[i]
        if i > prev_i + 1 and hidden > 0:
            lines.append(DIM(f"    … ({hidden} messages hidden) …"))
            lines.append("")
        try:
            p = _JSON_LOADS(parts_json)
        except Exception:
            p = []
        text = extract_text(p)
        if text:
            lines.append(render_message(role, text, fmt_ts(created)))
            lines.append("")
        prev_i = i
    return "\n".join(lines)


# ---------------------------------------------------------------
def main() -> int:
    if len(sys.argv) < 4:
        print("usage: fzf-ai-preview <agent> <session_id> <source> [query]", file=sys.stderr)
        return 2
    agent, sid, source = sys.argv[1], sys.argv[2], sys.argv[3]
    query = sys.argv[4] if len(sys.argv) > 4 else ""
    try:
        if agent == "opencode":
            out = preview_opencode(sid, source)
        else:
            out = preview_jsonl(Path(source), agent, sid, query=query)
        sys.stdout.write(out)
        sys.stdout.write("\n")
    except BrokenPipeError:
        pass
    except Exception as e:
        print(RED(f"preview error: {e}"))
    return 0


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