#!/usr/bin/env python3
"""tg — поиск по выгрузке Telegram (группы второго аккаунта владельца) и ежедневная сводка по темам.

  tg search ЗАПРОС [--chat ЧАСТЬ_НАЗВАНИЯ] [--since 30d|2026-01-01] [-n 12]
                      гибридный поиск (слова + смысл) по кускам переписки, со ссылками t.me на сообщения
  tg topic add "ТЕМА" [--query "что искать"] | list | del НОМЕР
                      темы ежедневной сводки (query — поисковая формулировка, по умолчанию — сама тема)
  tg report [--hours 24] [--dry]   сводка по темам за период → владельцу в Telegram (--dry — только напечатать)
  tg stats            сколько чатов/сообщений/кусков, когда обновлялось
  tg index            (сервис tg-export, юзер tgexport) — дописать в индекс новые сообщения выгрузки
  tg serve            (служба tg-embed, юзер astra) — держать модель в памяти, векторы запросов по unix-сокету

Индекс — /mnt/backup/tg/index.db (SQLite: FTS5 + эмбеддинги potion-multilingual-128M), пишет tgexport,
читают все из группы astra. Кусок = подряд идущие сообщения одного чата с паузами < 20 мин (до ~800 знаков).
Темы и состояние сводки — /var/lib/tg-report/ (astra). Сводку пишет Codex (подписка Astra), не Claude.
"""
import argparse
import datetime as dt
import json
import os
import re
import sqlite3
import subprocess
import sys
import tempfile
import time
from pathlib import Path
from zoneinfo import ZoneInfo

OUT = Path(os.environ.get("TG_OUT", "/mnt/backup/tg"))
DB = Path(os.environ.get("TG_INDEX", OUT / "index.db"))
MODELS = os.environ.get("TG_MODELS", "/opt/tg-export/models")
# potion — статические эмбеддинги (model2vec): на Pi 4 в ~16 раз быстрее MiniLM при близком качестве поиска
# (сравнение 2026-10-18 на 8 тыс. сообщений: 57 с против 915 с), без ограничения длины текста
EMB_MODEL = os.environ.get("TG_EMB_MODEL", "minishlab/potion-multilingual-128M")
REPORT = Path(os.environ.get("TG_REPORT_DIR", "/var/lib/tg-report"))
OUTBOX = Path(os.environ.get("ASTRA_OUTBOX", "/var/lib/astra-google/outbox"))
TZ = ZoneInfo("Asia/Nicosia")
GAP = 20 * 60          # пауза, после которой начинается новый кусок
CHUNK_CHARS = int(os.environ.get("TG_CHUNK_CHARS", "800"))
TG_MAX = 3800          # одно сообщение Telegram (лимит 4096) — сводку режем на части
# векторы всех кусков одним файлом рядом с индексом: чтение 85 тыс. BLOB из SQLite на HDD — ~40 с на запрос,
# а np.load(mmap) — доли секунды; пересобирается после каждой индексации
CACHE = DB.with_suffix(".emb.npz")

_emb = None
SOCK = os.environ.get("TG_EMBED_SOCK", "/run/tg-embed/embed.sock")


def _embed_remote(texts):
    """Векторы от службы tg-embed (модель уже в памяти) — иначе None. Загрузка модели в процесс — ~11 с на Pi 4."""
    import socket
    import numpy as np
    if not os.path.exists(SOCK):
        return None
    try:
        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
            s.settimeout(30)
            s.connect(SOCK)
            s.sendall(json.dumps({"texts": texts}).encode() + b"\n")
            buf = b""
            while not buf.endswith(b"\n"):
                part = s.recv(1 << 20)
                if not part:
                    break
                buf += part
        r = json.loads(buf)
        return np.frombuffer(bytes.fromhex(r["vecs"]), dtype=np.float32).reshape(r["n"], r["dim"])
    except (OSError, ValueError, KeyError):
        return None


