#!/usr/bin/env python3
"""Embed arxiv abstracts using static-retrieval-mrl-en-v1 (1024-dim)."""

import subprocess
import time

MODEL_NAME = "sentence-transformers/static-retrieval-mrl-en-v1"
BATCH_SIZE = 512
TSV_PATH = "/tmp/embeddings.tsv"

def psql(sql, timeout=60):
    r = subprocess.run(
        ['podman', 'exec', 'arxiv-pg', 'psql', '-U', 'postgres', '-d', 'arxiv',
         '-t', '-A', '-c', sql],
        capture_output=True, text=True, timeout=timeout
    )
    return r.stdout.strip(), r.stderr.strip()

def main():
    from sentence_transformers import SentenceTransformer

    print(f"Model: {MODEL_NAME}")
    model = SentenceTransformer(MODEL_NAME)
    dims = model.get_sentence_embedding_dimension()
    print(f"Dims: {dims}")

    # Count remaining
    count, _ = psql("SELECT COUNT(*) FROM arxiv_papers WHERE embedding IS NULL")
    total = int(count)
    print(f"Papers to embed: {total}")

    if total == 0:
        print("All done.")
        return

    # Fetch all at once
    print("Fetching papers...")
    out, _ = psql(
        "SELECT paper_id || E'\\t' || COALESCE(title,'') || ' ' || COALESCE(abstract,'') "
        "FROM arxiv_papers WHERE embedding IS NULL ORDER BY paper_id",
        timeout=600
    )
    
    pids, texts = [], []
    for line in out.strip().split('\n'):
        parts = line.split('\t', 1)
        if len(parts) == 2:
            pids.append(parts[0])
            texts.append(parts[1].strip())
    
    print(f"Fetched {len(pids)} papers")

    # Embed and write TSV
    start = time.time()
    f = open(TSV_PATH, 'w')
    embedded = 0

    for i in range(0, len(pids), BATCH_SIZE):
        batch_pids = pids[i:i+BATCH_SIZE]
        batch_texts = texts[i:i+BATCH_SIZE]

        embeddings = model.encode(batch_texts, show_progress_bar=False, 
                                  normalize_embeddings=True, batch_size=256)

        for pid, emb in zip(batch_pids, embeddings):
            emb_str = '[' + ','.join(f'{x:.6f}' for x in emb) + ']'
            f.write(f"{pid}\t{emb_str}\n")

        embedded += len(batch_pids)
        elapsed = time.time() - start
        rate = embedded / elapsed if elapsed > 0 else 0
        eta = (len(pids) - embedded) / rate if rate > 0 else 0
        
        if embedded % (BATCH_SIZE * 10) == 0 or embedded >= len(pids):
            print(f"  {embedded}/{len(pids)} ({embedded*100/len(pids):.1f}%) | "
                  f"{rate:.0f}/s | ETA: {eta/60:.1f} min")

    f.close()
    elapsed = time.time() - start
    print(f"Embedding done: {embedded} in {elapsed/60:.1f} min ({embedded/elapsed:.0f}/s)")

    # Bulk load
    print("Loading into database...")
    subprocess.run(['podman', 'cp', TSV_PATH, 'arxiv-pg:/tmp/embeddings.tsv'],
                  capture_output=True, timeout=300)

    psql(f"CREATE TEMP TABLE IF NOT EXISTS tmp_emb (pid text, emb vector({dims})); TRUNCATE tmp_emb;",
         timeout=30)

    r = subprocess.run(
        ['podman', 'exec', '-i', 'arxiv-pg', 'psql', '-U', 'postgres', '-d', 'arxiv',
         '-c', f"\\copy tmp_emb FROM '/tmp/embeddings.tsv' WITH (FORMAT csv, DELIMITER E'\\t')"],
        capture_output=True, text=True, timeout=600
    )
    print(f"COPY: {r.stdout.strip() or r.stderr.strip()}")

    psql("UPDATE arxiv_papers SET embedding = t.emb FROM tmp_emb t WHERE arxiv_papers.paper_id = t.pid;",
         timeout=600)

    count, _ = psql("SELECT COUNT(*) FROM arxiv_papers WHERE embedding IS NOT NULL")
    print(f"Total embedded: {count}")
    print(f"Total time: {(time.time() - start)/60:.1f} min")

if __name__ == "__main__":
    main()