import argparse
from contextlib import closing
from dataclasses import dataclass
from hashlib import sha256
import json
from pathlib import Path
import re
import sqlite3
import textwrap

from pypdf.errors import PdfReadError
from pdf_pages import extract_pages
from semantic_notes import cosine, embed, model_digest

DEFAULT_DB = Path("handbook.sqlite3")


@dataclass(frozen=True)
class Passage:
    id: int
    source: str
    page: int
    text: str


def load_pdf_files(path):
    data = path.read_bytes()
    ocr_files = {}
    pattern = re.compile(re.escape(path.stem) + r"\.page-([1-9]\d*)\.txt")
    for sibling in sorted(path.parent.iterdir()):
        match = pattern.fullmatch(sibling.name)
        if match and not sibling.is_symlink() and sibling.is_file():
            ocr_files[int(match[1])] = sibling.read_bytes()
    checksum = sha256()
    checksum.update(sha256(data).digest())
    for page, text in sorted(ocr_files.items()):
        checksum.update(f"{page}:".encode("ascii") + sha256(text).digest())
    return data, ocr_files, checksum.hexdigest()


def replace_pdf(db, source, data, ocr_files, checksum):
    pages = extract_pages(data)
    if ocr_files.keys() - {number for number, _ in pages}:
        raise ValueError("An OCR text filename refers to a page outside the PDF.")
    pages = [
        (page, text or ocr_files.get(page, b"").decode("utf-8").strip())
        for page, text in pages
    ]
    chunks = [
        (page, chunk)
        for page, text in pages
        for chunk in textwrap.wrap(" ".join(text.split()), width=480)
    ]
    db.execute("INSERT INTO files VALUES (?, ?)", (source, checksum))
    for start in range(0, len(chunks), 16):
        batch = chunks[start:start + 16]
        vectors = embed([text for _, text in batch])
        for (page, text), vector in zip(batch, vectors):
            row = db.execute(
                "INSERT INTO passages(source, page, text, vector) VALUES (?, ?, ?, ?)",
                (source, page, text, json.dumps(vector)),
            )
            db.execute("INSERT INTO words(rowid, text) VALUES (?, ?)", (row.lastrowid, text))
    return [page for page, text in pages if not text]


def remove_pdf(db, source):
    db.execute("DELETE FROM words WHERE rowid IN (SELECT id FROM passages WHERE source = ?)", (source,))
    db.execute("DELETE FROM passages WHERE source = ?", (source,))
    db.execute("DELETE FROM files WHERE source = ?", (source,))


def update(folder, database):
    if not folder.is_dir():
        raise ValueError(f"PDF folder not found: {folder}")
    files = {
        path.relative_to(folder).as_posix(): path for path in sorted(folder.rglob("*"))
        if path.is_file() and not path.is_symlink() and path.suffix.lower() == ".pdf"
    }
    digest = model_digest()
    counts = {"indexed": 0, "reused": 0, "removed": 0}
    empty_pages = []
    with closing(sqlite3.connect(database)) as db, db:
        db.execute("BEGIN IMMEDIATE")
        db.execute("CREATE TABLE IF NOT EXISTS settings (id INTEGER PRIMARY KEY, digest TEXT NOT NULL)")
        db.execute("CREATE TABLE IF NOT EXISTS files (source TEXT PRIMARY KEY, checksum TEXT NOT NULL)")
        db.execute(
            "CREATE TABLE IF NOT EXISTS passages (id INTEGER PRIMARY KEY, source TEXT NOT NULL, "
            "page INTEGER NOT NULL, text TEXT NOT NULL, vector TEXT NOT NULL)"
        )
        db.execute("CREATE VIRTUAL TABLE IF NOT EXISTS words USING fts5(text)")
        stored = dict(db.execute("SELECT source, checksum FROM files"))
        revision = db.execute("SELECT digest FROM settings WHERE id = 1").fetchone()
        for source in sorted(stored.keys() - files.keys()):
            remove_pdf(db, source)
            counts["removed"] += 1
        for source, path in files.items():
            data, ocr_files, checksum = load_pdf_files(path)
            if stored.get(source) == checksum and revision == (digest,):
                counts["reused"] += 1
                continue
            remove_pdf(db, source)
            try:
                skipped = replace_pdf(db, source, data, ocr_files, checksum)
            except (PdfReadError, ValueError) as error:
                raise ValueError(f"{source}: {error}") from error
            empty_pages.extend((source, page) for page in skipped)
            counts["indexed"] += 1
        if model_digest() != digest:
            raise ValueError("The embedding model changed during the update. Run it again.")
        db.execute("INSERT OR REPLACE INTO settings VALUES (1, ?)", (digest,))

    for source, page in empty_pages:
        print(f"Skipped {source}, page {page}: no embedded text or OCR text file.")
    return counts