def embed(texts, remote=True):
    global _emb
    import numpy as np
    if remote and len(texts) <= 64:
        v = _embed_remote(list(texts))
        if v is not None:
            return v
    if _emb is None:
        from fastembed import TextEmbedding
        _emb = TextEmbedding(EMB_MODEL, cache_dir=MODELS, threads=4)
    v = np.array(list(_emb.embed(texts, batch_size=32)), dtype=np.float32)
    v /= np.linalg.norm(v, axis=1, keepdims=True) + 1e-9
    return v


def db(write=False):
    if write:
        DB.parent.mkdir(parents=True, exist_ok=True)
        c = sqlite3.connect(DB, timeout=120)
        # rollback-журнал, не WAL: читателям (astra) не нужна запись в -shm/-wal
        c.execute("PRAGMA journal_mode=DELETE")
        c.executescript("""
        CREATE TABLE IF NOT EXISTS chat(id INTEGER PRIMARY KEY, title TEXT, username TEXT, kind TEXT, folder TEXT, last_id INTEGER DEFAULT 0);
        CREATE TABLE IF NOT EXISTS chunk(id INTEGER PRIMARY KEY, chat_id INTEGER, first_id INTEGER, last_id INTEGER,
            ts0 INTEGER, ts1 INTEGER, n INTEGER, text TEXT, emb BLOB);
        CREATE INDEX IF NOT EXISTS chunk_chat ON chunk(chat_id, last_id);
        CREATE INDEX IF NOT EXISTS chunk_ts ON chunk(ts1);
        CREATE VIRTUAL TABLE IF NOT EXISTS chunk_fts USING fts5(text, content='chunk', content_rowid='id',
            tokenize='unicode61 remove_diacritics 2');
        CREATE TRIGGER IF NOT EXISTS chunk_ai AFTER INSERT ON chunk BEGIN
            INSERT INTO chunk_fts(rowid, text) VALUES (new.id, new.text); END;
        CREATE TRIGGER IF NOT EXISTS chunk_ad AFTER DELETE ON chunk BEGIN
            INSERT INTO chunk_fts(chunk_fts, rowid, text) VALUES ('delete', old.id, old.text); END;
        CREATE TABLE IF NOT EXISTS meta(k TEXT PRIMARY KEY, v TEXT);
        """)
    else:
        if not DB.exists():
            raise SystemExit("индекса ещё нет — выгрузка Telegram не запускалась (tg-export)")
        c = sqlite3.connect(f"file:{DB}?mode=ro", uri=True, timeout=120)
    c.row_factory = sqlite3.Row
    return c


# ---------------------------------------------------------------- индекс

def _ts(s):
    return int(dt.datetime.fromisoformat(s).timestamp())


def _line(m):
    text = (m.get("text") or "").strip()
    if not text:
        return None
    who = m.get("from") or "?"
    return f"{who}: {re.sub(r'\s+', ' ', text)}"


def _windows(msgs):
    """Сообщения (по возрастанию id) → куски разговора."""
    cur, out = [], []
    for m in msgs:
        line = _line(m)
        if not line:
            continue
        ts = _ts(m["date"])
        if cur and (ts - cur[-1][1] > GAP or sum(len(x[2]) for x in cur) + len(line) > CHUNK_CHARS):
            out.append(cur)
            cur = []
        cur.append((m["id"], ts, line))
    if cur:
        out.append(cur)
    return out


