# utils/memory_fts_index.py — Derived SQLCipher FTS5 index for workspace memory files.
# Separate from chat search.db; rebuildable from on-disk markdown.
from __future__ import annotations

import hashlib
import logging
import os
import re
import sqlite3
import threading
from typing import Any, Iterable

from utils.fts5_query import safe_fts_query
from utils.memory_scopes import SCOPE_MESH_SAFE, SCOPE_USER_DURABLE, list_mesh_safe_memory_files
from utils.memory_search_fusion import read_memory_file_capped
from utils.memory_shard_layout import list_shard_abs_paths
from utils.path_safety import safe_resolve_path

_log = logging.getLogger(__name__)

try:
    from persistence import MUGI_DIR
except ImportError:
    MUGI_DIR = os.path.expanduser("~/.mugi")

MEMORY_FTS_DB_PATH = os.path.join(MUGI_DIR, "memory_fts.db")

# Keep in sync with tools.memory_tools._MEMORY_SEARCH_DAYS
_MEMORY_SEARCH_SUMMARY_FILES = 14

_conn_lock = threading.Lock()
_rebuild_lock = threading.Lock()
_rebuild_inflight: set[str] = set()
_CORPUS_SUBPATH_RES = (
    re.compile(r"^MEMORY\.md$"),
    re.compile(r"^memory/shards/[^/]+\.md$"),
    re.compile(r"^memory/summaries/\d{4}-\d{2}-\d{2}\.md$"),
    re.compile(r"^memory/archives/MEMORY_ARCHIVE_.*\.md$"),
    re.compile(r"^memory/mesh_safe/[^/]+\.md$"),
)


def _workspace_key(workspace: str) -> str:
    w = os.path.realpath(os.path.expanduser(workspace or ""))
    return hashlib.sha256(w.encode("utf-8")).hexdigest()


def _get_key() -> str:
    try:
        from session_search_db import _get_key as chat_key

        return chat_key()
    except Exception:
        pass
    raise RuntimeError("Cannot retrieve encryption key from Keychain for memory FTS DB.")


def _open_connection() -> sqlite3.Connection | None:
    try:
        ddir = os.path.dirname(MEMORY_FTS_DB_PATH)
        if ddir:
            os.makedirs(ddir, exist_ok=True, mode=0o700)
        try:
            import sqlcipher3

            conn = sqlcipher3.connect(MEMORY_FTS_DB_PATH)
            key = _get_key()
            esc = (key or "").replace("'", "''")
            conn.execute(f"PRAGMA key = '{esc}'")
        except ImportError:
            _log.debug("sqlcipher3 not available; using plain SQLite for memory FTS index.")
            conn = sqlite3.connect(MEMORY_FTS_DB_PATH)
        except Exception as e:
            _log.warning("SQLCipher connection failed for memory FTS (%s); falling back to plain SQLite.", e)
            conn = sqlite3.connect(MEMORY_FTS_DB_PATH)
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute(
            """CREATE VIRTUAL TABLE IF NOT EXISTS memory_fts USING fts5(
                workspace_id UNINDEXED,
                scope UNINDEXED,
                rel_path UNINDEXED,
                is_archive UNINDEXED,
                body,
                tokenize='unicode61'
            )"""
        )
        conn.execute(
            """CREATE TABLE IF NOT EXISTS memory_fts_meta (
                workspace_id TEXT NOT NULL,
                scope TEXT NOT NULL,
                rel_path TEXT NOT NULL,
                mtime_ns INTEGER NOT NULL,
                content_sha256 TEXT,
                PRIMARY KEY (workspace_id, scope, rel_path)
            )"""
        )
        conn.commit()
        return conn
    except Exception as e:
        _log.warning("memory FTS schema init failed (%s)", e)
        return None


def _connection() -> sqlite3.Connection | None:
    return _open_connection()


def _is_indexable_relative_path(rel: str) -> bool:
    s = (rel or "").replace("\\", "/").strip()
    if s == "memory/insights.md":
        return False
    for rx in _CORPUS_SUBPATH_RES:
        if rx.match(s):
            return True
    return False


