import time
import json
import torch
import mysql.connector
from sentence_transformers import SentenceTransformer
from config import DB, EMBEDDING_MODEL

def main():
    print("Connecting to DB...")
    cnx = mysql.connector.connect(**DB)
    cur = cnx.cursor(dictionary=True)

    cur.execute("SELECT id, law_name, article_number, content FROM rag_chunks WHERE rag_source_id IN (30, 74) ORDER BY id ASC")
    chunks = cur.fetchall()
    print(f"Loaded {len(chunks)} chunks to re-embed.")

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

    passages = []
    for c in chunks:
        law_name = c.get('law_name') or ''
        art_num = c.get('article_number') or ''
        prefix = f"{law_name} مادة {art_num}: " if (law_name or art_num) else ""
        content = c.get('content') or ''
        passages.append(f"passage: {prefix}{content}")

    print("Encoding passages...")
    t0 = time.time()
    with torch.inference_mode():
        embeddings = model.encode(passages, batch_size=64, normalize_embeddings=True, convert_to_numpy=True)
    print(f"Encoded in {time.time() - t0:.2f}s")

    print("Updating database...")
    update_cur = cnx.cursor()
    for chunk, emb in zip(chunks, embeddings):
        emb_json = json.dumps(emb.tolist())
        update_cur.execute("UPDATE rag_chunks SET embedding = %s WHERE id = %s", (emb_json, chunk["id"]))

    cnx.commit()
    print("✅ All embeddings updated successfully!")
    cur.close()
    update_cur.close()
    cnx.close()

if __name__ == '__main__':
    main()