import importlib.util
import json
import sys
import tempfile
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from test_archive_performance import MeasuredDatabase, seed_archive
from archive_store import ArchiveStore

class TimedDatabase(MeasuredDatabase):
    deadline = 0
    def connect(self):
        db = super().connect()
        if self.deadline:
            db.set_progress_handler(lambda: time.perf_counter() > self.deadline, 1000)
        return db

spec=importlib.util.spec_from_file_location("archive_before", Path(__file__).parent/"archive-performance-before/archive_store.py")
old=importlib.util.module_from_spec(spec)
spec.loader.exec_module(old)
report={"messages":1_000_000,"conversations":2000,"people":2000,"data":"synthetic"}
with tempfile.TemporaryDirectory(prefix="archive-million-") as tmp:
    database=TimedDatabase(Path(tmp)/"benchmark.db")
    database.initialize("", seed_admin=False)
    before=old.ArchiveStore(database)
    before.initialize()
    start=time.perf_counter()
    seed_archive(database,report["messages"])
    report["seed_seconds"]=round(time.perf_counter()-start,3)
    print("Seeded one million synthetic messages",flush=True)
    def measure(store):
        result={}
        for name,fn in [("conversations",store.conversations),("people",lambda:store.people(200)),("stats",store.stats),("exports",store.export_jobs),("media",store.media_items)]:
            database.deadline=time.perf_counter()+5
            start=time.perf_counter()
            try:
                fn()
                result[name]={"seconds":round(time.perf_counter()-start,4),"status":"ok"}
            except Exception as exc:
                result[name]={"seconds":round(time.perf_counter()-start,4),"status":str(exc)}
            finally:
                database.deadline=0
            print(name,result[name],flush=True)
        return result
    report["before"]=measure(before)
    after=ArchiveStore(database)
    start=time.perf_counter()
    after.initialize()
    report["index_migration_seconds"]=round(time.perf_counter()-start,3)
    report["after"]=measure(after)
    report["after_warm"]=measure(after)
Path(__file__).with_name("archive-million-benchmark.json").write_text(json.dumps(report,ensure_ascii=False,indent=2),encoding="utf-8")
print(json.dumps(report,ensure_ascii=False,indent=2))