import os
import re
import json
import math
import time
import sqlite3
import threading
from dataclasses import dataclass
from collections import deque, Counter
from typing import Dict, List, Optional, Tuple, Set

import numpy as np
import faiss
from flask import Flask, request, render_template_string, redirect, url_for, jsonify
from huggingface_hub import snapshot_download

# =========================
# CONFIG
# =========================
HF_REPO = "ArieLLL123/otzaria-embeddings"
DEFAULT_DB_PATH = r"C:\אוצריא\אוצריא\seforim.db"

EDITION_PATHS = {
    "v1": "editions/otzaria_embeddings_v1",
    "v2": "editions/otzaria_embeddings_v2",
    "v3": "editions/otzaria_embeddings_v3",
}

CACHE_DIR = os.path.join(os.path.dirname(__file__), "hf_cache")
RUNTIME_DIR = os.path.join(os.path.dirname(__file__), "runtime")
os.makedirs(CACHE_DIR, exist_ok=True)
os.makedirs(RUNTIME_DIR, exist_ok=True)

# פרמטרים לחיתוך טקסט
DEFAULT_WINDOW_LINES = 6
DEFAULT_STRIDE = 3

# =========================
# TEXT TOOLS & HEBREW NLP
# =========================
NIQQUD_RE   = re.compile(r"[\u0591-\u05C7]")
HTML_TAG_RE = re.compile(r"<[^>]+>")
# שימור גרשיים ומרכאות עבור ראשי תיבות, הסרת שאר הסימנים
NON_WORD_RE = re.compile(r"[^0-9A-Za-z\u0590-\u05FF\"']+")

def clean_text(s: str) -> str:
    """ניקוי טקסט בסיסי - מסיר ניקוד וHTML"""
    if not s: return ""
    s = HTML_TAG_RE.sub(" ", s)
    s = NIQQUD_RE.sub("", s)
    # נורמליזציה של גרשיים
    s = s.replace('״', '"').replace('׳', "'")
    s = NON_WORD_RE.sub(" ", s)
    return " ".join(s.split())

def hebrew_stem(word: str) -> str:
    """
    מסיר קידומות נפוצות בעברית (סטמינג נאיבי).
    הופך 'וכשהאדם' -> 'אדם'
    """
    if len(word) < 4: return word # אל תיגע במילים קצרות
    
    # קידומות: ו, ה, ב, ל, מ, ש, כ, כש, וש, וה, וב...
    # סדר הבדיקה חשוב (מהארוך לקצר)
    prefixes = ['וכש', 'וש', 'וה', 'וב', 'ול', 'ומ', 'כש', 'שב', 'שה', 'מש', 'מה', 'ו', 'ה', 'ב', 'ל', 'מ', 'ש', 'כ']
    
    for p in prefixes:
        if word.startswith(p) and len(word) > len(p) + 2:
            return word[len(p):]
    return word

def get_tokens(text: str) -> Set[str]:
    """מחזיר סט של מילים לאחר סטמינג לחיפוש מדויק יותר"""
    words = clean_text(text).split()
    return {hebrew_stem(w) for w in words}

# =========================
# DATABASE & STREAMING
# =========================
def get_book_titles(db_path: str) -> Dict[int, str]:
    titles = {}
    if not os.path.exists(db_path): return titles
    try:
        con = sqlite3.connect(db_path)
        # תמיכה במבני DB שונים של אוצריא
        try:
            cur = con.execute("SELECT id, heTitle FROM books")
        except:
            cur = con.execute("SELECT id, title FROM books")
            
        for r in cur:
            titles[r[0]] = r[1]
        con.close()
    except Exception as e:
        print(f"Error loading titles: {e}")
    return titles

