#!/usr/bin/env python
"""Embed chunked MTSamples content with NVIDIA NIM embeddings."""

from __future__ import annotations

import argparse
import json
import os
from datetime import datetime, timezone
from typing import Dict, Iterator, List

import numpy as np

def iter_jsonl(path: str) -> Iterator[Dict]:
    """Iterate over JSONL records from a file."""
    with open(path, "r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                yield json.loads(line)


def load_client(model_name: str, api_key: str, truncate: str):
    try:
        from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
    except Exception as exc:  # pragma: no cover - import guard for missing deps
        raise SystemExit(
            "Failed to import langchain_nvidia_ai_endpoints. "
            "Verify the env and run `python -c \"from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings\"`.\n"
            f"Import error: {exc}"
        ) from exc

    return NVIDIAEmbeddings(model=model_name, api_key=api_key, truncate=truncate)


def normalize_embeddings(embeddings: np.ndarray) -> np.ndarray:
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    return embeddings / np.clip(norms, a_min=1e-12, a_max=None)


def embed_chunks(
    input_path: str,
    output_dir: str,
    model_name: str,
    api_key: str,
    truncate: str,
    batch_size: int,
    normalize: bool,
) -> Dict:
    os.makedirs(output_dir, exist_ok=True)
    output_path = os.path.join(output_dir, "embeddings_mtsamples.jsonl")

    client = load_client(model_name, api_key, truncate)

    total_records = 0
    total_embedded = 0
    skipped = 0

    batch_records: List[Dict] = []
    batch_texts: List[str] = []

    def flush_batch(out_handle) -> None:
        nonlocal total_embedded
        if not batch_records:
            return

        embeddings = client.embed_documents(batch_texts)
        embeddings = np.asarray(embeddings, dtype=float)
        if normalize:
            embeddings = normalize_embeddings(embeddings)
        if embeddings.shape[0] != len(batch_records):
            raise RuntimeError(
                f"Embedding count mismatch: {embeddings.shape[0]} vs {len(batch_records)}"
            )

        for record, embedding in zip(batch_records, embeddings):
            out_record = dict(record)
            out_record["embedding"] = [float(x) for x in embedding.tolist()]
            out_handle.write(json.dumps(out_record, ensure_ascii=True) + "\n")
            total_embedded += 1

        batch_records.clear()
        batch_texts.clear()

    with open(output_path, "w", encoding="utf-8") as out_handle:
        for record in iter_jsonl(input_path):
            total_records += 1
            content = record.get("content", "")
            if not content or not content.strip():
                skipped += 1
                continue

            batch_records.append(record)
            batch_texts.append(content)

            if len(batch_records) >= batch_size:
                flush_batch(out_handle)

        flush_batch(out_handle)

    stats = {
        "input_file": os.path.abspath(input_path),
        "output_file": os.path.abspath(output_path),
        "generated_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "model": model_name,
        "truncate": truncate,
        "batch_size": batch_size,
        "normalize": normalize,
        "counts": {
            "input_records": total_records,
            "embedded_records": total_embedded,
            "skipped_records": skipped,
        },
    }

    stats_path = os.path.join(output_dir, "embedding_stats.json")
    with open(stats_path, "w", encoding="utf-8") as stats_handle:
        json.dump(stats, stats_handle, indent=2, ensure_ascii=True)

    return stats


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Embed chunked MTSamples JSONL with NVIDIA NIM embeddings."
    )
    base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
    parser.add_argument(
        "--input",
        default=os.path.join(base_dir, "02_chunking", "output", "chunked_mtsamples.jsonl"),
        help="Path to chunked JSONL input file.",
    )
    parser.add_argument(
        "--output-dir",
        default="output",
        help="Output directory for embeddings JSONL and stats.",
    )
    parser.add_argument(
        "--model",
        default="nvidia/llama-3.2-nv-embedqa-1b-v2",
        help="NVIDIA embedding model name.",
    )
    parser.add_argument(
        "--api-key",
        default=os.getenv("NVIDIA_API_KEY"),
        help="NVIDIA API key (or set NVIDIA_API_KEY).",
    )
    parser.add_argument(
        "--truncate",
        default="NONE",
        help="Truncation mode for the NVIDIA endpoint (e.g., NONE, START, END).",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=32,
        help="Batch size for embedding inference.",
    )
    parser.add_argument(
        "--normalize",
        action="store_true",
        help="Normalize embeddings for cosine similarity.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if not args.api_key:
        raise SystemExit(
            "Missing NVIDIA API key. Set NVIDIA_API_KEY or pass --api-key."
        )
    stats = embed_chunks(
        input_path=args.input,
        output_dir=args.output_dir,
        model_name=args.model,
        api_key=args.api_key,
        truncate=args.truncate,
        batch_size=args.batch_size,
        normalize=args.normalize,
    )

    print("Embedding complete!")
    print(f"  Input records:   {stats['counts']['input_records']}")
    print(f"  Embedded:        {stats['counts']['embedded_records']}")
    print(f"  Skipped:         {stats['counts']['skipped_records']}")
    print(f"  Output: {stats['output_file']}")


if __name__ == "__main__":
    main()