def cmd_index(a):
    c = db(write=True)
    chats = sorted(p for p in (OUT / "chats").glob("*") if (p / "_pretty.jsonl").exists())
    total = 0
    for folder in chats:
        meta = json.loads((folder / "meta.json").read_text()) if (folder / "meta.json").exists() else {}
        cid = int(meta.get("id") or folder.name.split("_")[0])
        row = c.execute("SELECT last_id FROM chat WHERE id=?", (cid,)).fetchone()
        last = row["last_id"] if row else 0
        c.execute("INSERT INTO chat(id, title, username, kind, folder, last_id) VALUES (?,?,?,?,?,?) "
                  "ON CONFLICT(id) DO UPDATE SET title=excluded.title, username=excluded.username, kind=excluded.kind, "
                  "folder=excluded.folder", (cid, meta.get("name") or folder.name, meta.get("username"), meta.get("kind"),
                                             folder.name, last))
        # последний кусок мог быть недописан (разговор продолжился) — пересобираем его вместе с новыми сообщениями
        tail = c.execute("SELECT id, first_id, ts1 FROM chunk WHERE chat_id=? ORDER BY last_id DESC LIMIT 1", (cid,)).fetchone()
        msgs = []
        with (folder / "_pretty.jsonl").open(encoding="utf-8") as f:
            floor = tail["first_id"] - 1 if tail else last
            for raw in f:
                mid = re.match(r'\{"id": (\d+)', raw)          # дёшево отсеять уже проиндексированное без разбора JSON
                if not mid or int(mid.group(1)) <= floor:
                    continue
                m = json.loads(raw)
                if m.get("date"):
                    msgs.append(m)
        msgs.sort(key=lambda m: m["id"])
        new = [m for m in msgs if m["id"] > last]
        if not new:
            continue
        if tail and _ts(new[0]["date"]) - tail["ts1"] <= GAP:
            c.execute("DELETE FROM chunk WHERE id=?", (tail["id"],))
        else:
            msgs = new
        wins = _windows(msgs)
        title = meta.get("name") or folder.name
        for i in range(0, len(wins), 256):
            part = wins[i:i + 256]
            texts = [f"Чат: {title}\n" + "\n".join(x[2] for x in w) for w in part]
            vecs = embed(texts)
            c.executemany("INSERT INTO chunk(chat_id, first_id, last_id, ts0, ts1, n, text, emb) VALUES (?,?,?,?,?,?,?,?)",
                          [(cid, w[0][0], w[-1][0], w[0][1], w[-1][1], len(w), "\n".join(x[2] for x in w), v.tobytes())
                           for w, v in zip(part, vecs)])
            c.commit()
            if not a.quiet:
                print(f"  {title}: {min(i + 256, len(wins))}/{len(wins)} кусков", flush=True)
        c.execute("UPDATE chat SET last_id=? WHERE id=?", (max(m["id"] for m in new), cid))
        c.commit()
        total += len(new)
    c.execute("INSERT OR REPLACE INTO meta VALUES ('indexed', ?)", (dt.datetime.now(TZ).isoformat(timespec="seconds"),))
    c.commit()
    if total or not CACHE.exists():
        build_cache(c)
    print(f"проиндексировано новых сообщений: {total}")


def build_cache(c):
    import numpy as np
    rows = c.execute("SELECT id, chat_id, ts1, emb FROM chunk ORDER BY id").fetchall()
    if not rows:
        return
    tmp = CACHE.with_name(CACHE.name + ".tmp.npz")
    np.savez(tmp, ids=np.array([r["id"] for r in rows], dtype=np.int64),
             chat=np.array([r["chat_id"] for r in rows], dtype=np.int64),
             ts=np.array([r["ts1"] for r in rows], dtype=np.int64),
             emb=np.frombuffer(b"".join(r["emb"] for r in rows), dtype=np.float32).reshape(len(rows), -1))
    os.replace(tmp, CACHE)


# ---------------------------------------------------------------- поиск

def link(chat, msg_id):
    if chat["username"]:
        return f"https://t.me/{chat['username']}/{msg_id}"
    cid = str(chat["id"])
    return f"https://t.me/c/{cid[4:] if cid.startswith('-100') else cid.lstrip('-')}/{msg_id}"


