#!/usr/bin/env python3
"""
generate_embeddings_async.py

Async embedding generator:
- Takes a directory containing .jsonl files
- For each, loads comments
- If <embeddings_dir>/embeddings_<basename>.npy exists → skip
- Else computes embeddings asynchronously and saves them

Independent from semantic_variable_extractor.py.
"""

import os
import json
import argparse
import numpy as np
import asyncio
from tqdm.asyncio import tqdm_asyncio
from openai import AsyncOpenAI

# -----------------------------
# CONFIG
# -----------------------------
EMBED_MODEL = "text-embedding-3-small"
MIN_LEN = 50
MIN_WORDS = 6
BATCH_SIZE = 128
CONCURRENCY = 8

async_client = AsyncOpenAI()


# -----------------------------
# ASYNC EMBEDDING
# -----------------------------
async def embed_texts_async(texts, batch_size=BATCH_SIZE, concurrency=CONCURRENCY):
    """Fast async batch embedding (correct implementation)."""

    batches = [texts[i:i + batch_size] for i in range(0, len(texts), batch_size)]
    sem = asyncio.Semaphore(concurrency)

    async def embed_batch(batch):
        async with sem:
            resp = await async_client.embeddings.create(
                model=EMBED_MODEL,
                input=batch
            )
            return [np.array(d.embedding, dtype=np.float32) for d in resp.data]

    tasks = [embed_batch(b) for b in batches]

    # 🟢 Correct: run all async tasks with tqdm
    results = await tqdm_asyncio.gather(*tasks, desc="Async embedding")

    # flatten
    all_vecs = []
    for r in results:
        all_vecs.extend(r)

    return np.vstack(all_vecs)


# -----------------------------
# LOAD REVIEWS
# -----------------------------
def load_reviews(jsonl_path):
    reviews = []
    with open(jsonl_path, "r", encoding="utf-8") as f:
        for line in f:
            try:
                obj = json.loads(line)
                text = obj.get("review_text")
                if not text:
                    continue
                text = text.strip()
                if len(text) < MIN_LEN or len(text.split()) < MIN_WORDS:
                    continue
                reviews.append(text)
            except json.JSONDecodeError:
                continue

    # dedupe while preserving order
    return list(dict.fromkeys(reviews))


# -----------------------------
# PROCESS ONE FILE
# -----------------------------
async def process_file(jsonl_path, embeddings_dir):
    base = os.path.splitext(os.path.basename(jsonl_path))[0]
    emb_path = os.path.join(embeddings_dir, f"embeddings_{base}.npy")

    if os.path.exists(emb_path):
        print(f"⏩ Skipping {jsonl_path} — embeddings already exist at {emb_path}")
        return

    print(f"📄 Loading reviews from: {jsonl_path}")
    reviews = load_reviews(jsonl_path)

    if not reviews:
        print(f"⚠️ No valid reviews in {jsonl_path}")
        return

    print(f"🧠 Computing embeddings for {len(reviews)} reviews...")
    vecs = await embed_texts_async(reviews)

    print(f"💾 Saving → {emb_path}")
    np.save(emb_path, vecs)


# -----------------------------
# MAIN
# -----------------------------
async def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("jsonl_dir", help="Directory containing .jsonl files")
    parser.add_argument(
        "--embeddings_dir",
        required=True,
        help="Directory where embeddings_*.npy files will be stored"
    )
    args = parser.parse_args()

    jsonl_dir = args.jsonl_dir
    embeddings_dir = args.embeddings_dir

    if not os.path.isdir(jsonl_dir):
        print(f"❌ Not a directory: {jsonl_dir}")
        return

    # ensure embeddings directory exists
    os.makedirs(embeddings_dir, exist_ok=True)

    jsonl_files = [
        os.path.join(jsonl_dir, f)
        for f in os.listdir(jsonl_dir)
        if f.endswith(".jsonl")
    ]

    if not jsonl_files:
        print("❌ No .jsonl files found in directory.")
        return

    print(f"\n📁 Found {len(jsonl_files)} .jsonl files")
    print(f"📦 Embeddings will be saved to: {embeddings_dir}\n")

    # Process each file sequentially (embedding itself is concurrent)
    for path in jsonl_files:
        await process_file(path, embeddings_dir)

    print("\n✅ All embeddings generated.")


if __name__ == "__main__":
    asyncio.run(main())