#!/usr/bin/env python3
"""Launch v2 data collection pipeline.

Runs the full pipeline:
  1. Generate unified QA data (datasets → CSV)
  2. Estimate API cost
  3. Collect LLM responses (8 models × N queries)

Usage:
    # Full pipeline
    python scripts/launch_v2_datacollection.py

    # Skip QA generation (already done), go straight to LLM collection
    python scripts/launch_v2_datacollection.py --skip-qa

    # Smoke test with 5 samples per task
    python scripts/launch_v2_datacollection.py --smoke-test
"""

import argparse
import json
import logging
import sys
from pathlib import Path

import yaml

_MODULE_DIR = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_MODULE_DIR.parent))  # packages/
sys.path.insert(0, str(_MODULE_DIR.parent.parent / "src"))  # src/

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("v2_datacollection")

CONFIG_PATH = _MODULE_DIR / "configs" / "config_v2_2026.yaml"


def main():
    parser = argparse.ArgumentParser(description="v2 data collection pipeline")
    parser.add_argument(
        "--skip-qa",
        action="store_true",
        help="Skip QA generation (use existing unified_qa_data_v2.csv)",
    )
    parser.add_argument(
        "--smoke-test",
        action="store_true",
        help="Run with 5 samples per task (quick verification)",
    )
    parser.add_argument(
        "--skip-embeddings",
        action="store_true",
        help="Skip embedding generation (defer to separate step)",
    )
    parser.add_argument(
        "--max-query-workers",
        type=int,
        default=10,
        help="Max parallel queries (default: 10)",
    )
    parser.add_argument(
        "--max-llm-workers",
        type=int,
        default=8,
        help="Max parallel LLM calls per query (default: 8)",
    )
    parser.add_argument(
        "--max-concurrent",
        type=int,
        default=80,
        help="Max concurrent API calls across all providers (default: 80)",
    )
    args = parser.parse_args()

    # Load v2 config
    logger.info("Loading config from %s", CONFIG_PATH)
    config = yaml.safe_load(CONFIG_PATH.read_text(encoding="utf-8"))

    # Resolve relative paths to absolute
    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])

    sample_size = 5 if args.smoke_test else int(config["dataset_sample_size"])
    logger.info("Sample size per task: %d", sample_size)

    # Load LLM descriptions for verification
    llm_desc = json.loads(Path(config["llm_description_path"]).read_text())
    llm_names = list(llm_desc.keys())
    logger.info("LLM pool (%d models):", len(llm_names))
    for name in llm_names:
        pricing = llm_desc[name]
        logger.info(
            "  %s: $%.2f/$%.2f per 1M tokens (%s)",
            name,
            pricing["input_price"],
            pricing["output_price"],
            pricing["model"],
        )

    # ── Step 1: Generate unified QA data ──
    if not args.skip_qa:
        logger.info("=" * 60)
        logger.info("Step 1: Generating unified QA dataset")
        logger.info("=" * 60)

        from embergraphrouter.data_processing.multidata_unify import (
            generate_unified_qa_dataset,
        )

        df = generate_unified_qa_dataset(
            output_path=config["unified_qa_data_path"],
            sample_size=sample_size,
            show_progress=True,
            wandb_run=None,
        )
        logger.info(
            "QA dataset: %d rows, %d tasks",
            len(df),
            df["task_id"].nunique(),
        )
    else:
        import pandas as pd

        qa_path = config["unified_qa_data_path"]
        if not Path(qa_path).exists():
            logger.error("QA data not found at %s (run without --skip-qa)", qa_path)
            sys.exit(1)
        df = pd.read_csv(qa_path)
        logger.info(
            "Loaded existing QA data: %d rows, %d tasks",
            len(df),
            df["task_id"].nunique(),
        )

    # ── Cost estimate ──
    total_queries = len(df)
    total_calls = total_queries * len(llm_names)
    avg_input_tokens = 200  # Conservative estimate
    avg_output_tokens = 300
    estimated_cost = 0.0
    for name in llm_names:
        pricing = llm_desc[name]
        model_cost = (
            (avg_input_tokens / 1_000_000) * pricing["input_price"]
            + (avg_output_tokens / 1_000_000) * pricing["output_price"]
        ) * total_queries
        estimated_cost += model_cost

    logger.info("=" * 60)
    logger.info("Cost estimate:")
    logger.info("  Total queries: %d", total_queries)
    logger.info("  Total API calls: %d (queries × %d models)", total_calls, len(llm_names))
    logger.info("  Estimated cost: $%.2f", estimated_cost)
    logger.info("=" * 60)

    # ── Step 2: Collect LLM responses ──
    logger.info("Step 2: Collecting LLM responses")
    logger.info("=" * 60)

    from embergraphrouter.data_processing.construct_router_data import data_building

    builder = data_building(
        qa_path=config["unified_qa_data_path"],
        llm_path=config["llm_description_path"],
        config=config,
        max_query_workers=args.max_query_workers,
        max_llm_workers=args.max_llm_workers,
        max_concurrent_api_calls=args.max_concurrent,
        skip_embeddings=args.skip_embeddings,
    )
    builder.construct_data_with_LLM()

    logger.info("=" * 60)
    logger.info("Data collection complete!")
    logger.info("Output: %s", config["saved_router_data_path"])
    logger.info("=" * 60)


if __name__ == "__main__":
    main()
