#!/usr/bin/env python3
"""
Bezalel's Memory Well — Conversation History Search

Converts Claude Code JSONL conversation transcripts into a searchable
SQLite database with FTS5 full-text search. This is the bridge between
sessions — the builder's ability to remember what was said, what broke,
and what was decided.

This is a construction tool. It runs when the builder needs to search
past conversations. It is not a running service.

Usage:
    # Index all transcripts
    python conversation_search.py index

    # Search for a term
    python conversation_search.py search "shoresh derivation"

    # Search with context (shows surrounding messages)
    python conversation_search.py search "siluq meteg" --context 3

    # List all sessions
    python conversation_search.py sessions

    # Re-index (rebuild from scratch)
    python conversation_search.py reindex
"""

import json
import sqlite3
import sys
import os
import argparse
from pathlib import Path
from datetime import datetime
from typing import Optional


# Default paths
CLAUDE_PROJECTS_DIR = Path.home() / ".claude" / "projects"
DB_PATH = Path.home() / ".claude" / "conversation_history.db"


def create_db(db_path: Path) -> sqlite3.Connection:
    """Create the SQLite database with FTS5 tables."""
    conn = sqlite3.connect(str(db_path))
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA synchronous=NORMAL")

    conn.executescript("""
        CREATE TABLE IF NOT EXISTS sessions (
            session_id TEXT PRIMARY KEY,
            source_file TEXT NOT NULL,
            first_timestamp TEXT,
            last_timestamp TEXT,
            message_count INTEGER DEFAULT 0,
            user_message_count INTEGER DEFAULT 0,
            assistant_message_count INTEGER DEFAULT 0
        );

        CREATE TABLE IF NOT EXISTS messages (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT NOT NULL,
            uuid TEXT UNIQUE,
            parent_uuid TEXT,
            type TEXT NOT NULL,
            role TEXT,
            timestamp TEXT,
            content TEXT,
            model TEXT,
            cwd TEXT,
            has_tool_use INTEGER DEFAULT 0,
            has_thinking INTEGER DEFAULT 0,
            FOREIGN KEY (session_id) REFERENCES sessions(session_id)
        );

        CREATE INDEX IF NOT EXISTS idx_messages_session
            ON messages(session_id);
        CREATE INDEX IF NOT EXISTS idx_messages_timestamp
            ON messages(timestamp);
        CREATE INDEX IF NOT EXISTS idx_messages_type
            ON messages(type);

        -- FTS5 virtual table for full-text search
        CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
            content,
            session_id UNINDEXED,
            uuid UNINDEXED,
            role UNINDEXED,
            timestamp UNINDEXED,
            content='messages',
            content_rowid='id',
            tokenize='porter unicode61'
        );

        -- Triggers to keep FTS in sync
        CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN
            INSERT INTO messages_fts(rowid, content, session_id, uuid, role, timestamp)
            VALUES (new.id, new.content, new.session_id, new.uuid, new.role, new.timestamp);
        END;

        CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN
            INSERT INTO messages_fts(messages_fts, rowid, content, session_id, uuid, role, timestamp)
            VALUES ('delete', old.id, old.content, old.session_id, old.uuid, old.role, old.timestamp);
        END;

        CREATE TRIGGER IF NOT EXISTS messages_au AFTER UPDATE ON messages BEGIN
            INSERT INTO messages_fts(messages_fts, rowid, content, session_id, uuid, role, timestamp)
            VALUES ('delete', old.id, old.content, old.session_id, old.uuid, old.role, old.timestamp);
            INSERT INTO messages_fts(rowid, content, session_id, uuid, role, timestamp)
            VALUES (new.id, new.content, new.session_id, new.uuid, new.role, new.timestamp);
        END;
    """)

    conn.commit()
    return conn


def extract_text_content(message_data: dict) -> Optional[str]:
    """Extract readable text from a message's content field.

    Handles both string content and array content blocks.
    Extracts text, tool_use inputs, and tool_result content.
    Skips thinking blocks (those are private to the model).
    """
    msg = message_data.get("message", {})
    content = msg.get("content")

    if content is None:
        return None

    if isinstance(content, str):
        return content if content.strip() else None

    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, str):
                parts.append(block)
            elif isinstance(block, dict):
                block_type = block.get("type", "")

                if block_type == "text":
                    text = block.get("text", "")
                    if text.strip():
                        parts.append(text)

                elif block_type == "tool_use":
                    tool_name = block.get("name", "unknown")
                    tool_input = block.get("input", {})
                    # Include tool name and key inputs for searchability
                    input_str = json.dumps(tool_input, ensure_ascii=False)
                    # Truncate very long tool inputs
                    if len(input_str) > 2000:
                        input_str = input_str[:2000] + "..."
                    parts.append(f"[tool:{tool_name}] {input_str}")

                elif block_type == "tool_result":
                    result_content = block.get("content", "")
                    if isinstance(result_content, str) and result_content.strip():
                        # Truncate very long tool results
                        if len(result_content) > 3000:
                            result_content = result_content[:3000] + "..."
                        parts.append(f"[result] {result_content}")

                # Skip thinking blocks — those are the model's private reasoning

        return "\n".join(parts) if parts else None

    return None