def iter_rows_ordered(db_path: str, chunk_rows: int = 20000):
    con = sqlite3.connect(db_path)
    con.row_factory = sqlite3.Row
    con.execute("PRAGMA journal_mode=OFF;")
    
    # זיהוי אוטומטי של שם הטבלה (line או lines)
    table_name = "line"
    try:
        con.execute(f"SELECT 1 FROM lines LIMIT 1")
        table_name = "lines"
    except:
        pass

    try:
        # בדיקה שהטבלה קיימת
        con.execute(f"SELECT 1 FROM {table_name} LIMIT 1")
    except:
        return

    q = f"SELECT id, bookId, lineIndex, content FROM {table_name} WHERE content IS NOT NULL AND content != '' ORDER BY bookId, lineIndex"
    cur = con.execute(q)
    while True:
        rows = cur.fetchmany(chunk_rows)
        if not rows: break
        yield rows
    con.close()

def iter_chunks(db_path: str, max_chunks: int, window_lines: int, stride: int):
    rows_iter = iter_rows_ordered(db_path)
    buf = deque()
    cur_book = None
    produced = 0

    for batch in rows_iter:
        for r in batch:
            b_id = r["bookId"]
            if b_id != cur_book:
                cur_book = b_id
                buf.clear()

            txt = str(r["content"])
            if len(txt) < 3: continue
            
            buf.append({"id": r["id"], "idx": r["lineIndex"], "txt": txt})
            
            if len(buf) >= window_lines:
                # בדיקת חפיפה - האם החלון מכיל מספיק טקסט?
                full_text = " ".join(w["txt"] for w in window_lines < len(buf) and buf or list(buf))
                if len(buf) > window_lines: # למקרה של באג
                    window = list(buf)[-window_lines:]
                    full_text = " ".join(w["txt"] for w in window)
                else:
                    full_text = " ".join(w["txt"] for w in buf)

                cln_text = clean_text(full_text)
                
                if len(cln_text) > 30: 
                    yield {
                        "bookId": cur_book,
                        "startLine": buf[0]["idx"],
                        "text": full_text,
                        "clean": cln_text
                    }
                    produced += 1
                    if produced >= max_chunks: return
                
                for _ in range(stride):
                    if buf: buf.popleft()

# =========================
# ENGINE CORE
# =========================
@dataclass
class LoadedModel:
    edition: str
    vocab: Dict[str, int]
    emb_norm: np.ndarray
    idf: np.ndarray

@dataclass
class BuiltIndex:
    faiss_index: faiss.Index
    meta_db_path: str
    count: int