def fts_query(q):
    words = [re.sub(r"[^\w]", "", w) for w in q.split()]
    out = []
    for w in words:
        if len(w) < 3:
            continue
        stem = re.sub(r"(ами|ями|ого|его|ому|ему|ой|ей|ий|ый|ая|яя|ое|ее|ые|ие|ам|ям|ах|ях|ом|ем|ов|ев|а|я|ы|и|у|ю|е|о)$", "", w) \
            if re.fullmatch(r"[а-яё]{5,}", w.lower()) else w
        out.append(f'"{stem}"*')
    return " OR ".join(out)


def parse_since(s):
    if not s:
        return None
    m = re.fullmatch(r"(\d+)\s*(d|д|w|н|m|м)", s.strip().lower())
    if m:
        days = int(m.group(1)) * {"d": 1, "д": 1, "w": 7, "н": 7, "m": 30, "м": 30}[m.group(2)]
        return int(time.time()) - days * 86400
    return int(dt.datetime.combine(dt.date.fromisoformat(s), dt.time(0), TZ).timestamp())


def search(q, chat=None, since=None, n=12, min_sim=None):
    """Гибрид: BM25 (FTS5) + косинус эмбеддингов, слияние по рангам (RRF). → [(row, sim)]"""
    import numpy as np
    c = db()
    where, args = ["1=1"], []
    if chat:
        where.append("c.chat_id IN (SELECT id FROM chat WHERE title LIKE ? OR username LIKE ?)")
        args += [f"%{chat}%", f"%{chat}%"]
    if since:
        where.append("c.ts1 >= ?")
        args.append(since)
    w = " AND ".join(where)
    ranks, sims, by_words = {}, {}, set()
    fq = fts_query(q)
    if fq:
        for r, row in enumerate(c.execute(f"SELECT c.id FROM chunk_fts f JOIN chunk c ON c.id=f.rowid WHERE chunk_fts MATCH ? "
                                          f"AND {w} ORDER BY bm25(chunk_fts) LIMIT 60", [fq] + args)):
            ranks[row["id"]] = ranks.get(row["id"], 0) + 1 / (60 + r)
            by_words.add(row["id"])
    qv = embed([q])[0]
    if CACHE.exists():
        z = np.load(CACHE)
        ids, emb, mask = z["ids"], z["emb"], np.ones(len(z["ids"]), dtype=bool)
        if since:
            mask &= z["ts"] >= since
        if chat:
            cids = [r[0] for r in c.execute("SELECT id FROM chat WHERE title LIKE ? OR username LIKE ?", (f"%{chat}%", f"%{chat}%"))]
            mask &= np.isin(z["chat"], cids)
        ids, emb = ids[mask], emb[mask]
    else:                                  # кэша нет (старый индекс) — медленно, из базы
        rows = c.execute(f"SELECT c.id, c.emb FROM chunk c WHERE {w}", args).fetchall()
        ids = np.array([r["id"] for r in rows], dtype=np.int64)
        emb = np.frombuffer(b"".join(r["emb"] for r in rows), dtype=np.float32).reshape(len(rows), -1) if rows else None
    if len(ids):
        s = emb @ qv
        top_i = np.argpartition(-s, min(60, len(s) - 1))[:60]
        for r, i in enumerate(top_i[np.argsort(-s[top_i])]):
            sims[int(ids[i])] = float(s[i])
            ranks[int(ids[i])] = ranks.get(int(ids[i]), 0) + 1 / (60 + r)
    top = sorted(ranks, key=ranks.get, reverse=True)
    if min_sim is not None:      # для сводки: найденное словами — всегда, по смыслу — только достаточно близкое
        top = [i for i in top if i in by_words or sims.get(i, 0) >= min_sim]
    top = top[:n]
    if not top:
        return []
    rows = {r["id"]: r for r in c.execute(
        f"SELECT c.*, ch.title, ch.username, ch.id AS cid FROM chunk c JOIN chat ch ON ch.id=c.chat_id "
        f"WHERE c.id IN ({','.join('?' * len(top))})", top)}
    return [(rows[i], sims.get(i)) for i in top if i in rows]