def index_jsonl_file(conn: sqlite3.Connection, jsonl_path: Path) -> dict:
    """Index a single JSONL transcript file into the database.

    Returns stats about what was indexed.
    """
    stats = {
        "file": str(jsonl_path),
        "total_lines": 0,
        "messages_indexed": 0,
        "skipped": 0,
        "errors": 0,
    }

    session_id = jsonl_path.stem  # filename without extension
    timestamps = []
    user_count = 0
    assistant_count = 0

    with open(jsonl_path, "r", encoding="utf-8") as f:
        for line_num, line in enumerate(f, 1):
            stats["total_lines"] += 1
            line = line.strip()
            if not line:
                continue

            try:
                record = json.loads(line)
            except json.JSONDecodeError:
                stats["errors"] += 1
                continue

            record_type = record.get("type")

            # Only index user and assistant messages
            if record_type not in ("user", "assistant"):
                stats["skipped"] += 1
                continue

            uuid = record.get("uuid")
            if not uuid:
                stats["skipped"] += 1
                continue

            text_content = extract_text_content(record)
            if not text_content:
                stats["skipped"] += 1
                continue

            parent_uuid = record.get("parentUuid")
            timestamp = record.get("timestamp")
            msg = record.get("message", {})
            role = msg.get("role", record_type)
            model = msg.get("model")
            cwd = record.get("cwd")

            # Detect tool use and thinking
            has_tool_use = 0
            has_thinking = 0
            content = msg.get("content", [])
            if isinstance(content, list):
                for block in content:
                    if isinstance(block, dict):
                        if block.get("type") == "tool_use":
                            has_tool_use = 1
                        if block.get("type") == "thinking":
                            has_thinking = 1

            if timestamp:
                timestamps.append(timestamp)

            if role == "user":
                user_count += 1
            elif role == "assistant":
                assistant_count += 1

            try:
                conn.execute("""
                    INSERT OR REPLACE INTO messages
                    (session_id, uuid, parent_uuid, type, role, timestamp,
                     content, model, cwd, has_tool_use, has_thinking)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """, (session_id, uuid, parent_uuid, record_type, role,
                      timestamp, text_content, model, cwd,
                      has_tool_use, has_thinking))
                stats["messages_indexed"] += 1
            except sqlite3.Error as e:
                stats["errors"] += 1

    # Update session record
    first_ts = min(timestamps) if timestamps else None
    last_ts = max(timestamps) if timestamps else None

    conn.execute("""
        INSERT OR REPLACE INTO sessions
        (session_id, source_file, first_timestamp, last_timestamp,
         message_count, user_message_count, assistant_message_count)
        VALUES (?, ?, ?, ?, ?, ?, ?)
    """, (session_id, str(jsonl_path), first_ts, last_ts,
          stats["messages_indexed"], user_count, assistant_count))

    conn.commit()
    return stats


def find_jsonl_files(projects_dir: Path) -> list[Path]:
    """Find all JSONL conversation transcripts."""
    files = []
    if projects_dir.exists():
        for jsonl in projects_dir.rglob("*.jsonl"):
            files.append(jsonl)
    return sorted(files, key=lambda p: p.stat().st_mtime, reverse=True)


def cmd_index(args):
    """Index all JSONL transcripts into the search database."""
    projects_dir = Path(args.projects_dir)
    db_path = Path(args.db)

    jsonl_files = find_jsonl_files(projects_dir)
    if not jsonl_files:
        print(f"No JSONL files found in {projects_dir}")
        return

    print(f"Found {len(jsonl_files)} transcript(s)")

    conn = create_db(db_path)

    for jsonl_path in jsonl_files:
        size_mb = jsonl_path.stat().st_size / (1024 * 1024)
        print(f"\nIndexing: {jsonl_path.name} ({size_mb:.1f} MB)")
        stats = index_jsonl_file(conn, jsonl_path)
        print(f"  Lines: {stats['total_lines']}")
        print(f"  Indexed: {stats['messages_indexed']}")
        print(f"  Skipped: {stats['skipped']} (non-message records)")
        if stats["errors"]:
            print(f"  Errors: {stats['errors']}")

    # Print total stats
    total = conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0]
    sessions = conn.execute("SELECT COUNT(*) FROM sessions").fetchone()[0]
    print(f"\nTotal: {total} messages across {sessions} session(s)")
    print(f"Database: {db_path}")

    conn.close()


def cmd_reindex(args):
    """Drop and rebuild the entire database."""
    db_path = Path(args.db)
    if db_path.exists():
        os.remove(db_path)
        print(f"Removed existing database: {db_path}")
    cmd_index(args)