def collect_legacy_ingest_queue(
    workspace: str,
    scope_eff: str,
    *,
    include_archives: bool,
    mem: dict[str, Any],
) -> list[tuple[str, str]]:
    """Return ordered (rel_path, abs_path) exactly like ``search_memory`` full scan."""
    ws = os.path.expanduser(workspace or "")
    out: list[tuple[str, str]] = []
    max_shards_scan = int(mem.get("memory_search_max_shards_scan") or 200)

    if scope_eff == SCOPE_MESH_SAFE:
        for rel, abs_p in list_mesh_safe_memory_files(ws):
            out.append((rel, abs_p))
        return out

    mem_path, _ = safe_resolve_path(ws, "MEMORY.md", allow_write=False, create_parents=False)
    if mem_path:
        out.append(("MEMORY.md", mem_path))

    shard_paths = sorted(list_shard_abs_paths(ws))[: max(0, max_shards_scan)]
    for sap in shard_paths:
        rel = os.path.relpath(sap, ws).replace("\\", "/")
        out.append((rel, sap))

    summaries_dir_path, _ = safe_resolve_path(ws, "memory/summaries", allow_write=False, create_parents=False)
    if summaries_dir_path and os.path.isdir(summaries_dir_path):
        try:
            files = sorted(
                [f for f in os.listdir(summaries_dir_path) if re.match(r"^\d{4}-\d{2}-\d{2}\.md$", f)],
                reverse=True,
            )[:_MEMORY_SEARCH_SUMMARY_FILES]
        except OSError:
            files = []
        for fname in files:
            rel = f"memory/summaries/{fname}"
            out.append((rel, os.path.join(summaries_dir_path, fname)))

    if include_archives:
        archives_path, _ = safe_resolve_path(ws, "memory/archives", allow_write=False, create_parents=False)
        if archives_path and os.path.isdir(archives_path):
            try:
                arch_files = sorted(
                    [f for f in os.listdir(archives_path) if f.startswith("MEMORY_ARCHIVE_") and f.endswith(".md")],
                    reverse=True,
                )
            except OSError:
                arch_files = []
            for fname in arch_files:
                rel = f"memory/archives/{fname}"
                out.append((rel, os.path.join(archives_path, fname)))

    return out


def collect_all_indexable_files(workspace: str, mem: dict[str, Any]) -> list[tuple[str, str, str, str]]:
    """Rows: (scope, rel_path, abs_path, is_archive '0'|'1') for rebuild (both namespaces)."""
    rows: list[tuple[str, str, str, str]] = []
    ws = os.path.expanduser(workspace or "")

    for rel, abs_p in list_mesh_safe_memory_files(ws):
        rows.append((SCOPE_MESH_SAFE, rel, abs_p, "0"))

    for rel, abs_p in collect_legacy_ingest_queue(
        ws, SCOPE_USER_DURABLE, include_archives=True, mem=mem
    ):
        isa = "1" if rel.startswith("memory/archives/") else "0"
        rows.append((SCOPE_USER_DURABLE, rel, abs_p, isa))
    return rows


def memory_fts_row_count(workspace: str) -> int:
    wid = _workspace_key(workspace)
    with _conn_lock:
        conn = _connection()
        if not conn:
            return 0
        try:
            row = conn.execute("SELECT COUNT(*) FROM memory_fts WHERE workspace_id = ?", (wid,)).fetchone()
            return int(row[0]) if row else 0
        except sqlite3.OperationalError:
            return 0
        finally:
            conn.close()


