import asyncio
import sys

import dotenv
import pandas as pd
from tqdm.auto import tqdm

dotenv.load_dotenv()


async def embed_texts(texts: list[str], model: str = "", delay=0.03):
    from openai import AsyncOpenAI

    client = AsyncOpenAI()
    tasks = []

    async def _embed_one(text: str, pbar):
        text = text[:20_000]
        res = await client.embeddings.create(input=text, model="text-embedding-3-small")
        pbar.update(1)
        return res.data[0].embedding

    pbar = tqdm(total=len(texts))
    with pbar:
        for text in texts:
            task = asyncio.create_task(_embed_one(text, pbar))
            tasks.append(task)
            await asyncio.sleep(delay)

        results = await asyncio.gather(*tasks)

    return results


if __name__ == "__main__":
    input_jsonl = sys.argv[1]
    df = pd.read_json(input_jsonl, lines=True, orient="records")
    results = asyncio.run(embed_texts(df["content"].tolist()))

    df["embedding"] = results

    new_file = input_jsonl.replace(".jsonl", "_embeddings.jsonl")

    df.to_json(new_file, lines=True, orient="records")
