#!/usr/bin/env python3
"""Read-only EPUB inspection. Reports structure, segment counts, estimates.
Does NOT modify the EPUB and does NOT call the model."""
import sys, zipfile, shutil, re, warnings
from pathlib import Path
from bs4 import BeautifulSoup, NavigableString, XMLParsedAsHTMLWarning
warnings.filterwarnings("ignore", category=XMLParsedAsHTMLWarning)

ROOT = Path("/opt/epub-translate")
INPUT_DIR = ROOT / "input"
WORK = ROOT / "work"

# Tags whose text we will NOT translate
SKIP_TAGS = {"script", "style", "code", "pre", "head", "title", "meta", "link"}
# Block-ish tags we treat as "paragraph" segment boundaries
BLOCK_TAGS = {"p", "h1", "h2", "h3", "h4", "h5", "h6",
              "li", "blockquote", "figcaption", "td", "th", "dd", "dt"}

def find_epub():
    epubs = list(INPUT_DIR.glob("*.epub"))
    if not epubs:
        sys.exit(f"No .epub found in {INPUT_DIR}")
    return epubs[0]

def unzip(epub):
    if WORK.exists():
        shutil.rmtree(WORK)
    WORK.mkdir(parents=True)
    with zipfile.ZipFile(epub) as z:
        z.extractall(WORK)
    return sorted(WORK.rglob("*"))

def is_translatable(node):
    txt = (node or "").strip()
    if not txt:
        return False
    # skip purely numeric / punctuation-only nodes
    if not re.search(r"[A-Za-z\u00C0-\u024F]", txt):
        return False
    return True

def parse(path):
    raw = path.read_text(encoding="utf-8", errors="replace")
    return BeautifulSoup(raw, "xml" if path.suffix.lower()==".xhtml" else "lxml")

def has_letters(s):
    return bool(s) and bool(re.search(r"[A-Za-z\u00C0-\u024F]", s))

def in_skip(node):
    for parent in node.parents:
        if parent.name in SKIP_TAGS:
            return True
    return False

def text_nodes(soup):
    """Yield translatable NavigableString nodes (the real unit of work)."""
    for node in soup.find_all(string=True):
        if isinstance(node, NavigableString) and has_letters(str(node)) and not in_skip(node):
            yield node





def main():
    epub = find_epub()
    size_mb = epub.stat().st_size / 1_048_576
    print(f"EPUB: {epub.name}  ({size_mb:.2f} MB)\n")

    members = unzip(epub)
    xhtml = sorted([p for p in members
                    if p.suffix.lower() in (".xhtml", ".html", ".htm")])
    print(f"Total files in container: {len([m for m in members if m.is_file()])}")
    print(f"XHTML/HTML content files: {len(xhtml)}\n")

    total_segments = 0
    total_words = 0
    print(f"{'file':45} {'segments':>9} {'words':>8}")
    print("-" * 64)



    per_file = []
    total_nodes = total_words = total_chars = 0
    print(f"{'file':40} {'textnodes':>10} {'words':>8} {'chars':>8}")
    print("-"*70)
    for x in xhtml:
        soup = parse(x)
        nodes = list(text_nodes(soup))
        words = sum(len(str(n).split()) for n in nodes)
        chars = sum(len(str(n).strip()) for n in nodes)
        rel = str(x.relative_to(WORK))
        per_file.append((rel, len(nodes), words, chars))
        total_nodes += len(nodes); total_words += words; total_chars += chars
        print(f"{rel[:40]:40} {len(nodes):>10} {words:>8} {chars:>8}")
    print("-"*70)
    print(f"{'TOTAL':40} {total_nodes:>10} {total_words:>8} {total_chars:>8}\n")



    # --- Sanity sample: show first real paragraphs of the biggest chapter ---
    biggest = max(per_file, key=lambda t: t[2]) if per_file else None
    if biggest:
        bx = WORK / biggest[0]
        soup = parse(bx)
        for bad in soup.find_all(SKIP_TAGS):
            bad.extract()
        print(f"\n--- SAMPLE from {biggest[0]} (first 5 text blocks) ---")
        shown = 0
        for el in soup.find_all(BLOCK_TAGS):
            t = el.get_text(" ", strip=True)
            if is_translatable(t):
                print(f"[{el.name}] {t[:200]}")
                shown += 1
                if shown >= 5:
                    break



    # --- Estimates (CPU-only throughput is the bottleneck) ---
    # in-tokens ~1.3/word EN; Russian out ~1.6x the English tokens
    est_out_tok = total_words * 1.3 * 1.6
    print(f"~{total_words:,} EN words  ->  ~{est_out_tok:,.0f} RU output tokens (rough)\n")
    for rate in (4, 8, 15):  # CPU tok/s scenarios for gemma4:26b MoE
        print(f"  @ {rate:>2} tok/s  -> ~{est_out_tok/rate/3600:5.1f} h output gen "
              f"(plus prompt-eval time, so real wall-time is higher)")



    print("\nNOTE: measure the REAL rate with the sample step before the full run.")
    print(f"Unzipped tree is in: {WORK}")
    print("This script made NO changes to the EPUB and called NO model.")

if __name__ == "__main__":
    main()