def delete_memory_fts_document(workspace: str, scope: str, rel_path: str) -> None:
    wid = _workspace_key(workspace)
    rel = rel_path.replace("\\", "/")
    with _conn_lock:
        conn = _connection()
        if not conn:
            return
        try:
            conn.execute(
                "DELETE FROM memory_fts WHERE workspace_id = ? AND scope = ? AND rel_path = ?",
                (wid, scope, rel),
            )
            conn.execute(
                "DELETE FROM memory_fts_meta WHERE workspace_id = ? AND scope = ? AND rel_path = ?",
                (wid, scope, rel),
            )
            conn.commit()
        except sqlite3.OperationalError as e:
            _log.debug("memory FTS delete failed: %s", e)
        finally:
            conn.close()


def upsert_memory_fts_file(
    workspace: str,
    scope: str,
    rel_path: str,
    abs_path: str,
    *,
    mem: dict[str, Any],
) -> None:
    if not _is_indexable_relative_path(rel_path.replace("\\", "/")):
        return
    wid = _workspace_key(workspace)
    rel = rel_path.replace("\\", "/")
    isa = "1" if rel.startswith("memory/archives/") else "0"
    per_file = int(mem.get("memory_search_max_bytes_per_file") or 200_000)

    if not abs_path or not os.path.isfile(abs_path):
        delete_memory_fts_document(workspace, scope, rel)
        return
    try:
        st = os.stat(abs_path)
        mtime_ns = int(getattr(st, "st_mtime_ns", int(st.st_mtime * 1_000_000_000)))
    except OSError:
        delete_memory_fts_document(workspace, scope, rel)
        return

    body, _trunc = read_memory_file_capped(abs_path, per_file)
    sha = hashlib.sha256(body.encode("utf-8", errors="replace")).hexdigest()

    with _conn_lock:
        conn = _connection()
        if not conn:
            return
        try:
            prev = conn.execute(
                "SELECT mtime_ns, content_sha256 FROM memory_fts_meta WHERE workspace_id = ? AND scope = ? AND rel_path = ?",
                (wid, scope, rel),
            ).fetchone()
            if prev and int(prev[0]) == mtime_ns and str(prev[1]) == sha:
                return
            conn.execute(
                "DELETE FROM memory_fts WHERE workspace_id = ? AND scope = ? AND rel_path = ?",
                (wid, scope, rel),
            )
            conn.execute(
                """INSERT INTO memory_fts(workspace_id, scope, rel_path, is_archive, body)
                   VALUES (?,?,?,?,?)""",
                (wid, scope, rel, isa, body),
            )
            conn.execute(
                """INSERT INTO memory_fts_meta(workspace_id, scope, rel_path, mtime_ns, content_sha256)
                   VALUES (?,?,?,?,?)
                   ON CONFLICT(workspace_id, scope, rel_path) DO UPDATE SET
                     mtime_ns=excluded.mtime_ns,
                     content_sha256=excluded.content_sha256""",
                (wid, scope, rel, mtime_ns, sha),
            )
            conn.commit()
        except sqlite3.OperationalError as e:
            _log.warning("memory FTS upsert failed (%s)", e)
            try:
                conn.rollback()
            except Exception:
                pass
        finally:
            try:
                conn.close()
            except Exception:
                pass


def rebuild_memory_fts_index(workspace: str, mem: dict[str, Any] | None = None) -> None:
    if mem is None:
        from core.config import get_memory_config

        mem = get_memory_config()
    wid = _workspace_key(workspace)
    with _conn_lock:
        conn = _open_connection()
        if not conn:
            return
        try:
            conn.execute("DELETE FROM memory_fts WHERE workspace_id = ?", (wid,))
            conn.execute("DELETE FROM memory_fts_meta WHERE workspace_id = ?", (wid,))
            conn.commit()
            for scope, rel, abs_p, isa in collect_all_indexable_files(workspace, mem):
                if not os.path.isfile(abs_p):
                    continue
                body, _ = read_memory_file_capped(abs_p, int(mem.get("memory_search_max_bytes_per_file") or 200_000))
                try:
                    st = os.stat(abs_p)
                    mtime_ns = int(getattr(st, "st_mtime_ns", int(st.st_mtime * 1_000_000_000)))
                except OSError:
                    continue
                sha = hashlib.sha256(body.encode("utf-8", errors="replace")).hexdigest()
                conn.execute(
                    """INSERT INTO memory_fts(workspace_id, scope, rel_path, is_archive, body)
                       VALUES (?,?,?,?,?)""",
                    (wid, scope, rel, isa, body),
                )
                conn.execute(
                    """INSERT INTO memory_fts_meta(workspace_id, scope, rel_path, mtime_ns, content_sha256)
                       VALUES (?,?,?,?,?)""",
                    (wid, scope, rel, mtime_ns, sha),
                )
            conn.commit()
        except Exception:
            _log.exception("memory FTS rebuild aborted")
            try:
                conn.rollback()
            except Exception:
                pass
        finally:
            try:
                conn.close()
            except Exception:
                pass


