#!/usr/bin/env python3
"""Regenerate LLM description embeddings using OpenAI text-embedding-3-large (512-dim).

This script regenerates LLM embeddings using OpenAI's text-embedding-3-large model
with 512 dimensions, matching the query embeddings used in training.

Uses Ember's model registry for credentials (reads from ~/.ember/config.yaml).
"""

import json
import os
import pickle
import sys
from pathlib import Path

import numpy as np

_SCRIPT_DIR = Path(__file__).resolve().parent
_MODULE_DIR = _SCRIPT_DIR.parent

# Configure embedding model via env vars (used by data_processing/utils.py)
os.environ["EGR_TEXT_ENCODER"] = "openai/text-embedding-3-large"
os.environ["EGR_EMBEDDING_DIM"] = "512"

# Add parent to path for imports
sys.path.insert(0, str(_MODULE_DIR))
from data_processing.utils import get_embedding

EMBEDDING_DIM = 512


def main():
    # Load LLM descriptions
    llm_desc_path = _MODULE_DIR / "configs" / "LLM_Descriptions.json"
    with open(llm_desc_path) as f:
        llm_descriptions = json.load(f)

    print(f"Loaded {len(llm_descriptions)} LLM descriptions")
    for name in llm_descriptions:
        print(f"  - {name}")

    # Generate embeddings for each LLM's feature description
    llm_names = list(llm_descriptions.keys())
    texts = [llm_descriptions[name]["feature"] for name in llm_names]

    print(f"\nGenerating embeddings for {len(texts)} LLMs using openai/text-embedding-3-large ({EMBEDDING_DIM}-dim)...")
    embeddings = get_embedding(texts, dimensions=EMBEDDING_DIM)
    print(f"Embeddings shape: {embeddings.shape}")

    # Verify dimension
    assert embeddings.shape == (len(llm_names), EMBEDDING_DIM), (
        f"Expected ({len(llm_names)}, {EMBEDDING_DIM}), got {embeddings.shape}"
    )

    # Save new embeddings
    output_path = _MODULE_DIR / "configs" / f"llm_description_embedding_openai_{EMBEDDING_DIM}.pkl"
    with open(output_path, "wb") as f:
        pickle.dump(embeddings, f)
    print(f"Saved to {output_path}")

    # Also update the default embedding file
    default_path = _MODULE_DIR / "configs" / "llm_description_embedding.pkl"
    backup_path = _MODULE_DIR / "configs" / "llm_description_embedding_backup.pkl"

    # Backup old embeddings
    if default_path.exists():
        import shutil
        shutil.copy(default_path, backup_path)
        print(f"Backed up old embeddings to {backup_path}")

    with open(default_path, "wb") as f:
        pickle.dump(embeddings, f)
    print(f"Updated {default_path}")

    # Verify
    with open(default_path, "rb") as f:
        loaded = pickle.load(f)
    print(f"Verified: shape = {loaded.shape}")


if __name__ == "__main__":
    main()
