"""
Builds rag_index/ from the crawled corpus in the sibling v:\\crawler project
(79 digital-marketing books + the web crawl across SEO / Social Media
Marketing / Online Advertising / Digital Marketing General), instead of this
project's own crawled_*/books-2 folders.

Adds one metadata field beyond the original build_index.py: "category" (one
of "seo", "social_media", "online_ads", "general") — a coarse topical bucket,
not identifying information — so seo.py/social.py/ads.py can each retrieve
from their own domain instead of the whole mixed corpus. Title/author/URL/
domain are still never stored, preserving the "Source N" anonymization
guarantee.

Usage:
    python build_index_from_crawler.py --limit 50   # quick correctness check
    python build_index_from_crawler.py --reset      # full rebuild
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import sqlite3
from datetime import datetime, timezone
from pathlib import Path

import tiktoken
from tqdm import tqdm

CRAWLER_DB = Path(r"V:\crawler\data\default\meta.db")
ENCODING = tiktoken.get_encoding("cl100k_base")
CHUNK_TOKENS = 800
OVERLAP_TOKENS = 100
EMBED_MODEL = "BAAI/bge-small-en-v1.5"
EMBED_BATCH = 64

CATEGORY_MAP = {
    "SEO": "seo",
    "SOCIAL MEDIA MARKETING": "social_media",
    "ONLINE ADVERTISING": "online_ads",
    "DIGITAL MARKETING (General)": "general",
}

BOOK_PREFIX_MAP = [
    (re.compile(r"^digital marketing seo", re.I), "seo"),
    (re.compile(r"^digital marketing social", re.I), "social_media"),
    (re.compile(r"^digital marketing ads", re.I), "online_ads"),
]

BOOK_KEYWORD_MAP = [
    (re.compile(r"\bseo\b|search engine optimi[sz]ation|link[- ]building", re.I), "seo"),
    (re.compile(r"social media|facebook|instagram|tiktok|youtube|influencer", re.I), "social_media"),
    (re.compile(r"\bads?\b|advertising|adwords|google ads|\bppc\b", re.I), "online_ads"),
]


def classify_book(filename: str) -> str:
    for pattern, category in BOOK_PREFIX_MAP:
        if pattern.search(filename):
            return category
    for pattern, category in BOOK_KEYWORD_MAP:
        if pattern.search(filename):
            return category
    return "general"


def opaque_id(text: str) -> str:
    return hashlib.sha1(text.encode("utf-8")).hexdigest()[:16]


def chunk_text(text: str) -> list[str]:
    tokens = ENCODING.encode(text)
    if not tokens:
        return []
    chunks = []
    step = CHUNK_TOKENS - OVERLAP_TOKENS
    for start in range(0, len(tokens), step):
        window = tokens[start : start + CHUNK_TOKENS]
        if len(window) < 50:
            break
        chunks.append(ENCODING.decode(window))
        if start + CHUNK_TOKENS >= len(tokens):
            break
    return chunks


def _document_text(conn: sqlite3.Connection, file_id: int) -> str:
    pages = conn.execute(
        "SELECT text FROM pages WHERE file_id=? ORDER BY page_no", (file_id,)
    ).fetchall()
    return "\n\n".join(p["text"] for p in pages if p["text"])


def iter_documents(crawler_db: Path, limit: int | None = None):
    """Yields (text, category) for every crawled file (books + web articles)."""
    conn = sqlite3.connect(str(crawler_db))
    conn.row_factory = sqlite3.Row

    web_rows = conn.execute(
        """
        SELECT f.id as file_id, s.category as category
        FROM candidates c
        JOIN sources s ON s.id = c.source_id
        JOIN files f ON f.path = c.url
        WHERE c.selected = 1
        """
    ).fetchall()

    book_rows = conn.execute(
        "SELECT id as file_id, path FROM files WHERE path NOT LIKE 'http%'"
    ).fetchall()

    count = 0
    for row in web_rows:
        if limit and count >= limit:
            break
        text = _document_text(conn, row["file_id"])
        if text.strip():
            yield text, CATEGORY_MAP.get(row["category"], "general")
            count += 1

    for row in book_rows:
        if limit and count >= limit:
            break
        text = _document_text(conn, row["file_id"])
        if text.strip():
            yield text, classify_book(Path(row["path"]).name)
            count += 1

    conn.close()


def build_chunks(crawler_db: Path, limit: int | None = None):
    all_chunks: list[tuple[str, str]] = []
    seen_ids: set[str] = set()

    for doc_text, category in tqdm(list(iter_documents(crawler_db, limit=limit)), desc="chunking"):
        for chunk in chunk_text(doc_text):
            cid = opaque_id(chunk)
            if cid in seen_ids:
                continue
            seen_ids.add(cid)
            all_chunks.append((chunk, category))
    return all_chunks


def embed_and_write(chunks_with_category: list[tuple[str, str]], out_dir: Path, reset: bool = False):
    import chromadb
    import torch
    from sentence_transformers import SentenceTransformer

    if torch.cuda.is_available():
        device = "cuda"
        batch_size = 128
    else:
        device = "cpu"
        batch_size = EMBED_BATCH
        torch.set_num_threads(os.cpu_count() or 4)  # CPU default under-uses available cores
    print(f"[embed] using device={device}, batch_size={batch_size}")

    out_dir.mkdir(parents=True, exist_ok=True)
    client = chromadb.PersistentClient(path=str(out_dir))
    if reset:
        try:
            client.delete_collection("dm_rag")
        except Exception:
            pass
    collection = client.get_or_create_collection(name="dm_rag", metadata={"hnsw:space": "cosine"})

    model = SentenceTransformer(EMBED_MODEL, device=device)

    texts = [c for c, _ in chunks_with_category]
    categories = [cat for _, cat in chunks_with_category]
    ids = [opaque_id(t) for t in texts]

    for i in tqdm(range(0, len(texts), batch_size), desc="embedding+writing"):
        batch_texts = texts[i : i + batch_size]
        batch_ids = ids[i : i + batch_size]
        batch_cats = categories[i : i + batch_size]
        batch_emb = model.encode(
            batch_texts, batch_size=batch_size, normalize_embeddings=True, show_progress_bar=False
        ).tolist()
        collection.upsert(
            ids=batch_ids,
            documents=batch_texts,
            embeddings=batch_emb,
            metadatas=[{"category": c} for c in batch_cats],
        )

    counts_by_cat: dict[str, int] = {}
    for c in categories:
        counts_by_cat[c] = counts_by_cat.get(c, 0) + 1

    meta = {
        "model": EMBED_MODEL,
        "chunk_tokens": CHUNK_TOKENS,
        "overlap_tokens": OVERLAP_TOKENS,
        "count": len(texts),
        "counts_by_category": counts_by_cat,
        "built_at": datetime.now(timezone.utc).isoformat(),
        "source": "v:/crawler (79 books + web crawl across 4 categories)",
    }
    (out_dir / "index_meta.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
    print(f"Wrote {len(texts)} chunks to {out_dir}")
    print(f"By category: {counts_by_cat}")


def main():
    parser = argparse.ArgumentParser(description="Build rag_index/ from v:/crawler's crawled corpus")
    parser.add_argument("--crawler-db", default=str(CRAWLER_DB))
    parser.add_argument("--out", default=str(Path(__file__).resolve().parent.parent / "rag_index"))
    parser.add_argument("--limit", type=int, default=None, help="Limit documents processed (for testing)")
    parser.add_argument("--reset", action="store_true", help="Wipe the existing collection first")
    args = parser.parse_args()

    chunks = build_chunks(Path(args.crawler_db), limit=args.limit)
    if not chunks:
        print("No chunks produced — check --crawler-db path.")
        return
    embed_and_write(chunks, Path(args.out), reset=args.reset)


if __name__ == "__main__":
    main()