import sys, os, json, time, torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import mysql.connector
import urllib.request
from sentence_transformers import SentenceTransformer
from config import DB, EMBEDDING_MODEL

def main():
    device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Loading embedding model '{EMBEDDING_MODEL}' on {device}...")
    model = SentenceTransformer(EMBEDDING_MODEL, device=device)
    model.max_seq_length = 256

    cnx = mysql.connector.connect(**DB)
    cur = cnx.cursor(dictionary=True)

    print("Fetching circular chunks (sources: 256, 257, 258, 259)...")
    cur.execute("SELECT id, content FROM rag_chunks WHERE rag_source_id IN (256, 257, 258, 259)")
    rows = cur.fetchall()
    total = len(rows)
    print(f"Found {total} circular chunks to encode.")

    batch_size = 64
    updated = 0
    t0 = time.time()

    for i in range(0, total, batch_size):
        batch = rows[i:i + batch_size]
        passages = ["passage: " + (r["content"][:1000] if r["content"] else "") for r in batch]
        
        with torch.inference_mode():
            embs = model.encode(passages, batch_size=batch_size, normalize_embeddings=True, show_progress_bar=False)

        update_cur = cnx.cursor()
        for r, emb in zip(batch, embs):
            emb_json = json.dumps(emb.tolist())
            update_cur.execute("UPDATE rag_chunks SET embedding = %s WHERE id = %s", (emb_json, r["id"]))
        cnx.commit()
        update_cur.close()

        updated += len(batch)
        print(f"  Processed {updated}/{total} chunks ({(updated/total)*100:.1f}%) in {time.time() - t0:.1f}s")

    cur.close()
    cnx.close()
    print(f"✓ Completed encoding {updated} circular chunks in {time.time() - t0:.2f}s.")

    # Hot-reload search_api memory
    try:
        req = urllib.request.Request("http://127.0.0.1:8005/reload", method="POST")
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = json.loads(resp.read().decode())
            print(f"✓ Hot-reloaded search API: {data}")
    except Exception as e:
        print(f"Note: Could not auto-reload search API ({e}), will be reloaded on next query or restart.")

if __name__ == "__main__":
    main()