#!/usr/bin/env python3
"""Compute embeddings for router queries independently of LLM data collection.

This script runs as a separate process to pre-compute and cache embeddings
for all queries and task descriptions. It writes to embedding_cache.json
which the data_building class reads from.

Run this in parallel with collect_missing_bbeh.py:
    Terminal 1: python -m embergraphrouter.data_processing.compute_embeddings
    Terminal 2: python -m embergraphrouter.data_processing.collect_missing_bbeh

The embedding cache is independent from the response cache, so both processes
can run concurrently without blocking each other.

Usage:
    python -m embergraphrouter.data_processing.compute_embeddings [--watch]

Options:
    --watch     Continuously poll for new queries (useful during data collection)
    --batch-size N  Number of texts to embed per API call (default: 32)
"""

import argparse
import hashlib
import json
import logging
import time
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 atomically."""
    cache_file = cache_dir / "embedding_cache.json"
    temp_file = cache_file.with_suffix(".json.tmp")
    temp_file.write_text(json.dumps(cache, indent=2), encoding="utf-8")
    temp_file.rename(cache_file)


def compute_embeddings_for_dataset(
    qa_path: Path,
    cache_dir: Path,
    llm_description_path: Path,
    llm_embedding_path: Path,
    batch_size: int = 32,
) -> int:
    """Compute embeddings for all queries and task descriptions in a dataset.

    Args:
        qa_path: Path to the unified QA CSV file.
        cache_dir: Directory for embedding cache.
        llm_description_path: Path to LLM descriptions JSON.
        llm_embedding_path: Path to save LLM embeddings pickle.
        batch_size: Number of texts to embed per API call.

    Returns:
        Number of new embeddings computed.
    """
    cache_dir.mkdir(parents=True, exist_ok=True)
    cache = load_embedding_cache(cache_dir)
    initial_count = len(cache)

    # Load dataset
    df = pd.read_csv(qa_path, low_memory=False)
    logger.info("Loaded %d rows from %s", len(df), qa_path)

    # Collect unique texts that need embedding
    all_queries = df["query"].drop_duplicates().tolist()
    all_tasks = df["task_description"].drop_duplicates().tolist()
    all_texts = list(set(all_queries + all_tasks))

    # Filter to texts not in cache
    texts_to_embed = [t for t in all_texts if _text_hash(t) not in cache]

    if not texts_to_embed:
        logger.info("All %d texts already cached. Nothing to compute.", len(all_texts))
    else:
        logger.info(
            "Computing embeddings for %d/%d texts (%d cached)",
            len(texts_to_embed),
            len(all_texts),
            len(all_texts) - len(texts_to_embed),
        )

        # Compute in batches
        for i in tqdm(range(0, len(texts_to_embed), batch_size), desc="Embedding batches"):
            batch = texts_to_embed[i : i + batch_size]
            try:
                embeddings = get_embedding(batch)
                for text, emb in zip(batch, embeddings):
                    cache[_text_hash(text)] = str(emb.tolist())
                # Save after each batch for resilience
                save_embedding_cache(cache, cache_dir)
            except Exception as e:
                logger.error("Error embedding batch %d: %s", i // batch_size, e)
                # Continue with next batch

    new_count = len(cache) - initial_count
    logger.info("Computed %d new embeddings. Cache now has %d entries.", new_count, len(cache))

    # 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)

    return new_count


def watch_and_compute(
    qa_path: Path,
    cache_dir: Path,
    llm_description_path: Path,
    llm_embedding_path: Path,
    batch_size: int = 32,
    poll_interval: int = 30,
) -> None:
    """Continuously compute embeddings, polling for new data.

    Args:
        qa_path: Path to the unified QA CSV file.
        cache_dir: Directory for embedding cache.
        llm_description_path: Path to LLM descriptions JSON.
        llm_embedding_path: Path to save LLM embeddings pickle.
        batch_size: Number of texts to embed per API call.
        poll_interval: Seconds between polls.
    """
    logger.info("Starting embedding worker in watch mode (poll every %ds)", poll_interval)
    logger.info("Press Ctrl+C to stop")

    try:
        while True:
            new_count = compute_embeddings_for_dataset(
                qa_path=qa_path,
                cache_dir=cache_dir,
                llm_description_path=llm_description_path,
                llm_embedding_path=llm_embedding_path,
                batch_size=batch_size,
            )
            if new_count == 0:
                logger.info("No new texts to embed. Sleeping %ds...", poll_interval)
            time.sleep(poll_interval)
    except KeyboardInterrupt:
        logger.info("Embedding worker stopped.")


def main():
    parser = argparse.ArgumentParser(
        description="Compute embeddings for router queries"
    )
    parser.add_argument(
        "--watch",
        action="store_true",
        help="Continuously poll for new queries",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=32,
        help="Number of texts to embed per API call (default: 32)",
    )
    parser.add_argument(
        "--qa-path",
        type=str,
        default=None,
        help="Path to QA data CSV (default: from config)",
    )
    args = parser.parse_args()

    # 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 (
        "unified_qa_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])

    qa_path = Path(args.qa_path) if args.qa_path else Path(config["unified_qa_data_path"])
    cache_dir = _MODULE_DIR / "data" / ".router_cache"
    llm_description_path = Path(config["llm_description_path"])
    llm_embedding_path = Path(config["llm_embedding_path"])

    if args.watch:
        watch_and_compute(
            qa_path=qa_path,
            cache_dir=cache_dir,
            llm_description_path=llm_description_path,
            llm_embedding_path=llm_embedding_path,
            batch_size=args.batch_size,
        )
    else:
        compute_embeddings_for_dataset(
            qa_path=qa_path,
            cache_dir=cache_dir,
            llm_description_path=llm_description_path,
            llm_embedding_path=llm_embedding_path,
            batch_size=args.batch_size,
        )


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

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