# Copyright (c) 2026 Varun Pratap Bhardwaj / Qualixar
# Licensed under AGPL-3.0-or-later - see LICENSE file
# Part of SuperLocalMemory v3.4.22 — LLD-06 §3

"""AST-extract binary hook entry.

LLD reference: ``.backup/active-brain/lld/LLD-06-windows-binary-and-legacy-migration.md``
Section 3 (Binary Entry — no code duplication).

This generator reads the two canonical source modules
(``src/superlocalmemory/core/topic_signature.py`` and
``src/superlocalmemory/core/context_cache.py``), extracts the functions
the hot-path binary needs (``compute_topic_signature`` and the small
read-only slice of ``read_entry_fast``'s logic), and emits a single
self-contained file that:

    * imports stdlib only (json, sys, os, sqlite3, hashlib, hmac, re,
      time, unicodedata)
    * renames the public functions with a ``_`` prefix
    * ships the ``main()`` contract from LLD-06 §3.2
    * writes ``# source_sha256: <hex>`` so drift is detectable

The output is deterministic: same inputs → byte-identical bytes.
Generated content is meant to be consumed by PyInstaller at release
time (or called during tests to assert the contract). It is NOT
checked into the repo — callers invoke ``emit_entry`` to write it.
"""
from __future__ import annotations

import ast
import hashlib
from dataclasses import dataclass
from pathlib import Path

# stdlib modules the emitted file is allowed to import. Any other import
# in the final file is a bug — the static-check step rejects it.
ALLOWED_STDLIB_IMPORTS: frozenset[str] = frozenset({
    "__future__",
    "json", "sys", "os", "sqlite3", "hashlib", "hmac",
    "re", "time", "unicodedata",
})

# Function names we extract by AST from each source module.
_TOPIC_FN_NAME = "compute_topic_signature"
_READER_FN_NAME = "read_entry_fast"  # referenced for parity docs only

# Deny-list: the emitted file must not reference these packages.
_NONSTDLIB_DENY: tuple[str, ...] = (
    "torch", "sentence_transformers", "fastapi", "uvicorn", "numpy",
    "scipy", "lightgbm", "onnxruntime", "transformers", "httpx",
    "mcp", "pydantic", "lark", "superlocalmemory",
)

BANNER_AUTOGEN = "# AUTO-GENERATED BY scripts/build_entry.py — DO NOT EDIT"


@dataclass(frozen=True, slots=True)
class GenResult:
    """Outcome of running the generator."""

    text: str
    source_sha256: str


# ---------------------------------------------------------------------------
# AST helpers
# ---------------------------------------------------------------------------

def _load_source(path: Path) -> str:
    return path.read_text(encoding="utf-8")


def _sha256_concat(paths: list[Path]) -> str:
    h = hashlib.sha256()
    for p in paths:
        h.update(p.read_bytes())
    return h.hexdigest()


def _extract_function(source: str, name: str) -> ast.FunctionDef:
    """Return the top-level ``def name(...)`` node from ``source``."""
    tree = ast.parse(source)
    for node in tree.body:
        if isinstance(node, ast.FunctionDef) and node.name == name:
            return node
    raise LookupError(f"function {name!r} not found in source")


def _extract_module_constants(
    source: str, names: set[str]
) -> list[ast.stmt]:
    """Extract top-level assignments where a target matches ``names``.

    Captures plain assignments (``X = ...``) and annotated assignments
    (``X: T = ...``). Returns the statements in source order.
    """
    tree = ast.parse(source)
    out: list[ast.stmt] = []
    for node in tree.body:
        if isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id in names:
                    out.append(node)
                    break
        elif isinstance(node, ast.AnnAssign):
            if isinstance(node.target, ast.Name) and node.target.id in names:
                out.append(node)
    return out


# ---------------------------------------------------------------------------
# Emitted main()
# ---------------------------------------------------------------------------