def cmd_search(args):
    """Search conversation history using FTS5."""
    db_path = Path(args.db)
    if not db_path.exists():
        print(f"Database not found: {db_path}")
        print("Run 'conversation_search.py index' first.")
        return

    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row

    query = args.query
    limit = args.limit
    context = args.context
    role_filter = args.role

    # Build the search query
    sql = """
        SELECT
            m.id,
            m.session_id,
            m.uuid,
            m.role,
            m.timestamp,
            m.content,
            m.model,
            m.has_tool_use,
            m.has_thinking,
            highlight(messages_fts, 0, '>>>', '<<<') as highlighted,
            rank
        FROM messages_fts
        JOIN messages m ON m.id = messages_fts.rowid
        WHERE messages_fts MATCH ?
    """
    params = [query]

    if role_filter:
        sql += " AND m.role = ?"
        params.append(role_filter)

    sql += " ORDER BY rank LIMIT ?"
    params.append(limit)

    results = conn.execute(sql, params).fetchall()

    if not results:
        print(f"No results for: {query}")
        return

    print(f"Found {len(results)} result(s) for: {query}\n")

    for i, row in enumerate(results):
        ts = row["timestamp"]
        if ts:
            try:
                dt = datetime.fromisoformat(ts.replace("Z", "+00:00"))
                ts_display = dt.strftime("%Y-%m-%d %H:%M")
            except (ValueError, AttributeError):
                ts_display = ts[:16] if ts else "?"
        else:
            ts_display = "?"

        role = row["role"] or "?"
        model = row["model"] or ""
        session = row["session_id"][:8]

        print(f"--- [{i+1}] {role} | {ts_display} | session:{session}... ---")
        if model:
            print(f"    model: {model}")

        # Show highlighted content
        highlighted = row["highlighted"]
        # Trim to reasonable length for display
        lines = highlighted.split("\n")
        display_lines = []
        for line in lines:
            if ">>>" in line or len(display_lines) < 5:
                display_lines.append(line)
            if len(display_lines) >= 20:
                display_lines.append("...")
                break

        for line in display_lines:
            # Truncate very long lines
            if len(line) > 200:
                line = line[:200] + "..."
            print(f"  {line}")

        # Show context messages if requested
        if context > 0:
            context_msgs = conn.execute("""
                SELECT role, timestamp, content
                FROM messages
                WHERE session_id = ?
                  AND id BETWEEN ? - ? AND ? + ?
                  AND id != ?
                ORDER BY id
            """, (row["session_id"], row["id"], context,
                  row["id"], context, row["id"])).fetchall()

            if context_msgs:
                print(f"  --- context ---")
                for ctx in context_msgs:
                    ctx_role = ctx["role"] or "?"
                    ctx_content = ctx["content"] or ""
                    if len(ctx_content) > 150:
                        ctx_content = ctx_content[:150] + "..."
                    print(f"  [{ctx_role}] {ctx_content}")

        print()

    conn.close()


def cmd_sessions(args):
    """List all indexed sessions."""
    db_path = Path(args.db)
    if not db_path.exists():
        print(f"Database not found: {db_path}")
        return

    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row

    sessions = conn.execute("""
        SELECT * FROM sessions ORDER BY first_timestamp DESC
    """).fetchall()

    if not sessions:
        print("No sessions indexed.")
        return

    print(f"{'Session ID':<40} {'Messages':>8} {'User':>6} {'Asst':>6} {'First':>20} {'Last':>20}")
    print("-" * 110)

    for s in sessions:
        sid = s["session_id"][:36]
        first = s["first_timestamp"][:16] if s["first_timestamp"] else "?"
        last = s["last_timestamp"][:16] if s["last_timestamp"] else "?"
        print(f"{sid:<40} {s['message_count']:>8} {s['user_message_count']:>6} "
              f"{s['assistant_message_count']:>6} {first:>20} {last:>20}")

    conn.close()


def main():
    parser = argparse.ArgumentParser(
        description="Bezalel's Memory Well — Search conversation history"
    )
    parser.add_argument(
        "--db", default=str(DB_PATH),
        help=f"Path to SQLite database (default: {DB_PATH})"
    )
    parser.add_argument(
        "--projects-dir", default=str(CLAUDE_PROJECTS_DIR),
        help=f"Path to Claude projects directory (default: {CLAUDE_PROJECTS_DIR})"
    )

    subparsers = parser.add_subparsers(dest="command")

    # index
    subparsers.add_parser("index", help="Index JSONL transcripts")

    # reindex
    subparsers.add_parser("reindex", help="Rebuild database from scratch")

    # search
    search_parser = subparsers.add_parser("search", help="Search conversations")
    search_parser.add_argument("query", help="FTS5 search query")
    search_parser.add_argument(
        "--limit", type=int, default=10,
        help="Max results (default: 10)"
    )
    search_parser.add_argument(
        "--context", type=int, default=0,
        help="Show N surrounding messages for context"
    )
    search_parser.add_argument(
        "--role", choices=["user", "assistant"],
        help="Filter by role"
    )

    # sessions
    subparsers.add_parser("sessions", help="List indexed sessions")

    args = parser.parse_args()

    if args.command == "index":
        cmd_index(args)
    elif args.command == "reindex":
        cmd_reindex(args)
    elif args.command == "search":
        cmd_search(args)
    elif args.command == "sessions":
        cmd_sessions(args)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()