def fmt_hit(r, sim=None, cap=900):
    chat = {"id": r["cid"], "username": r["username"]}
    t0 = dt.datetime.fromtimestamp(r["ts0"], TZ)
    text = r["text"] if len(r["text"]) <= cap else r["text"][:cap] + " …"
    return (f"── {r['title']} · {t0:%d.%m.%Y %H:%M} · {link(chat, r['first_id'])}"
            + (f" · близость {sim:.2f}" if sim is not None else "") + f"\n{text}\n")


def cmd_search(a):
    hits = search(" ".join(a.query), a.chat, parse_since(a.since), a.n)
    if not hits:
        print("Ничего не нашлось.")
    for r, sim in hits:
        print(fmt_hit(r, sim))


def cmd_stats(a):
    c = db()
    chats = c.execute("SELECT count(*) FROM chat").fetchone()[0]
    n, msgs, last = c.execute("SELECT count(*), coalesce(sum(n),0), max(ts1) FROM chunk").fetchone()
    upd = c.execute("SELECT v FROM meta WHERE k='indexed'").fetchone()
    print(f"чатов {chats}, кусков {n}, сообщений с текстом {msgs}; последнее сообщение "
          f"{dt.datetime.fromtimestamp(last, TZ):%d.%m.%Y %H:%M}" if last else "пусто",
          f"; индекс обновлён {upd[0] if upd else '—'}")
    for r in c.execute("SELECT ch.title, count(k.id) n, max(k.ts1) last FROM chat ch LEFT JOIN chunk k ON k.chat_id=ch.id "
                       "GROUP BY ch.id ORDER BY n DESC LIMIT 30"):
        print(f"  {r['title'][:50]:<50} кусков {r['n']:>6}"
              + (f", последнее {dt.datetime.fromtimestamp(r['last'], TZ):%d.%m %H:%M}" if r["last"] else ""))


# ---------------------------------------------------------------- темы и сводка

def topics():
    f = REPORT / "topics.json"
    return json.loads(f.read_text()) if f.exists() else []


def save_topics(t):
    REPORT.mkdir(parents=True, exist_ok=True)
    (REPORT / "topics.json").write_text(json.dumps(t, ensure_ascii=False, indent=1))


def cmd_topic(a):
    t = topics()
    if a.op == "add":
        t.append({"title": a.title.strip(), "query": (a.query or a.title).strip(), "added": dt.date.today().isoformat()})
        save_topics(t)
        print(f"Тема {len(t)}: {a.title}")
    elif a.op == "del":
        i = int(a.num) - 1
        if not 0 <= i < len(t):
            raise SystemExit("нет такой темы")
        print("Удалена:", t.pop(i)["title"])
        save_topics(t)
    else:
        for i, x in enumerate(t, 1):
            print(f"{i}. {x['title']}" + (f"  (ищу: {x['query']})" if x["query"] != x["title"] else ""))
        if not t:
            print("Тем нет. Добавить: tg topic add \"тема\"")


def codex(prompt):
    with tempfile.TemporaryDirectory() as d:
        out = Path(d) / "out.txt"
        subprocess.run(["codex", "exec", "--skip-git-repo-check", "--sandbox", "read-only", "--ephemeral",
                        "--color", "never", "-o", str(out), "-"], input=prompt, text=True, cwd=d,
                       capture_output=True, timeout=900, check=True)
        return out.read_text().strip()


REPORT_PROMPT = """Ты готовишь владельцу утреннюю сводку по Telegram-группам за последние {hours} ч по его темам.
Ниже по каждой теме — найденные поиском куски переписки (группа, время, ссылка, текст). Поиск грубый: часть кусков
к теме не относится — их пропусти.

Для каждой темы, где есть действительно относящееся: 2–6 пунктов по сути (что предлагают/советуют/спрашивают,
конкретика — места, цены, контакты, даты), у каждого пункта — ссылка на сообщение из данных. Ничего не выдумывай
и не добавляй от себя. Тему без относящегося — одной строкой «ничего нового». Без вступлений и выводов.
Формат — простой текст для Telegram без Markdown: заголовок темы строкой «▸ Тема», пункты — «• …».
Пиши по-русски.

{data}"""