class Engine:
    def __init__(self):
        self.model: Optional[LoadedModel] = None
        self.built: Optional[BuiltIndex] = None
        self.book_map: Dict[int, str] = {}
        self.status = {"state": "idle", "msg": "המערכת מוכנה", "progress": 0}
        self._lock = threading.RLock()

    def _update(self, state, msg, progress):
        with self._lock:
            self.status = {"state": state, "msg": msg, "progress": int(progress)}
        print(f"[{state}] {msg} ({progress}%)")

    def load_resources(self, db_path, edition="v3"):
        if db_path and os.path.exists(db_path):
            self.book_map = get_book_titles(db_path)
        
        try:
            self._update("downloading", f"טוען מודל {edition}...", 5)
            path = EDITION_PATHS.get(edition, EDITION_PATHS["v3"])
            
            local_dir = snapshot_download(repo_id=HF_REPO, repo_type="model", cache_dir=CACHE_DIR,
                                        allow_patterns=[f"{path}/vocab.json", f"{path}/embeddings_last.npy"])
            
            base = os.path.join(local_dir, path)
            with open(f"{base}/vocab.json", "r", encoding="utf-8") as f:
                meta = json.load(f)
            
            emb = np.load(f"{base}/embeddings_last.npy").astype(np.float32)
            
            norms = np.linalg.norm(emb, axis=1, keepdims=True)
            norms[norms == 0] = 1
            emb_norm = emb / norms
            
            vocab = meta["vocab"]
            freqs = np.array(meta.get("freqs", []), dtype=np.float64)
            if len(freqs) == len(vocab):
                idf = np.log((np.sum(freqs) + 1) / (freqs + 1)) + 1
            else:
                idf = np.ones(len(vocab), dtype=np.float32)

            self.model = LoadedModel(edition, vocab, emb_norm, idf.astype(np.float32))
            self._update("idle", "המודל נטען בהצלחה", 100)
            
        except Exception as e:
            self._update("error", f"שגיאה בטעינת מודל: {e}", 0)
            raise e

    def build_index(self, db_path, max_chunks):
        if not self.model: return
        
        stamp = f"{self.model.edition}_N{max_chunks}"
        idx_path = os.path.join(RUNTIME_DIR, f"{stamp}.index")
        meta_db_path = os.path.join(RUNTIME_DIR, f"{stamp}.sqlite")

        if os.path.exists(idx_path) and os.path.exists(meta_db_path):
            self._update("loading", "טוען אינדקס קיים...", 50)
            idx = faiss.read_index(idx_path)
            self.built = BuiltIndex(idx, meta_db_path, idx.ntotal)
            if not self.book_map: self.book_map = get_book_titles(db_path)
            self._update("ready", f"מוכן לחיפוש ({idx.ntotal:,} רשומות)", 100)
            return

        self._update("indexing", "מתחיל בבניית אינדקס (זה יקח זמן)...", 0)
        
        if os.path.exists(meta_db_path): os.remove(meta_db_path)
        con = sqlite3.connect(meta_db_path)
        con.execute("CREATE TABLE chunks (rowid INTEGER PRIMARY KEY, bookId INTEGER, startLine INTEGER, text TEXT)")
        # אינדקס על bookId לסינון מהיר
        con.execute("CREATE INDEX idx_book ON chunks(bookId)") 
        con.execute("PRAGMA synchronous = OFF")
        con.execute("PRAGMA journal_mode = MEMORY")

        d = self.model.emb_norm.shape[1]
        # שימוש ב-IDMap כדי לשמור על סנכרון מלא עם ה-DB
        index = faiss.IndexIDMap(faiss.IndexFlatIP(d))
        
        vectors = []
        ids = []
        db_buffer = []
        batch_size = 5000
        total_processed = 0
        
        start_time = time.time()
        
        # שימוש בפרמטרים אופטימליים לחלון
        for chunk in iter_chunks(db_path, max_chunks, DEFAULT_WINDOW_LINES, DEFAULT_STRIDE):
            vec = self._text_to_vec(chunk["clean"])
            if vec is None: continue
            
            # FAISS IDs חייבים להיות int64
            # נשתמש במונה רץ כ-ID
            current_id = total_processed
            
            vectors.append(vec)
            ids.append(current_id)
            
            # שמירה ל-DB עם אותו ID
            db_buffer.append((current_id, chunk["bookId"], chunk["startLine"], chunk["text"]))
            
            total_processed += 1
            
            if len(vectors) >= batch_size:
                index.add_with_ids(np.vstack(vectors), np.array(ids).astype('int64'))
                con.executemany("INSERT INTO chunks VALUES (?,?,?,?)", db_buffer)
                con.commit()
                vectors, ids, db_buffer = [], [], []
                
                elapsed = time.time() - start_time
                rate = total_processed / (elapsed + 0.1)
                pct = min(95, int((total_processed / max_chunks) * 100))
                self._update("indexing", f"עובדו {total_processed:,} רשומות ({int(rate)} לשנייה)", pct)

        if vectors:
            index.add_with_ids(np.vstack(vectors), np.array(ids).astype('int64'))
            con.executemany("INSERT INTO chunks VALUES (?,?,?,?)", db_buffer)
            con.commit()

        con.close()
        faiss.write_index(index, idx_path)
        
        self.built = BuiltIndex(index, meta_db_path, total_processed)
        self._update("ready", "הבנייה הושלמה בהצלחה!", 100)

    def _text_to_vec(self, text: str):
        words = text.split()
        if not words: return None
        
        # שימוש במילון בצורה וקטורית מהירה
        indices = [self.model.vocab[w] for w in words if w in self.model.vocab]
        if not indices: return None
        
        idfs = self.model.idf[indices]
        vecs = self.model.emb_norm[indices]
        
        weighted = vecs * idfs[:, None]
        avg_vec = np.sum(weighted, axis=0)
        
        norm = np.linalg.norm(avg_vec)
        if norm < 1e-9: return None
        return avg_vec / norm

    def search(self, query: str, book_filter: Optional[int] = None, top_k=20):
        if not self.model or not self.built: return []
        
        q_clean = clean_text(query)
        q_vec = self._text_to_vec(q_clean)
        if q_vec is None: return []
        
        # 1. חיפוש וקטורי רחב יותר (פי 10 מהדרוש כדי לאפשר סינון טוב)
        candidates_k = top_k * 10
        scores, ids = self.built.faiss_index.search(np.array([q_vec]), candidates_k)
        
        found_ids = [int(i) for i in ids[0] if i >= 0]
        if not found_ids: return []
        
        # 2. שליפת טקסטים + סינון ספרים ברמת ה-SQL
        con = sqlite3.connect(self.built.meta_db_path)
        con.row_factory = sqlite3.Row
        
        id_list_str = ",".join(map(str, found_ids))
        sql = f"SELECT * FROM chunks WHERE rowid IN ({id_list_str})"
        
        if book_filter:
            sql += f" AND bookId = {int(book_filter)}"
            
        rows = con.execute(sql).fetchall()
        con.close()
        
        if not rows: return []

        chunk_map = {r["rowid"]: dict(r) for r in rows}
        
        # 3. דירוג מחדש (Re-ranking) משופר
        results = []
        
        # הכנת טוקנים לשאילתה (כולל סטמינג)
        q_tokens = get_tokens(q_clean)
        q_original_words = q_clean.split()
        
        # מיפוי ציון וקטורי לפי ID
        vec_scores = {fid: float(scr) for fid, scr in zip(ids[0], scores[0])}

        for row in rows:
            rid = row["rowid"]
            chunk_txt = row["text"]
            chunk_clean = clean_text(chunk_txt)
            chunk_tokens = get_tokens(chunk_clean)
            chunk_words = chunk_clean.split()
            
            # --- חישוב ציונים ---
            
            # 1. ציון וקטורי (בסיס) - נרמול גס ל 0-1
            base_score = vec_scores.get(rid, 0)
            
            # 2. חפיפת מילים (Jaccard) עם סטמינג
            # כמה מתוך מילות החיפוש הייחודיות נמצאות בטקסט?
            intersection = len(q_tokens & chunk_tokens)
            overlap_score = (intersection / len(q_tokens)) if q_tokens else 0
            
            # 3. רצף מדויק (Exact Phrase)
            phrase_score = 0
            if q_clean in chunk_clean:
                phrase_score = 1.0
            
            # 4. צפיפות (Proximity) - האם המילים קרובות?
            proximity_score = 0
            if intersection > 1:
                # מוצאים את האינדקסים של מילות המפתח בטקסט
                found_indices = []
                for qw in q_tokens:
                    # חיפוש נאיבי של המילה בטקסט (אחרי סטמינג)
                    for i, cw in enumerate(chunk_words):
                        if hebrew_stem(cw) == qw:
                            found_indices.append(i)
                            break # מספיק מופע אחד לחישוב צפיפות בסיסי
                
                if found_indices:
                    span = max(found_indices) - min(found_indices)
                    # ככל שהטווח קטן יותר ביחס לאורך הטקסט, הציון גבוה
                    # (הוספת 1 למניעת חלוקה באפס)
                    density = len(found_indices) / (span + 1)
                    proximity_score = min(density, 1.0)

            # שקלול סופי (Weighted Sum)
            # נותנים משקל גבוה לחפיפת מילים כדי לסנן "רעש" וקטורי
            final_score = (
                (base_score * 0.4) + 
                (overlap_score * 0.35) + 
                (phrase_score * 0.15) + 
                (proximity_score * 0.1)
            )
            
            book_title = self.book_map.get(row["bookId"], f"ספר {row['bookId']}")
            
            results.append({
                "score": final_score,
                "text": chunk_txt,
                "source": f"{book_title}, שורה {row['startLine']}",
                "book_id": row["bookId"],
                "book_title": book_title # לצורך סינון בUI
            })
            
        # מיון יורד
        results.sort(key=lambda x: x["score"], reverse=True)
        return results[:top_k]

