#!/usr/bin/env python3
"""
Replace the fatwa block of QModel.index + metadata.json without re-embedding Quran/Hadith.

Fatwas are the last block of both files (build_index.py writes quran + hadith + fatwa),
so this keeps every earlier vector, drops the old fatwa rows, and appends freshly
embedded fatwas. By default it loads the same fatwas as build_index.py (FATWA_DIRS
minus FATWA_EXCLUDE), so a fatwa-only update matches what a full rebuild produces.

Usage:
    python scripts/update_fatwa_index.py                 # build_index.py's fatwa sources
    python scripts/update_fatwa_index.py --in-place
    python scripts/update_fatwa_index.py --fatwa-dir data/fatwas_scraped   # override sources

Without --in-place, writes QModel.fatwa-update.index / metadata.fatwa-update.json next to
the originals; with it, replaces them and keeps .bak copies.
"""

import argparse
import json
import shutil
import sys
from pathlib import Path

import faiss
import numpy as np
from sentence_transformers import SentenceTransformer

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from build_index import (  # noqa: E402
    DEFAULT_EMBED_MODEL,
    EMBED_BATCH_SIZE,
    FATWA_DIRS,
    FATWA_EXCLUDE,
    INDEX_PATH,
    METADATA_PATH,
    load_fatwa_entries,
    passage_text,
)


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--fatwa-dir", action="append", type=Path,
                        help="Directory of fatwa JSON files (repeatable; default: build_index.FATWA_DIRS)")
    parser.add_argument("--exclude-file", action="append", type=Path,
                        help="JSON file to skip (repeatable; default: build_index.FATWA_EXCLUDE)")
    parser.add_argument("--model", default=DEFAULT_EMBED_MODEL)
    parser.add_argument("--in-place", action="store_true", help="Overwrite QModel.index / metadata.json (keeps .bak)")
    args = parser.parse_args()

    print(f"Loading {INDEX_PATH.name} and {METADATA_PATH.name} …")
    index = faiss.read_index(str(INDEX_PATH))
    metadata = json.loads(METADATA_PATH.read_text(encoding="utf-8"))

    first_fatwa = next((i for i, d in enumerate(metadata) if d.get("type") == "fatwa"), len(metadata))
    if any(d.get("type") == "fatwa" for d in metadata[:first_fatwa]) or \
       any(d.get("type") != "fatwa" for d in metadata[first_fatwa:]):
        sys.exit("Fatwas are not a contiguous tail block of metadata.json — run build_index.py instead.")
    if index.ntotal < first_fatwa:
        sys.exit(f"Index has {index.ntotal} vectors but {first_fatwa} non-fatwa records — run build_index.py instead.")
    print(f"  keeping {first_fatwa:,} Quran/Hadith rows; dropping {len(metadata) - first_fatwa} fatwa records "
          f"and {index.ntotal - first_fatwa} fatwa vectors")

    fatwas = load_fatwa_entries(args.fatwa_dir or FATWA_DIRS, args.exclude_file or FATWA_EXCLUDE)
    if not fatwas:
        sys.exit("No fatwas found.")

    import torch
    device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Embedding {len(fatwas):,} fatwas with {args.model} on {device} …")
    model = SentenceTransformer(args.model, device=device)
    vectors = model.encode(
        [passage_text(f) for f in fatwas], batch_size=EMBED_BATCH_SIZE,
        convert_to_numpy=True, show_progress_bar=True,
    ).astype(np.float32)
    faiss.normalize_L2(vectors)
    if vectors.shape[1] != index.d:
        sys.exit(f"Embedding dim {vectors.shape[1]} != index dim {index.d} — use the model the index was built with.")

    new_index = faiss.IndexFlatIP(index.d)
    new_index.add(index.reconstruct_n(0, first_fatwa))
    new_index.add(vectors)
    new_metadata = metadata[:first_fatwa] + fatwas
    assert new_index.ntotal == len(new_metadata)

    if args.in_place:
        out_index, out_meta = INDEX_PATH, METADATA_PATH
        shutil.copy2(INDEX_PATH, INDEX_PATH.with_suffix(".index.bak"))
        shutil.copy2(METADATA_PATH, METADATA_PATH.with_suffix(".json.bak"))
    else:
        out_index = INDEX_PATH.with_name("QModel.fatwa-update.index")
        out_meta = METADATA_PATH.with_name("metadata.fatwa-update.json")

    faiss.write_index(new_index, str(out_index))
    out_meta.write_text(json.dumps(new_metadata, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"Wrote {out_index.name} ({new_index.ntotal:,} vectors) and {out_meta.name} "
          f"({len(new_metadata):,} records, {len(fatwas):,} fatwas)")


if __name__ == "__main__":
    main()