"""Библиотека: разбор PDF/EPUB на главы, хранение в SQLite, гибридный поиск (FTS5 + эмбеддинги)."""
import datetime as dt
from html.parser import HTMLParser
import os
import re
import sqlite3
import threading
from pathlib import Path

import fitz  # PyMuPDF: PDF и EPUB одним API

fitz.TOOLS.mupdf_display_errors(False)   # шум про шрифты OCR-слоя в stderr — не ошибки разбора
import numpy as np

DATA = Path(os.environ.get("BOOKS_DATA", "/data"))
FILES = DATA / "files"
DB_PATH = DATA / "books.db"
EMB_MODEL = os.environ.get("EMB_MODEL", "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
CHUNK = 1200          # символов на фрагмент поиска
OVERLAP = 200
NO_TOC_PAGES = 20     # без оглавления — режем по 20 страниц

_local = threading.local()
_emb_lock = threading.Lock()
_emb = None


def db():
    c = getattr(_local, "c", None)
    if c is None:
        DATA.mkdir(parents=True, exist_ok=True)
        c = sqlite3.connect(DB_PATH, timeout=60, check_same_thread=False)
        c.row_factory = sqlite3.Row
        c.execute("PRAGMA journal_mode=WAL")
        c.execute("PRAGMA foreign_keys=ON")
        c.executescript(SCHEMA)
        _local.c = c
    return c


SCHEMA = """
CREATE TABLE IF NOT EXISTS book(id INTEGER PRIMARY KEY, title TEXT, author TEXT, lang TEXT, file TEXT, orig TEXT, fmt TEXT,
    pages INTEGER, added TEXT, status TEXT, note TEXT, last_ch INTEGER);
CREATE TABLE IF NOT EXISTS chapter(id INTEGER PRIMARY KEY, book_id INTEGER REFERENCES book(id) ON DELETE CASCADE,
    idx INTEGER, title TEXT, level INTEGER, p1 INTEGER, p2 INTEGER, text TEXT, chars INTEGER,
    read INTEGER DEFAULT 0, read_at TEXT, UNIQUE(book_id, idx));
CREATE TABLE IF NOT EXISTS recap(chapter_id INTEGER PRIMARY KEY REFERENCES chapter(id) ON DELETE CASCADE,
    md TEXT, ai_md TEXT, updated TEXT);
CREATE TABLE IF NOT EXISTS chunk(id INTEGER PRIMARY KEY, book_id INTEGER REFERENCES book(id) ON DELETE CASCADE,
    chapter_id INTEGER REFERENCES chapter(id) ON DELETE CASCADE, kind TEXT, page INTEGER, text TEXT, emb BLOB);
CREATE INDEX IF NOT EXISTS chunk_book ON chunk(book_id, kind);
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 qa(id INTEGER PRIMARY KEY, book_id INTEGER REFERENCES book(id) ON DELETE CASCADE,
    chapter_idx INTEGER, kind TEXT, q TEXT, a TEXT, ts TEXT);
"""


def now():
    return dt.datetime.now().isoformat(timespec="seconds")


# ---------------------------------------------------------------- эмбеддинги

def embedder():
    global _emb
    with _emb_lock:
        if _emb is None:
            from fastembed import TextEmbedding
            _emb = TextEmbedding(EMB_MODEL, cache_dir=str(DATA / "models"), threads=2)
        return _emb


def embed(texts):
    if not texts:
        return np.zeros((0, 384), dtype=np.float32)
    v = np.array(list(embedder().embed(texts, batch_size=32)), dtype=np.float32)
    v /= np.linalg.norm(v, axis=1, keepdims=True) + 1e-9
    return v


# ---------------------------------------------------------------- разбор книги

def detect_lang(text):
    cyr = len(re.findall(r"[а-яёА-ЯЁ]", text[:20000]))
    lat = len(re.findall(r"[a-zA-Z]", text[:20000]))
    return "ru" if cyr > lat else "en"


def clean(t):
    t = t.replace("­", "")                       # мягкие переносы
    t = re.sub(r"(\w)-\n(\w)", r"\1\2", t)            # перенос слова по слогам в PDF
    t = re.sub(r"[ \t]+\n", "\n", t)
    t = re.sub(r"\n{3,}", "\n\n", t)
    return t.strip()


CHAPTER_RE = re.compile(r"^\s*(глава|chapter|часть|part|раздел)\s+([0-9IVXLC]+|[а-яё]+)\b", re.I)
MAX_AVG_PAGES = 45    # «главы» длиннее в среднем — это части книги, берём уровень глубже


def _num(tok):
    tok = tok.upper()
    if tok.isdigit():
        return int(tok)
    roman = {"I": 1, "V": 5, "X": 10, "L": 50, "C": 100}
    if tok and all(ch in roman for ch in tok):
        total = 0
        for a, b in zip(tok, tok[1:] + " "):
            total += -roman[a] if b in roman and roman[b] > roman[a] else roman[a]
        return total
    return None


def headings_from_pages(texts):
    """Без оглавления: страницы, которые начинаются с «Глава N» / «Chapter N».

    Колонтитулы повторяют «Глава N» на каждой странице, а страницы оглавления перечисляют все главы подряд, поэтому:
    берём первое появление каждого номера и самую длинную возрастающую по номеру цепочку (по порядку страниц)."""
    found, seen = [], set()
    for i, t in enumerate(texts):
        head = [l.strip() for l in t.splitlines()[:4] if l.strip()]
        for j, l in enumerate(head[:3]):
            m = CHAPTER_RE.match(l)
            if m and len(l) < 90:
                kind, n = m.group(1).lower(), _num(m.group(2))
                if kind not in ("глава", "chapter") or n is None or n in seen:
                    break
                seen.add(n)
                title = l if len(l) > 12 or j + 1 >= len(head) else f"{l}. {head[j + 1][:70]}"
                title = re.sub(r"\.{2,}", ".", title)
                found.append((n, i + 1, title))
                break
    # самая длинная возрастающая подпоследовательность по номеру главы (O(n²), глав немного)
    best, prev = [1] * len(found), [-1] * len(found)
    for a in range(len(found)):
        for b in range(a):
            if found[b][0] < found[a][0] and best[b] + 1 > best[a]:
                best[a], prev[a] = best[b] + 1, b
    if not found:
        return []
    k = max(range(len(found)), key=best.__getitem__)
    chain = []
    while k >= 0:
        chain.append(found[k])
        k = prev[k]
    chain.reverse()
    return [[1, title, page] for _, page, title in chain] if len(chain) >= 3 else []


def plan_chapters(toc, pages, texts=None):
    """Из оглавления — главы (верхний уровень, если он не из огромных «частей»). [(title, level, p1, p2)]."""
    toc = [(lvl, (title or "").strip(), max(1, min(pages, p))) for lvl, title, p, *_ in toc if p and p > 0]
    if not toc and texts:
        toc = [tuple(x) for x in headings_from_pages(texts)]
    if toc:
        lvl1 = [t for t in toc if t[0] == 1]
        lvl2 = [t for t in toc if t[0] <= 2]
        use = lvl1 if len(lvl1) >= 3 and (pages / len(lvl1) <= MAX_AVG_PAGES or len(lvl2) <= len(lvl1)) else lvl2
        use = sorted({(p, i): (lvl, title, p) for i, (lvl, title, p) in enumerate(use)}.items())
        use = [v for _, v in use]
        out = []
        if use and use[0][2] > 1:
            out.append(("Начало книги", 1, 1, use[0][2] - 1))
        for i, (lvl, title, p) in enumerate(use):
            p2 = use[i + 1][2] - 1 if i + 1 < len(use) else pages
            out.append((title or f"Часть {i + 1}", lvl, p, max(p, p2)))
        # слишком мелкие подряд (одна страница) оставляем как есть — это видно в оглавлении
        return out
    return [(f"Страницы {p}–{min(pages, p + NO_TOC_PAGES - 1)}", 1, p, min(pages, p + NO_TOC_PAGES - 1))
            for p in range(1, pages + 1, NO_TOC_PAGES)]


# ---------------------------------------------------------------- EPUB: свой разбор (у PyMuPDF нет оглавления EPUB 3)

EPUB_PAGE = 2000      # у EPUB нет страниц — «страница» ≈ 2000 символов, для ссылок [с. N]


class _Text(HTMLParser):
    BLOCK = {"p", "div", "br", "li", "h1", "h2", "h3", "h4", "h5", "h6", "tr", "pre", "blockquote", "section", "dt", "dd"}

    def __init__(self):
        super().__init__(convert_charrefs=True)
        self.out, self.skip = [], 0

    def handle_starttag(self, tag, attrs):
        if tag in ("script", "style", "head"):
            self.skip += 1
        elif tag in self.BLOCK:
            self.out.append("\n\n")

    def handle_endtag(self, tag):
        if tag in ("script", "style", "head"):
            self.skip = max(0, self.skip - 1)
        elif tag in self.BLOCK:
            self.out.append("\n\n")

    def handle_data(self, d):
        if not self.skip:
            self.out.append(d)


def html_text(raw):
    t = _Text()
    t.feed(raw)
    return clean(re.sub(r"[ \t\r\f\v]*\n[ \t\r\f\v]*", "\n", re.sub(r"[ \t]+", " ", "".join(t.out))))


def parse_epub(path):
    """→ (title, author, [(глава, уровень, [тексты «страниц»])]). Главы — по оглавлению (nav/NCX), текст — по spine."""
    import posixpath
    import urllib.parse
    import xml.etree.ElementTree as ET
    import zipfile
    z = zipfile.ZipFile(path)
    tag = lambda e: e.tag.rsplit("}", 1)[-1]
    cont = ET.fromstring(z.read("META-INF/container.xml"))
    opf_path = next(e.get("full-path") for e in cont.iter() if tag(e) == "rootfile")
    base = posixpath.dirname(opf_path)
    norm = lambda href, rel=base: posixpath.normpath(posixpath.join(rel, urllib.parse.unquote(href.split("#")[0])))
    opf = ET.fromstring(z.read(opf_path))
    meta = {tag(e): (e.text or "").strip() for e in opf.iter() if tag(e) in ("title", "creator") and (e.text or "").strip()}
    manifest, nav_href, ncx_href = {}, None, None
    for e in opf.iter():
        if tag(e) == "item":
            manifest[e.get("id")] = e.get("href")
            if "nav" in (e.get("properties") or "").split():
                nav_href = e.get("href")
            if e.get("media-type") == "application/x-dtbncx+xml":
                ncx_href = e.get("href")
    spine = [norm(manifest[e.get("idref")]) for e in opf.iter() if tag(e) == "itemref" and e.get("idref") in manifest]
    titles = {}                                       # файл spine → (название, уровень) первой записи оглавления
    def add(href, title, lvl, rel):
        f = norm(href, rel)
        if title and f not in titles:
            titles[f] = (title.strip(), lvl)
    try:
        if nav_href:
            rel = posixpath.dirname(norm(nav_href))
            nav = ET.fromstring(z.read(norm(nav_href)))
            def walk(ol, lvl):
                for li in ol:
                    if tag(li) != "li":
                        continue
                    a = next((x for x in li.iter() if tag(x) == "a"), None)
                    if a is not None and a.get("href"):
                        add(a.get("href"), "".join(a.itertext()), lvl, rel)
                    for sub in li:
                        if tag(sub) == "ol":
                            walk(sub, lvl + 1)
            for n in nav.iter():
                if tag(n) == "nav" and "toc" in (n.get("{http://www.idpf.org/2007/ops}type") or "toc"):
                    for ol in n:
                        if tag(ol) == "ol":
                            walk(ol, 1)
                    break
        if not titles and ncx_href:
            rel = posixpath.dirname(norm(ncx_href))
            ncx = ET.fromstring(z.read(norm(ncx_href)))
            def walkp(node, lvl):
                for np_ in node:
                    if tag(np_) != "navPoint":
                        continue
                    label = next(("".join(x.itertext()) for x in np_.iter() if tag(x) == "text"), "")
                    src = next((x.get("src") for x in np_ if tag(x) == "content"), None)
                    if src:
                        add(src, label, lvl, rel)
                    walkp(np_, lvl + 1)
            for n in ncx.iter():
                if tag(n) == "navMap":
                    walkp(n, 1)
    except (KeyError, ET.ParseError):
        titles = {}
    # верхний уровень, если его достаточно — иначе первые два
    lvl1 = [f for f, (t, l) in titles.items() if l == 1]
    keep_lvl = 1 if len(lvl1) >= 3 else 2
    chapters = []
    for f in spine:
        try:
            text = html_text(z.read(f).decode("utf-8", "replace"))
        except KeyError:
            continue
        t = titles.get(f)
        if not chapters or (t and t[1] <= keep_lvl) or (not titles and text):
            first = next((l for l in text.splitlines() if l.strip()), "")[:80]
            chapters.append([t[0] if t else (first or f"Часть {len(chapters) + 1}"), t[1] if t else 1, []])
        if text:
            chapters[-1][2].append(text)
    out = []
    for title, lvl, texts in chapters:
        body = "\n\n".join(texts)
        if not body.strip():
            continue
        pages, buf = [], ""
        for para in re.split(r"\n\s*\n", body):
            if buf and len(buf) + len(para) > EPUB_PAGE:
                pages.append(buf)
                buf = ""
            buf = (buf + "\n\n" + para) if buf else para
        if buf:
            pages.append(buf)
        out.append((title, lvl, pages))
    return meta.get("title", ""), meta.get("creator", ""), out


def chunks_of(page_texts, p1):
    """Фрагменты для поиска с номером страницы начала."""
    out = []
    buf, start = "", p1
    for i, t in enumerate(page_texts):
        page = p1 + i
        for para in re.split(r"\n\s*\n", t):
            para = para.strip()
            if not para:
                continue
            if not buf:
                start = page
            if len(buf) + len(para) > CHUNK and buf:
                out.append((start, buf))
                buf = buf[-OVERLAP:] + "\n" + para
                start = page
            else:
                buf = (buf + "\n\n" + para) if buf else para
    if buf.strip():
        out.append((start, buf))
    return out


def import_book(book_id):
    """Фоновая обработка загруженного файла: главы, текст, фрагменты, эмбеддинги."""
    c = db()
    b = c.execute("SELECT * FROM book WHERE id=?", (book_id,)).fetchone()
    try:
        if b["fmt"] == "epub":
            et, ea, eps = parse_epub(FILES / b["file"])
            texts, chapters = [], []
            for ctitle, lvl, ptexts in eps:
                chapters.append((ctitle, lvl, len(texts) + 1, len(texts) + len(ptexts)))
                texts += ptexts
            pages = len(texts)
            meta = {"title": et, "author": ea}
        else:
            doc = fitz.open(FILES / b["file"])
            if doc.needs_pass:
                raise ValueError("файл защищён паролем")
            pages = doc.page_count
            texts = [clean(doc[i].get_text("text")) for i in range(pages)]
            meta = doc.metadata or {}
            chapters = plan_chapters(doc.get_toc(simple=True), pages, texts)
        total = sum(len(t) for t in texts)
        mt = (meta.get("title") or "").strip()
        title = b["title"] or (mt if len(mt) > 2 and not mt.lower().startswith(("untitled", "microsoft")) else "") \
            or Path(b["orig"] or b["file"]).stem
        author = (meta.get("author") or "").strip()
        note = ""
        if total < pages * 80:
            note = "Похоже на скан: текста почти нет (распознавание пока не поддерживается)."
        if not pages or not chapters:
            raise ValueError("в файле не нашлось текста")
        c.execute("DELETE FROM chapter WHERE book_id=?", (book_id,))
        all_chunks = []
        for idx, (ctitle, lvl, p1, p2) in enumerate(chapters, 1):
            ptexts = texts[p1 - 1:p2]
            body = "\n\n".join(f"[[p{p1 + i}]]\n{t}" for i, t in enumerate(ptexts))
            cur = c.execute("INSERT INTO chapter(book_id, idx, title, level, p1, p2, text, chars) VALUES (?,?,?,?,?,?,?,?)",
                            (book_id, idx, ctitle, lvl, p1, p2, body, sum(len(t) for t in ptexts)))
            for page, text in chunks_of(ptexts, p1):
                all_chunks.append((cur.lastrowid, page, text))
        c.execute("UPDATE book SET title=?, author=?, lang=?, pages=?, note=?, status='indexing' WHERE id=?",
                  (title, author, detect_lang("".join(texts[:50])), pages, note, book_id))
        c.commit()
        for i in range(0, len(all_chunks), 256):
            part = all_chunks[i:i + 256]
            vecs = embed([t for _, _, t in part])
            c.executemany("INSERT INTO chunk(book_id, chapter_id, kind, page, text, emb) VALUES (?,?,?,?,?,?)",
                          [(book_id, ch, "book", pg, t, v.tobytes()) for (ch, pg, t), v in zip(part, vecs)])
            c.commit()
        c.execute("UPDATE book SET status='ready' WHERE id=?", (book_id,))
        c.commit()
    except Exception as e:  # отдаём ошибку в интерфейс, не роняем приложение
        c.rollback()
        c.execute("UPDATE book SET status='error', note=? WHERE id=?", (f"Не удалось разобрать: {e}", book_id))
        c.commit()


def add_book(filename, data):
    FILES.mkdir(parents=True, exist_ok=True)
    ext = Path(filename).suffix.lower()
    if ext not in (".pdf", ".epub"):
        raise ValueError("Нужен PDF или EPUB")
    c = db()
    bid = c.execute("INSERT INTO book(orig, fmt, added, status) VALUES (?,?,?, 'parsing')",
                    (Path(filename).name, ext[1:], now())).lastrowid
    name = f"{bid}{ext}"
    (FILES / name).write_bytes(data)
    c.execute("UPDATE book SET file=? WHERE id=?", (name, bid))
    c.commit()
    threading.Thread(target=import_book, args=(bid,), daemon=True).start()
    return bid


def delete_book(bid):
    c = db()
    b = c.execute("SELECT file FROM book WHERE id=?", (bid,)).fetchone()
    if b and b["file"]:
        (FILES / b["file"]).unlink(missing_ok=True)
    c.execute("DELETE FROM book WHERE id=?", (bid,))
    c.commit()


# ---------------------------------------------------------------- рекапы в индекс

def index_recap(chapter_id):
    c = db()
    ch = c.execute("SELECT * FROM chapter WHERE id=?", (chapter_id,)).fetchone()
    r = c.execute("SELECT md FROM recap WHERE chapter_id=?", (chapter_id,)).fetchone()
    c.execute("DELETE FROM chunk WHERE chapter_id=? AND kind='recap'", (chapter_id,))
    if r and (r["md"] or "").strip():
        text = f"Рекап главы «{ch['title']}»:\n{r['md']}"
        v = embed([text])[0]
        c.execute("INSERT INTO chunk(book_id, chapter_id, kind, page, text, emb) VALUES (?,?,?,?,?,?)",
                  (ch["book_id"], chapter_id, "recap", ch["p1"], text, v.tobytes()))
    c.commit()


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

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


def search(q, book_id=None, k=8):
    """Гибрид: BM25 (FTS5) + косинус эмбеддингов, слияние по рангам (RRF)."""
    c = db()
    where, args = ("AND c.book_id=?", [book_id]) if book_id else ("", [])
    ranks = {}
    fq = fts_query(q)
    if fq:
        rows = c.execute(f"SELECT c.id FROM chunk_fts f JOIN chunk c ON c.id=f.rowid WHERE chunk_fts MATCH ? {where} "
                         f"ORDER BY bm25(chunk_fts) LIMIT 40", [fq] + args).fetchall()
        for r, row in enumerate(rows):
            ranks[row["id"]] = ranks.get(row["id"], 0) + 1 / (60 + r)
    rows = c.execute(f"SELECT c.id, c.emb FROM chunk c WHERE c.emb IS NOT NULL {where}", args).fetchall()
    if rows:
        qv = embed([q])[0]
        m = np.frombuffer(b"".join(r["emb"] for r in rows), dtype=np.float32).reshape(len(rows), -1)
        sims = m @ qv
        for r, i in enumerate(np.argsort(-sims)[:40]):
            ranks[rows[i]["id"]] = ranks.get(rows[i]["id"], 0) + 1 / (60 + r)
    top = sorted(ranks, key=ranks.get, reverse=True)[:k]
    if not top:
        return []
    res = c.execute(f"SELECT c.*, ch.idx AS ch_idx, ch.title AS ch_title, b.title AS book_title FROM chunk c "
                    f"JOIN chapter ch ON ch.id=c.chapter_id JOIN book b ON b.id=c.book_id "
                    f"WHERE c.id IN ({','.join('?' * len(top))})", top).fetchall()
    order = {cid: i for i, cid in enumerate(top)}
    return sorted(res, key=lambda r: order[r["id"]])


def find_selection(book_id, chapter_idx, text):
    """Окружение выделенного фрагмента в главе — для точного перевода/объяснения."""
    c = db()
    ch = c.execute("SELECT text FROM chapter WHERE book_id=? AND idx=?", (book_id, chapter_idx)).fetchone()
    if not ch:
        return ""
    body = re.sub(r"\[\[p\d+\]\]\n", "", ch["text"])
    probe = re.sub(r"\s+", " ", text.strip())[:60]
    flat = re.sub(r"\s+", " ", body)
    i = flat.find(probe)
    if i < 0:
        return ""
    return flat[max(0, i - 1500): i + len(text) + 1500]