def cmd_report(a):
    t = topics()
    if not t:
        print("Тем нет — сводку не делаю.")
        return
    if not DB.exists():
        print("Индекса ещё нет (выгрузка не запускалась) — сводку не делаю.")
        return
    since = int(time.time()) - a.hours * 3600
    blocks = []
    for x in t:
        hits = search(x["query"], since=since, n=15, min_sim=0.4)
        body = "\n".join(fmt_hit(r, cap=700) for r, _ in hits) or "(ничего не найдено)"
        blocks.append(f"### Тема: {x['title']}\n{body}")
    text = codex(REPORT_PROMPT.format(hours=a.hours, data="\n\n".join(blocks)))
    text = f"🗞 Telegram за {a.hours} ч — {dt.datetime.now(TZ):%d.%m}\n\n{text}"
    if a.dry:
        print(text)
        return
    OUTBOX.mkdir(parents=True, exist_ok=True)
    parts, cur = [], ""
    for para in text.split("\n"):
        if len(cur) + len(para) + 1 > TG_MAX and cur:
            parts.append(cur)
            cur = ""
        cur += para[:TG_MAX] + "\n"
    if cur.strip():
        parts.append(cur)
    old = os.umask(0o007)
    try:
        stamp = dt.datetime.now().strftime("%Y%m%d%H%M%S%f")
        for i, p in enumerate(parts):
            (OUTBOX / f"{stamp}-tg{i:02d}.txt").write_text(p.strip())
    finally:
        os.umask(old)
    REPORT.mkdir(parents=True, exist_ok=True)
    (REPORT / "last-report.txt").write_text(text)
    print(f"сводка поставлена в очередь: {len(parts)} сообщ.")


def cmd_serve(a):
    """Служба tg-embed: модель в памяти, векторы запросов по unix-сокету (одна строка JSON → одна строка JSON)."""
    import socketserver
    embed(["прогрев"], remote=False)

    class H(socketserver.StreamRequestHandler):
        def handle(self):
            try:
                req = json.loads(self.rfile.readline())
                v = embed(req["texts"][:64], remote=False)
                out = {"n": v.shape[0], "dim": v.shape[1], "vecs": v.astype("float32").tobytes().hex()}
            except Exception as e:        # плохой запрос не роняет службу
                out = {"error": str(e)}
            self.wfile.write(json.dumps(out).encode() + b"\n")

    if os.path.exists(SOCK):
        os.unlink(SOCK)
    old = os.umask(0o077)                 # сокет — только своему пользователю (astra)
    srv = socketserver.UnixStreamServer(SOCK, H)
    os.umask(old)
    print(f"tg-embed: {EMB_MODEL} в памяти, сокет {SOCK}", flush=True)
    srv.serve_forever()


def main():
    p = argparse.ArgumentParser(prog="tg", description=__doc__.split("\n")[0])
    sub = p.add_subparsers(dest="cmd", required=True)
    s = sub.add_parser("search"); s.add_argument("query", nargs="+"); s.add_argument("--chat"); s.add_argument("--since")
    s.add_argument("-n", type=int, default=12)
    s = sub.add_parser("topic"); ts = s.add_subparsers(dest="op", required=True)
    r = ts.add_parser("add"); r.add_argument("title"); r.add_argument("--query")
    r = ts.add_parser("del"); r.add_argument("num")
    ts.add_parser("list")
    s = sub.add_parser("report"); s.add_argument("--hours", type=int, default=24); s.add_argument("--dry", action="store_true")
    sub.add_parser("stats")
    s = sub.add_parser("index"); s.add_argument("--quiet", action="store_true")
    sub.add_parser("serve")
    a = p.parse_args()
    globals()["cmd_" + a.cmd](a)


if __name__ == "__main__":
    main()