#!/usr/bin/env python3
"""Backfill embeddings in router_data.csv from the embedding cache.

Run this after data collection to fill in any missing embeddings from the cache.
If embeddings are missing from the cache, it will compute them.

Usage:
    python -m embergraphrouter.data_processing.backfill_embeddings
"""

import hashlib
import json
import logging
from pathlib import Path

import pandas as pd
import yaml
from tqdm import tqdm

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

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

# Handle both package import and direct script execution
try:
    from .utils import get_embedding, savepkl, loadjson
except ImportError:
    from utils import get_embedding, savepkl, loadjson


def _text_hash(text: str) -> str:
    """Generate a unique hash for a text."""
    return hashlib.sha256(text.encode()).hexdigest()[:16]


def load_embedding_cache(cache_dir: Path) -> dict:
    """Load embedding cache from disk."""
    cache_file = cache_dir / "embedding_cache.json"
    if cache_file.exists():
        try:
            return json.loads(cache_file.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            return {}
    return {}


def save_embedding_cache(cache: dict, cache_dir: Path) -> None:
    """Save embedding cache to disk."""
    cache_file = cache_dir / "embedding_cache.json"
    cache_file.write_text(json.dumps(cache, indent=2), encoding="utf-8")


def backfill_embeddings(
    router_data_path: Path,
    cache_dir: Path,
    llm_description_path: Path,
    llm_embedding_path: Path,
    batch_size: int = 1024,
) -> None:
    """Fill in missing embeddings in router data from cache, computing if needed.

    Args:
        router_data_path: Path to router_data.csv.
        cache_dir: Directory containing embedding_cache.json.
        llm_description_path: Path to LLM_Descriptions.json.
        llm_embedding_path: Path to save LLM description embeddings.
        batch_size: Number of texts to embed per API call if computing.
    """
    logger.info("Loading router data from %s", router_data_path)
    df = pd.read_csv(router_data_path, low_memory=False)

    # Load embedding cache
    cache = load_embedding_cache(cache_dir)
    logger.info("Embedding cache has %d entries", len(cache))

    # Check which embeddings need filling
    needs_query_emb = df["query_embedding"].isna() | (df["query_embedding"] == "")
    needs_task_emb = (
        df["task_description_embedding"].isna()
        | (df["task_description_embedding"] == "")
    )

    logger.info(
        "Rows needing embeddings: %d query, %d task_description",
        needs_query_emb.sum(),
        needs_task_emb.sum(),
    )

    if not needs_query_emb.any() and not needs_task_emb.any():
        logger.info("All embeddings already present. Nothing to do.")
    else:
        # Get unique texts that need embedding
        unique_queries = df.loc[needs_query_emb, "query"].drop_duplicates().tolist()
        unique_tasks = (
            df.loc[needs_task_emb, "task_description"].drop_duplicates().tolist()
        )

        # Check which are in cache vs need computing
        queries_in_cache = [q for q in unique_queries if _text_hash(q) in cache]
        queries_to_compute = [q for q in unique_queries if _text_hash(q) not in cache]
        tasks_in_cache = [t for t in unique_tasks if _text_hash(t) in cache]
        tasks_to_compute = [t for t in unique_tasks if _text_hash(t) not in cache]

        logger.info(
            "Queries: %d in cache, %d to compute",
            len(queries_in_cache),
            len(queries_to_compute),
        )
        logger.info(
            "Tasks: %d in cache, %d to compute",
            len(tasks_in_cache),
            len(tasks_to_compute),
        )

        # Compute missing embeddings with incremental checkpointing
        texts_to_compute = queries_to_compute + tasks_to_compute
        if texts_to_compute:
            logger.info("Computing %d missing embeddings...", len(texts_to_compute))
            checkpoint_interval = 5  # Save every 5 batches (~5K embeddings)
            batch_count = 0
            for i in tqdm(range(0, len(texts_to_compute), batch_size), desc="Computing"):
                batch = texts_to_compute[i : i + batch_size]
                embeddings = get_embedding(batch)
                for text, emb in zip(batch, embeddings):
                    cache[_text_hash(text)] = str(emb.tolist())
                batch_count += 1
                # Checkpoint every N batches for resumption
                if batch_count % checkpoint_interval == 0:
                    save_embedding_cache(cache, cache_dir)
                    logger.info("Checkpoint: saved %d embeddings to cache", len(cache))
            # Final save
            save_embedding_cache(cache, cache_dir)
            logger.info("Updated embedding cache with %d new entries", len(texts_to_compute))

        # Build lookup maps
        query_emb_map = {q: cache[_text_hash(q)] for q in unique_queries}
        task_emb_map = {t: cache[_text_hash(t)] for t in unique_tasks}

        # Update dataframe
        logger.info("Updating dataframe...")
        df.loc[needs_query_emb, "query_embedding"] = df.loc[
            needs_query_emb, "query"
        ].map(query_emb_map)
        df.loc[needs_task_emb, "task_description_embedding"] = df.loc[
            needs_task_emb, "task_description"
        ].map(task_emb_map)

        # Save updated CSV
        df.to_csv(router_data_path, index=False)
        logger.info("Saved updated router data to %s", router_data_path)

    # Compute LLM description embeddings if needed
    if not llm_embedding_path.exists():
        logger.info("Computing LLM description embeddings...")
        llm_desc = loadjson(str(llm_description_path))
        llm_names = list(llm_desc.keys())
        all_llm_descriptions = [llm_desc[name]["feature"] for name in llm_names]
        llm_description_embedding = get_embedding(all_llm_descriptions)
        savepkl(llm_description_embedding, str(llm_embedding_path))
        logger.info("Saved LLM embeddings to %s", llm_embedding_path)
    else:
        logger.info("LLM description embeddings already exist at %s", llm_embedding_path)

    logger.info("Embedding backfill complete!")


def main():
    # Load config
    config_path = _MODULE_DIR / "configs" / "config.yaml"
    with open(config_path, "r", encoding="utf-8") as f:
        config = yaml.safe_load(f)

    # Resolve paths
    for key in (
        "saved_router_data_path",
        "llm_description_path",
        "llm_embedding_path",
    ):
        if key in config and not Path(config[key]).is_absolute():
            config[key] = str(_MODULE_DIR / config[key])

    cache_dir = _MODULE_DIR / "data" / ".router_cache"

    backfill_embeddings(
        router_data_path=Path(config["saved_router_data_path"]),
        cache_dir=cache_dir,
        llm_description_path=Path(config["llm_description_path"]),
        llm_embedding_path=Path(config["llm_embedding_path"]),
    )


if __name__ == "__main__":
    import os
    import sys

    sys.path.insert(0, str(Path(__file__).resolve().parent))
    os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
    main()
