import os
import time
import json
import torch
import mysql.connector
import numpy as np
from sentence_transformers import SentenceTransformer
from config import DB, EMBEDDING_MODEL

def run():
    print("=" * 60)
    print("🚀 Starting High-Speed Neural Semantic Embeddings Generator")
    print("=" * 60)

    device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device.upper()} | Model: {EMBEDDING_MODEL}")

    t0 = time.time()
    model = SentenceTransformer(EMBEDDING_MODEL, device=device)
    model.max_seq_length = 256
    dim = model.get_sentence_embedding_dimension() if hasattr(model, 'get_sentence_embedding_dimension') else model.get_embedding_dimension()
    print(f"✅ Model loaded successfully in {time.time() - t0:.2f}s (Embedding Dimension: {dim})")

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

    print("\nFetching chunks from database...")
    cur.execute("SELECT id, law_name, article_number, content FROM rag_chunks ORDER BY id ASC")
    chunks = cur.fetchall()
    total = len(chunks)
    print(f"Found {total} chunks to embed.")

    if total == 0:
        print("No chunks found.")
        cur.close()
        cnx.close()
        return

    # Prepare passages for e5 format: "passage: <text>"
    passages = []
    chunk_ids = []
    for c in chunks:
        chunk_ids.append(c["id"])
        law_title = c.get("law_name") or ""
        art = c.get("article_number") or ""
        text = c["content"] or ""
        prefix = f"{law_title} مادة {art}: " if (law_title or art) else ""
        passages.append(f"passage: {prefix}{text}")

    batch_size = 128
    print(f"\n⚡ Generating neural embeddings in batches of {batch_size} on {device.upper()}...")

    with torch.inference_mode():
        embeddings = model.encode(
            passages,
            batch_size=batch_size,
            show_progress_bar=True,
            normalize_embeddings=True,
            convert_to_numpy=True
        )

    print(f"\n✅ Generated {len(embeddings)} embeddings in {time.time() - t0:.2f}s.")

    print("\n💾 Updating MySQL database with neural embeddings...")
    update_cursor = cnx.cursor()
    update_query = "UPDATE rag_chunks SET embedding = %s WHERE id = %s"

    update_batch = []
    update_batch_size = 1000
    updated_count = 0

    for cid, emb in zip(chunk_ids, embeddings):
        emb_json = json.dumps([round(float(x), 5) for x in emb])
        update_batch.append((emb_json, cid))

        if len(update_batch) >= update_batch_size:
            update_cursor.executemany(update_query, update_batch)
            cnx.commit()
            updated_count += len(update_batch)
            print(f"  -> Saved {updated_count}/{total} chunks in database...")
            update_batch = []

    if update_batch:
        update_cursor.executemany(update_query, update_batch)
        cnx.commit()
        updated_count += len(update_batch)
        print(f"  -> Saved {updated_count}/{total} chunks in database...")

    update_cursor.close()
    cur.close()
    cnx.close()

    print("\n" + "=" * 60)
    print(f"🎉 SUCCESS: All {total} chunks now have real neural 768-dim semantic embeddings!")
    print(f"⏱️ Total time elapsed: {time.time() - t0:.2f} seconds")
    print("=" * 60)

if __name__ == "__main__":
    run()