_MAIN_SRC = '''

_TTL_SECONDS = 120
_HMAC_MATERIAL = b"active_brain_cache"
_HMAC_HEX_LEN = 32
_MAX_STDIN_BYTES = 65536
_MAX_PROMPT_CHARS = 8192
_MAX_SESSION_CHARS = 128


def _install_token_for(home_dir):
    token_path = os.path.join(home_dir, ".install_token")
    if not os.path.isfile(token_path):
        return None
    try:
        with open(token_path, "r", encoding="utf-8") as fh:
            token = fh.read().strip()
    except OSError:
        return None
    return token or None


def _binding_hmac(token):
    return hmac.new(
        token.encode("utf-8"),
        _HMAC_MATERIAL,
        hashlib.sha256,
    ).hexdigest()[:_HMAC_HEX_LEN]


def _read_entry_fast(session_id, topic_sig, db_path):
    try:
        home_dir = os.path.dirname(db_path)
        token = _install_token_for(home_dir)
        if token is None:
            return None
        conn = sqlite3.connect(
            "file:" + db_path + "?mode=ro", uri=True, timeout=0.5,
        )
        try:
            row = conn.execute(
                "SELECT value FROM slm_meta WHERE key='install_token_hmac'",
            ).fetchone()
            if row is None:
                return None
            if not hmac.compare_digest(row[0], _binding_hmac(token)):
                return None
            now = int(time.time())
            row = conn.execute(
                "SELECT content FROM context_entries "
                "WHERE session_id=? AND topic_sig=? "
                "  AND computed_at > ?",
                (session_id, topic_sig, now - _TTL_SECONDS),
            ).fetchone()
        finally:
            try:
                conn.close()
            except sqlite3.Error:
                pass
        if row is None:
            return None
        return {"content": row[0]}
    except Exception:
        return None


def main():
    try:
        data = sys.stdin.read(_MAX_STDIN_BYTES)
        payload = json.loads(data) if data else {}
    except (ValueError, UnicodeDecodeError):
        sys.stdout.write("{}")
        return 0

    session_id = str(payload.get("session_id", ""))[:_MAX_SESSION_CHARS]
    prompt = str(payload.get("prompt", ""))[:_MAX_PROMPT_CHARS]
    if not session_id or not prompt:
        sys.stdout.write("{}")
        return 0

    db_path = os.environ.get("SLM_CACHE_DB") or os.path.join(
        os.path.expanduser("~"),
        ".superlocalmemory",
        "active_brain_cache.db",
    )
    if not os.path.isfile(db_path):
        sys.stdout.write("{}")
        return 0

    try:
        sig = _compute_topic_signature(prompt)
        entry = _read_entry_fast(session_id, sig, db_path)
        if entry is None:
            sys.stdout.write("{}")
        else:
            out = {
                "hookSpecificOutput": {
                    "hookEventName": "UserPromptSubmit",
                    "additionalContext": entry["content"],
                }
            }
            sys.stdout.write(json.dumps(out, separators=(",", ":")))
    except Exception:
        sys.stdout.write("{}")
    return 0


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


# ---------------------------------------------------------------------------
# Generator
# ---------------------------------------------------------------------------


def generate(
    topic_signature_src: Path,
    context_cache_src: Path,
) -> GenResult:
    """Produce the emitted hook binary entry as a string.

    Deterministic: given the same input source files, returns identical
    bytes. The ``source_sha256`` banner captures the concatenated SHA
    so silent drift is rejected by the static-check step.
    """
    topic_src = _load_source(topic_signature_src)
    cache_src = _load_source(context_cache_src)

    fn_topic = _extract_function(topic_src, _TOPIC_FN_NAME)
    # Inline helpers/constants referenced by compute_topic_signature.
    topic_consts = _extract_module_constants(
        topic_src,
        {"_STOPWORDS", "MAX_SIG_INPUT_CHARS", "_SIG_LEN",
         "_CAMEL_PASCAL", "_URL", "_PATH", "_QUOTED_DOUBLE",
         "_QUOTED_SINGLE", "_WORD"},
    )
    fn_canon = _extract_function(topic_src, "_canon")

    # Rename compute_topic_signature → _compute_topic_signature, keep
    # call sites inside the function body consistent.
    fn_topic_renamed = _clone_and_rename(fn_topic, "_compute_topic_signature")

    consts_src = _unparse_many(topic_consts)
    fn_canon_src = ast.unparse(fn_canon)
    fn_topic_src = ast.unparse(fn_topic_renamed)

    sha = _sha256_concat([topic_signature_src, context_cache_src])

    pieces: list[str] = [
        BANNER_AUTOGEN,
        f"# source_sha256: {sha}",
        "# Extracted from:",
        f"#   - {topic_signature_src.name} (compute_topic_signature + helpers)",
        f"#   - {context_cache_src.name} (read_entry_fast logic)",
        "",
        "from __future__ import annotations",
        "",
        "import hashlib",
        "import hmac",
        "import json",
        "import os",
        "import re",
        "import sqlite3",
        "import sys",
        "import time",
        "import unicodedata",
        "",
        "# --- extracted constants -----------------------------------------------",
        consts_src,
        "",
        "# --- extracted helpers -------------------------------------------------",
        fn_canon_src,
        "",
        "# --- extracted topic signature ----------------------------------------",
        fn_topic_src,
        "",
        "# --- reader + main() (LLD-06 §3.2) ------------------------------------",
        _MAIN_SRC.strip("\n"),
        "",
    ]
    text = "\n".join(pieces) + "\n"
    return GenResult(text=text, source_sha256=sha)


def _clone_and_rename(fn: ast.FunctionDef, new_name: str) -> ast.FunctionDef:
    new_fn = ast.parse(ast.unparse(fn)).body[0]
    assert isinstance(new_fn, ast.FunctionDef)
    new_fn.name = new_name
    return new_fn


def _unparse_many(nodes: list[ast.stmt]) -> str:
    return "\n".join(ast.unparse(n) for n in nodes)


# ---------------------------------------------------------------------------
# Static checks on the emitted file
# ---------------------------------------------------------------------------


class GeneratorError(RuntimeError):
    """Raised when the emitted file fails a static check."""


def static_check(text: str) -> None:
    """Validate the emitted file. Raises :class:`GeneratorError` on fail.

    Checks:
      1. ``ast.parse`` succeeds.
      2. Every import is from ``ALLOWED_STDLIB_IMPORTS``; no package on
         the non-stdlib deny-list appears at all.
      3. No ``open(...)`` call in write mode anywhere (``w``, ``wb``,
         ``a``, ``a+``, ``w+``, ``x``, ``x+``).
      4. No socket / urllib / http / requests imports (no network).
    """
    try:
        tree = ast.parse(text)
    except SyntaxError as exc:
        raise GeneratorError(f"syntax error in emitted file: {exc}") from exc

    _check_imports(tree)
    _check_no_write_open(tree)
    _check_no_network(tree)


def _check_imports(tree: ast.AST) -> None:
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                root = alias.name.split(".")[0]
                _reject_if_denied(root)
                if root not in ALLOWED_STDLIB_IMPORTS:
                    raise GeneratorError(
                        f"non-stdlib import: {alias.name!r}",
                    )
        elif isinstance(node, ast.ImportFrom):
            module = (node.module or "").split(".")[0]
            _reject_if_denied(module)
            if module and module not in ALLOWED_STDLIB_IMPORTS:
                raise GeneratorError(
                    f"non-stdlib import: from {node.module!r}",
                )


def _reject_if_denied(name: str) -> None:
    if name in _NONSTDLIB_DENY:
        raise GeneratorError(
            f"forbidden non-stdlib package referenced: {name!r}",
        )


def _check_no_write_open(tree: ast.AST) -> None:
    write_mode_chars = {"w", "a", "x", "+"}
    for node in ast.walk(tree):
        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) \
                and node.func.id == "open":
            mode_arg = _open_mode(node)
            if mode_arg is None:
                continue  # default mode 'r' — safe
            for ch in mode_arg:
                if ch in write_mode_chars:
                    raise GeneratorError(
                        f"write-mode open() not allowed: mode={mode_arg!r}",
                    )


def _open_mode(call: ast.Call) -> str | None:
    # Positional: open(path, mode)
    if len(call.args) >= 2:
        arg = call.args[1]
        if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
            return arg.value
    for kw in call.keywords:
        if kw.arg == "mode" and isinstance(kw.value, ast.Constant) \
                and isinstance(kw.value.value, str):
            return kw.value.value
    return None


def _check_no_network(tree: ast.AST) -> None:
    net_roots = {"socket", "urllib", "http", "requests", "httpx",
                 "ftplib", "smtplib", "telnetlib"}
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                root = alias.name.split(".")[0]
                if root in net_roots:
                    raise GeneratorError(
                        f"network import not allowed: {alias.name!r}",
                    )
        elif isinstance(node, ast.ImportFrom):
            root = (node.module or "").split(".")[0]
            if root in net_roots:
                raise GeneratorError(
                    f"network import not allowed: from {node.module!r}",
                )


# ---------------------------------------------------------------------------
# Public CLI-ish entry (invoked from CI build script)
# ---------------------------------------------------------------------------


def emit_entry(
    topic_signature_src: Path,
    context_cache_src: Path,
    dest: Path,
) -> GenResult:
    """Generate, static-check, then write the emitted file to ``dest``."""
    result = generate(topic_signature_src, context_cache_src)
    static_check(result.text)
    dest.parent.mkdir(parents=True, exist_ok=True)
    dest.write_text(result.text, encoding="utf-8")
    return result


def default_source_paths(repo_root: Path) -> tuple[Path, Path]:
    base = repo_root / "src" / "superlocalmemory" / "core"
    return base / "topic_signature.py", base / "context_cache.py"


if __name__ == "__main__":  # pragma: no cover — invoked from CI
    import argparse

    ap = argparse.ArgumentParser(description="AST-extract slm-hook entry")
    ap.add_argument("--repo-root", type=Path, default=Path.cwd())
    ap.add_argument(
        "--dest", type=Path,
        default=None,
        help="output path (default: build/hook_binary_entry.py)",
    )
    ns = ap.parse_args()
    topic, cache = default_source_paths(ns.repo_root)
    dest = ns.dest or (ns.repo_root / "build" / "hook_binary_entry.py")
    res = emit_entry(topic, cache, dest)
    print(f"wrote {dest} (source_sha256={res.source_sha256[:12]}...)")