ENGINE = Engine()

# =========================
# FLASK WEB APP
# =========================
app = Flask(__name__)
app.secret_key = "otzaria_ai_secret_v2"

HTML_TEMPLATE = """
<!DOCTYPE html>
<html lang="he" dir="rtl">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>אוצריא AI - מנוע חכם</title>
    <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0/dist/css/bootstrap.min.css" rel="stylesheet">
    <link href="https://fonts.googleapis.com/css2?family=Heebo:wght@300;400;700&family=Frank+Ruhl+Libre:wght@400;700&display=swap" rel="stylesheet">
    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/bootstrap-icons@1.10.5/font/bootstrap-icons.css">
    <style>
        :root {
            --primary: #2c3e50;
            --accent: #d35400; /* כתום כהה יותר, תורני יותר */
            --bg-light: #fdfbf7; /* גוון נייר עדין */
            --text-serif: 'Frank Ruhl Libre', serif;
            --text-sans: 'Heebo', sans-serif;
        }
        body { background-color: var(--bg-light); font-family: var(--text-sans); color: #333; }
        
        .sidebar {
            background: #fff; height: 100vh; position: fixed; right: 0; top: 0; width: 300px;
            padding: 2rem 1.5rem; border-left: 1px solid #e0e0e0; overflow-y: auto;
        }
        .main-content { margin-right: 300px; padding: 2rem; max-width: 900px; }
        
        @media (max-width: 900px) {
            .sidebar { position: relative; width: 100%; height: auto; border-left: none; }
            .main-content { margin-right: 0; }
        }

        .brand { font-size: 1.8rem; font-weight: 700; color: var(--primary); margin-bottom: 2rem; display: block; text-decoration: none; }
        .brand span { color: var(--accent); }
        
        .search-box {
            background: #fff; padding: 1.5rem; border-radius: 12px; box-shadow: 0 4px 20px rgba(0,0,0,0.06);
            margin-bottom: 2rem; border: 1px solid #eee;
        }
        .search-input { border: 1px solid #ddd; padding: 0.8rem; font-size: 1.1rem; border-radius: 8px; }
        .search-input:focus { border-color: var(--accent); box-shadow: 0 0 0 3px rgba(211, 84, 0, 0.1); }

        .result-card {
            background: #fff; border-radius: 8px; padding: 1.2rem; margin-bottom: 1.2rem;
            border-right: 4px solid var(--accent); box-shadow: 0 2px 8px rgba(0,0,0,0.03);
            transition: transform 0.2s;
        }
        .result-card:hover { transform: translateY(-2px); box-shadow: 0 5px 15px rgba(0,0,0,0.08); }
        
        .res-header { display: flex; justify-content: space-between; align-items: center; margin-bottom: 0.8rem; }
        .res-source { font-size: 0.95rem; font-weight: 700; color: var(--primary); }
        .res-text { font-family: var(--text-serif); font-size: 1.2rem; line-height: 1.6; color: #222; }
        
        mark { background-color: #ffeaa7; padding: 0 2px; border-radius: 3px; }
        
        /* Stats Bar */
        .stats-pill { font-size: 0.8rem; background: #eee; padding: 4px 10px; border-radius: 20px; color: #666; }
    </style>
</head>
<body>

<div class="sidebar">
    <a href="/" class="brand">אוצריא <span>AI</span></a>
    
    <div class="mb-4">
        <label class="small fw-bold text-muted mb-2">סטטוס אינדקס</label>
        <div class="progress" style="height: 6px;">
            <div id="status-bar" class="progress-bar bg-success" style="width: 0%"></div>
        </div>
        <div class="d-flex justify-content-between mt-1">
            <span class="small text-muted" id="status-text">ממתין...</span>
            <span class="small fw-bold" id="idx-count">{{ idx_count }}</span>
        </div>
    </div>

    <hr>
    
    <form action="/setup" method="post" class="mt-4">
        <h6 class="fw-bold mb-3"><i class="bi bi-gear"></i> הגדרות בנייה</h6>
        <div class="mb-3">
            <label class="form-label small">נתיב בסיס נתונים</label>
            <input type="text" name="db_path" class="form-control form-control-sm" value="{{ db_path }}">
        </div>
        <div class="mb-3">
            <label class="form-label small">מספר רשומות (Chunks)</label>
            <input type="number" name="max_chunks" class="form-control form-control-sm" value="100000">
        </div>
        <button type="submit" class="btn btn-dark btn-sm w-100">בנה מחדש</button>
    </form>
</div>

<div class="main-content">
    <form action="/" method="get" class="search-box">
        <div class="row g-2">
            <div class="col-md-8">
                <input type="text" name="q" class="form-control search-input" placeholder="חיפוש חופשי (לדוגמה: אהבת ישראל ברמבם)" value="{{ query or '' }}" autofocus>
            </div>
            <div class="col-md-4">
                <select name="book_id" class="form-select" style="height: 100%; padding: 0.8rem;">
                    <option value="">כל הספרים</option>
                    {% for bid, title in books.items() %}
                        <option value="{{ bid }}" {% if selected_book|int == bid %}selected{% endif %}>{{ title }}</option>
                    {% endfor %}
                </select>
            </div>
            <div class="col-12 mt-2 text-end">
                <button type="submit" class="btn btn-warning text-white fw-bold px-4">חפש</button>
            </div>
        </div>
    </form>

    {% if query %}
        <div class="d-flex justify-content-between align-items-center mb-4">
            <h5 class="mb-0">תוצאות עבור: <strong>{{ query }}</strong></h5>
            <span class="stats-pill">{{ results|length }} תוצאות</span>
        </div>

        {% if not results %}
             <div class="text-center py-5 text-muted">
                <i class="bi bi-search display-4 opacity-25"></i>
                <p class="mt-3">לא נמצאו תוצאות מתאימות.</p>
                {% if selected_book %}
                <p class="small">נסה להסיר את הסינון לפי ספר.</p>
                {% endif %}
             </div>
        {% endif %}

        {% for r in results %}
        <div class="result-card">
            <div class="res-header">
                <span class="res-source"><i class="bi bi-book-fill me-1 opacity-50"></i> {{ r.source }}</span>
                <!-- <span class="badge bg-light text-muted border">{{ "%.2f"|format(r.score) }}</span> -->
            </div>
            <div class="res-text">
                {{ r.text | highlight(query) | safe }}
            </div>
        </div>
        {% endfor %}
        
    {% else %}
        <div class="text-center py-5 opacity-75">
            <h4>ברוכים הבאים למנוע החיפוש הסמנטי</h4>
            <p>המערכת מבינה הקשרים ולא רק מילים מדויקות.</p>
        </div>
    {% endif %}
</div>

<script>
    function updateStatus() {
        fetch('/status')
            .then(r => r.json())
            .then(data => {
                document.getElementById('status-text').innerText = data.msg;
                document.getElementById('status-bar').style.width = data.progress + '%';
                if(data.count) document.getElementById('idx-count').innerText = data.count.toLocaleString();
                
                let interval = (data.state === 'indexing' || data.state === 'downloading') ? 1000 : 5000;
                setTimeout(updateStatus, interval);
            })
            .catch(e => setTimeout(updateStatus, 5000));
    }
    document.addEventListener('DOMContentLoaded', updateStatus);
</script>

</body>
</html>
"""

