import os
from collections import Counter
from datetime import timedelta

from common import connect, URI, read_chunks, to_record

rows, model, dim = read_chunks()
records = [to_record(r) for r in rows]
db = connect()
NAME = "docs_upsert"  # a scratch copy so the main "docs" table stays untouched
table = db.create_table(NAME, data=records, mode="overwrite")


def stats():
    """File count and bytes of the table dir (local path only)."""
    n = size = 0
    for root, _, files in os.walk(f"{URI}/{NAME}.lance"):
        for f in files:
            n += 1
            size += os.path.getsize(os.path.join(root, f))
    return n, size


def upsert(recs, delete_filter=None):
    m = table.merge_insert("id").when_matched_update_all().when_not_matched_insert_all()
    if delete_filter is not None:
        m = m.when_not_matched_by_source_delete(delete_filter)
    return m.execute(recs)


print("start:", table.count_rows(), "rows,", len(table.list_versions()), "version(s)")

print("\n1. upsert the same 304 chunks again")
res = upsert(records)
print("  result:", res)
print("  count:", table.count_rows())

print("\n2. change one chunk's text")
changed = dict(records[10], text=records[10]["text"] + "\n\nEdited after re-crawl.")
res = upsert([changed])
print("  result:", res)
got = table.search().where(f"id = '{changed['id']}'").select(["text"]).limit(1).to_list()[0]["text"]
print("  edit visible:", got.endswith("Edited after re-crawl."), "| count:", table.count_rows())

print("\n3. a page shrinks")
by_url = Counter(r["url"] for r in records)
url = next(u for u, n in by_url.most_common() if n == 5)
page = [r for r in records if r["url"] == url]
print(f"  page {url.replace('https://webcrawlerapi.com', '')} has {len(page)} chunks; re-crawl returns the first 3")
new = page[:3]
upsert(new)
print("  after plain upsert, rows for that url:", table.count_rows(f"url = '{url}'"), "(2 stale)")
res = upsert(new, delete_filter=f"url = '{url}'")
print("  after upsert + when_not_matched_by_source_delete(url filter):", table.count_rows(f"url = '{url}'"))
print("  result:", res, "| total:", table.count_rows())

print("\n4. the same call without the url filter, on a scratch copy")
scratch = db.create_table("docs_scratch", data=records, mode="overwrite")
res = (scratch.merge_insert("id").when_matched_update_all().when_not_matched_insert_all()
       .when_not_matched_by_source_delete().execute(new))
print("  result:", res, "| rows left:", scratch.count_rows(), "of", len(records))
db.drop_table("docs_scratch")

print("\n5. versions and disk")
for i in range(5):
    upsert([dict(records[20 + i], text=records[20 + i]["text"] + f" v{i}")])
vs = table.list_versions()
print("  versions:", len(vs), "| last:", vs[-1])
n, size = stats()
print(f"  before optimize: {n} files, {size / 1024:.0f} KiB")
table.optimize(cleanup_older_than=timedelta(0))
n, size = stats()
print(f"  after optimize(cleanup_older_than=timedelta(0)): {n} files, {size / 1024:.0f} KiB, {len(table.list_versions())} version(s)")
print("  rows:", table.count_rows())