#!/usr/bin/env python3
"""Collect missing BBEH data only.

This script runs construct_router_data for only BBEH tasks, using the existing
cache to skip already-collected entries. Only missing entries will trigger API calls.

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

import json
import hashlib
from pathlib import Path
import pandas as pd
import yaml
import logging

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


def _query_hash(query: str, llm_name: str) -> str:
    """Generate a unique hash for a query+llm pair."""
    return hashlib.sha256(f"{query}||{llm_name}".encode()).hexdigest()[:16]


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 ("unified_qa_data_path", "llm_description_path", "saved_router_data_path", "llm_embedding_path"):
        if key in config and not Path(config[key]).is_absolute():
            config[key] = str(_MODULE_DIR / config[key])

    # Load BBEH-only data
    bbeh_path = _MODULE_DIR / "data" / "unified_qa_data_bbeh_only.csv"
    if not bbeh_path.exists():
        raise FileNotFoundError(f"BBEH-only data not found: {bbeh_path}")

    qa_df = pd.read_csv(bbeh_path)
    logger.info(f"Loaded BBEH data: {len(qa_df)} queries")

    # Load cache and LLM names
    cache_path = _MODULE_DIR / "data" / ".router_cache" / "response_cache.json"
    with open(cache_path) as f:
        cache = json.load(f)

    with open(config["llm_description_path"]) as f:
        llm_desc = json.load(f)
    llm_names = list(llm_desc.keys())

    # Count missing entries
    missing_count = 0
    for _, row in qa_df.iterrows():
        for llm in llm_names:
            cache_key = _query_hash(row['query'], llm)
            if cache_key not in cache:
                missing_count += 1

    logger.info(f"Missing cache entries: {missing_count}")
    logger.info(f"Expected API calls: {missing_count}")
    logger.info(f"Estimated cost: ${missing_count * 0.001:.2f}")

    if missing_count == 0:
        logger.info("All BBEH entries are cached. Nothing to do.")
        return

    # Import and run data_building with BBEH-only data
    # We need to temporarily override the qa_path in config
    logger.info("Starting data collection for missing BBEH entries...")

    from construct_router_data import data_building

    # Run the data builder - it will use cache for existing entries
    # New flat architecture: each (query, LLM) pair is independent
    # Per-provider limits: 5 OpenAI + 10 Google + 10 Anthropic = 25 max concurrent
    builder = data_building(
        qa_path=str(bbeh_path),
        llm_path=config["llm_description_path"],
        config=config,
        max_concurrent_api_calls=25,  # Workers to saturate provider semaphores
        skip_embeddings=True,  # Defer embedding computation to avoid API bottleneck
    )

    logger.info("BBEH data collection complete!")
    logger.info(f"Cache now has {len(builder.response_cache)} entries")


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