# =========================
# HELPER FILTERS
# =========================
def highlight_text(text, query):
    if not query: return text
    
    # הדגשה חכמה יותר שתופסת גם עם קידומות
    # אנו משתמשים בסטמינג כדי לזהות שורשים
    q_words = [hebrew_stem(w) for w in clean_text(query).split() if len(w) > 1]
    
    if not q_words: return text
    
    # בניית ביטוי רגולרי גמיש לעברית
    # תופס: (רווח או התחלה) + (אופציונלי: ו/מ/ש/ה/ל/ב) + המילה + (סוף מילה)
    patterns = []
    for w in q_words:
        # Regex שמתאים למילה, גם אם יש לה קידומת אותיות שימוש
        pat = r'(?:^|[\s\"\'\-])([ו|מ|ש|ה|ל|ב|כ]?' + re.escape(w) + r')(?=[\s\"\'\.\,\-]|$)'
        patterns.append(pat)
        
    combined_pattern = "|".join(patterns)
    
    def replacer(match):
        # מחזירים את הטקסט המקורי עטוף ב-mark, שומרים על התווים מסביב
        full_match = match.group(0)
        # מוצאים את המילה נטו בתוך ההתאמה כדי לא לצבוע רווחים
        word_match = re.search(r'[א-ת]+', full_match)
        if word_match:
            wm = word_match.group(0)
            return full_match.replace(wm, f'<mark>{wm}</mark>')
        return full_match

    try:
        # שימוש ב-Ignore Case לא רלוונטי בעברית אבל לא מזיק
        text = re.sub(combined_pattern, replacer, text)
    except:
        pass # במקרה של שגיאת regex, מחזיר טקסט רגיל
        
    return text