def fuse_rankings(*rankings):
    scores = {}
    for ranking in rankings:
        for rank, row_id in enumerate(ranking, start=1):
            scores[row_id] = scores.get(row_id, 0) + 1 / (60 + rank)
    return sorted(scores, key=lambda row_id: (-scores[row_id], row_id))


def search(database, query, mode="hybrid", limit=3):
    query = query.strip()
    if not query:
        raise ValueError("Enter a search question.")
    if not database.is_file():
        raise ValueError("Build the handbook index first.")
    terms = " OR ".join(f'"{word}"' for word in re.findall(r"[^\W_]+", query))
    with closing(sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True)) as db:
        db.execute("BEGIN")
        revision = db.execute("SELECT digest FROM settings WHERE id = 1").fetchone()
        keyword = [
            row[0] for row in db.execute(
                "SELECT rowid FROM words WHERE words MATCH ? ORDER BY rank, rowid LIMIT 8", (terms,)
            )
        ] if terms and mode != "vector" else []
        if mode == "keyword":
            hits = [db.execute(
                "SELECT id, source, page, text FROM passages WHERE id = ?", (row_id,)
            ).fetchone() for row_id in keyword[:limit]]
            return [Passage(*row) for row in hits]
        rows = db.execute("SELECT id, source, page, text, vector FROM passages ORDER BY id").fetchall()

    vector = []
    if rows:
        digest = model_digest()
        if revision != (digest,):
            raise ValueError("The embedding model changed. Run update before searching.")
        query_vector, = embed([query])
        if model_digest() != digest:
            raise ValueError("The embedding model changed during search. Run update and retry.")
        ranked = sorted(rows, key=lambda row: cosine(json.loads(row[4]), query_vector), reverse=True)
        vector = [row[0] for row in ranked[:8]]

    ids = fuse_rankings(keyword, vector) if mode == "hybrid" else vector
    candidates = {row[0]: Passage(*row[:4]) for row in rows}
    return [candidates[row_id] for row_id in ids[:limit]]


def main():
    parser = argparse.ArgumentParser(description="Update and search a small collection of PDF files.")
    parser.add_argument("--db", type=Path, default=DEFAULT_DB, help="Dedicated handbook SQLite file")
    commands = parser.add_subparsers(dest="command", required=True)
    rebuild = commands.add_parser("update")
    rebuild.add_argument("folder", type=Path)
    find = commands.add_parser("find")
    find.add_argument("query")
    find.add_argument("--mode", choices=("keyword", "vector", "hybrid"), default="hybrid")
    args = parser.parse_args()
    try:
        if args.command == "update":
            counts = update(args.folder, args.db)
            print(f"Indexed {counts['indexed']}, reused {counts['reused']}, removed {counts['removed']} PDFs.")
        else:
            matches = search(args.db, args.query, args.mode)
            for passage in matches:
                print(f"{passage.source}, page {passage.page}\n{passage.text}\n")
            if not matches:
                print("No matching passages.")
    except (OSError, sqlite3.Error, ValueError) as error:
        raise SystemExit(str(error)) from None


if __name__ == "__main__":
    main()