def schedule_memory_fts_rebuild(workspace: str) -> None:
    wid = _workspace_key(workspace)
    with _rebuild_lock:
        if wid in _rebuild_inflight:
            return
        _rebuild_inflight.add(wid)

    def run() -> None:
        try:
            from core.config import get_memory_config

            rebuild_memory_fts_index(workspace, get_memory_config())
        except Exception:
            _log.exception("background memory FTS rebuild failed")
        finally:
            with _rebuild_lock:
                _rebuild_inflight.discard(wid)

    threading.Thread(target=run, daemon=True, name="memory-fts-rebuild").start()


def search_memory_fts_candidates(
    workspace: str,
    scope_eff: str,
    query: str,
    *,
    include_archives: bool,
    mem: dict[str, Any],
    limit: int,
) -> tuple[list[str], str | None]:
    """Return ordered rel_paths from FTS (best rank first). Second value: error token or None."""
    fts_q = safe_fts_query(
        query,
        mode=str(mem.get("memory_fts_mode") or "and"),
    )
    if not fts_q:
        return [], None
    wid = _workspace_key(workspace)
    inc_arch = 1 if include_archives else 0
    conn = None
    with _conn_lock:
        conn = _connection()
        if not conn:
            return [], "no_conn"
        try:
            rows = conn.execute(
                """
                SELECT rel_path, bm25(memory_fts) AS rnk
                FROM memory_fts
                WHERE memory_fts MATCH ?
                  AND workspace_id = ?
                  AND scope = ?
                  AND (is_archive = '0' OR ? = 1)
                ORDER BY rnk
                LIMIT ?
                """,
                (fts_q, wid, scope_eff, inc_arch, limit),
            ).fetchall()
            return [str(r[0]) for r in rows], None
        except sqlite3.OperationalError as e:
            _log.warning("memory FTS query failed (%s); scheduling workspace rebuild", e)
            schedule_memory_fts_rebuild(workspace)
            return [], "fts_op_error"
        finally:
            if conn is not None:
                try:
                    conn.close()
                except Exception:
                    pass


def notify_memory_paths_changed(workspace: str, paths: Iterable[str], *, mem: dict[str, Any] | None = None) -> None:
    """Best-effort incremental index update after a successful memory write."""
    if mem is None:
        from core.config import get_memory_config

        mem = get_memory_config()
    if not bool(mem.get("memory_fts_enabled", True)):
        return
    ws_abs = os.path.realpath(os.path.expanduser(workspace or ""))
    for raw in paths:
        p = os.path.realpath(os.path.expanduser(str(raw or "")))
        if not p or not os.path.isfile(p):
            continue
        try:
            common = os.path.commonpath([ws_abs, p])
        except ValueError:
            continue
        if common != ws_abs:
            continue
        rel = os.path.relpath(p, ws_abs).replace("\\", "/")
        if rel.startswith(".."):
            continue
        if not _is_indexable_relative_path(rel):
            continue
        scope = SCOPE_MESH_SAFE if rel.startswith("memory/mesh_safe/") else SCOPE_USER_DURABLE
        try:
            upsert_memory_fts_file(workspace, scope, rel, p, mem=mem)
        except Exception:
            _log.debug("memory FTS notify failed for %s", rel, exc_info=True)