@app.template_filter('highlight')
def highlight_filter(text, query):
    return highlight_text(text, query)

# =========================
# ROUTES
# =========================
@app.route("/")
def index():
    q = request.args.get("q", "").strip()
    book_id_str = request.args.get("book_id", "")
    book_id = int(book_id_str) if book_id_str.isdigit() else None
    
    results = []
    
    # בניית רשימת ספרים ל-Dropdown (ממוין)
    # מגבילים ל-500 ספרים נפוצים או הכל אם יש מעט, כדי לא להעמיס על ה-HTML
    all_books = ENGINE.book_map.copy()
    sorted_books = dict(sorted(all_books.items(), key=lambda item: item[1])[:800])

    if q:
        if not ENGINE.built:
             # אם אין אינדקס, מתחיל לבנות ברקע ומחזיר דף ריק בינתיים
            threading.Thread(target=lambda: ENGINE.build_index(DEFAULT_DB_PATH, 10000), daemon=True).start()
        else:
            results = ENGINE.search(q, book_filter=book_id, top_k=20)

    idx_c = ENGINE.built.count if ENGINE.built else 0
    
    return render_template_string(HTML_TEMPLATE, 
                                  query=q, 
                                  results=results, 
                                  db_path=DEFAULT_DB_PATH, 
                                  idx_count=idx_c,
                                  books=sorted_books,
                                  selected_book=book_id)

@app.route("/setup", methods=["POST"])
def setup():
    db = request.form.get("db_path", DEFAULT_DB_PATH)
    mc = int(request.form.get("max_chunks", 100000))
    
    def task():
        ENGINE.load_resources(db)
        ENGINE.build_index(db, mc)

    threading.Thread(target=task, daemon=True).start()
    return redirect("/")

@app.route("/status")
def status_api():
    s = ENGINE.status.copy()
    if ENGINE.built:
        s["count"] = ENGINE.built.count
    return jsonify(s)

# =========================
# MAIN
# =========================
if __name__ == "__main__":
    # נסיון טעינה שקט בהתחלה
    threading.Thread(target=lambda: ENGINE.load_resources(DEFAULT_DB_PATH), daemon=True).start()

    print("Starting Enhanced Server at http://127.0.0.1:8000")
    app.run(host="127.0.0.1", port=8000, debug=True)