#!/usr/bin/env python3
"""Prepare data and LLM embeddings for Gemini embedding sweep.

CRITICAL: This script ensures ALL embeddings use the SAME source model.
- Query embeddings: Gemini text-embedding-004 at specified dimension
- Task embeddings: Gemini text-embedding-004 at specified dimension
- LLM embeddings: Gemini text-embedding-004 at specified dimension

NO mixing of embedding models (MiniLM, OpenAI, etc.) is allowed.
NO dimension padding or truncation is performed.

Each variant uses native Gemini embeddings at the exact dimension specified:
- gemini_256: All embeddings are 256-dim from Gemini
- gemini_512: All embeddings are 512-dim from Gemini
- gemini_768: All embeddings are 768-dim from Gemini
"""

import json
import logging
import os
import pickle
import sys
from pathlib import Path
from typing import Tuple

import numpy as np
import pandas as pd
import yaml

_SCRIPT_DIR = Path(__file__).resolve().parent
_MODULE_DIR = _SCRIPT_DIR.parent
sys.path.insert(0, str(_MODULE_DIR.parent))

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


# =============================================================================
# EMBEDDING CONFIGURATION - SINGLE SOURCE OF TRUTH
# =============================================================================

EMBEDDING_MODEL = "google/gemini-embedding-001"  # text-embedding-004
EMBEDDING_MODEL_NAME = "gemini"  # Short name for files

EMBEDDING_VARIANTS = [
    {"name": "gemini_256", "dim": 256},
    {"name": "gemini_512", "dim": 512},
    {"name": "gemini_768", "dim": 768},
]

# Column name patterns - MUST match compute_multi_embeddings.py output
def get_query_col(dim: int) -> str:
    return f"query_embedding_{EMBEDDING_MODEL_NAME}_{dim}"

def get_task_col(dim: int) -> str:
    return f"task_description_embedding_{EMBEDDING_MODEL_NAME}_{dim}"


# =============================================================================
# VALIDATION FUNCTIONS
# =============================================================================

def validate_embedding_dimension(
    embedding: np.ndarray,
    expected_dim: int,
    source: str,
) -> None:
    """Validate embedding has exactly the expected dimension.

    Raises AssertionError if dimension doesn't match - NO padding/truncation.
    """
    actual_dim = embedding.shape[-1]
    if actual_dim != expected_dim:
        raise AssertionError(
            f"EMBEDDING DIMENSION MISMATCH in {source}:\n"
            f"  Expected: {expected_dim}\n"
            f"  Actual: {actual_dim}\n"
            f"  This indicates wrong embedding model. Check data source."
        )


def validate_all_embeddings_consistent(
    query_dim: int,
    task_dim: int,
    llm_dim: int,
    variant_dim: int,
) -> None:
    """Validate all embeddings have the same dimension for a variant."""
    if not (query_dim == task_dim == llm_dim == variant_dim):
        raise AssertionError(
            f"EMBEDDING DIMENSIONS INCONSISTENT:\n"
            f"  Query: {query_dim}\n"
            f"  Task: {task_dim}\n"
            f"  LLM: {llm_dim}\n"
            f"  Expected: {variant_dim}\n"
            f"  All must match exactly. No mixing allowed."
        )


# =============================================================================
# DATA PREPARATION
# =============================================================================

def prepare_data_file(
    source_parquet: Path,
    output_parquet: Path,
    dim: int,
    router_csv: Path,
) -> None:
    """Create a data file with Gemini embeddings at specified dimension.

    VALIDATES that embeddings match expected dimension - no silent failures.
    """
    query_col = get_query_col(dim)
    task_col = get_task_col(dim)

    logger.info(f"Preparing {output_parquet.name}...")
    logger.info(f"  Embedding model: {EMBEDDING_MODEL}")
    logger.info(f"  Dimension: {dim}")
    logger.info(f"  Query column: {query_col}")
    logger.info(f"  Task column: {task_col}")

    # Load multi-embed parquet with only needed columns
    needed_cols = [
        "task_id", "query", "ground_truth", "metric", "task_description",
        query_col, task_col,
    ]

    try:
        df = pd.read_parquet(source_parquet, columns=needed_cols)
    except Exception as e:
        raise ValueError(
            f"Failed to load required columns from {source_parquet}:\n"
            f"  Missing column? {e}\n"
            f"  Run compute_multi_embeddings.py to generate Gemini embeddings."
        ) from e

    # VALIDATE: Check embedding dimensions on first row
    import ast
    sample_query = df[query_col].iloc[0]
    sample_task = df[task_col].iloc[0]

    if isinstance(sample_query, str):
        query_emb = np.array(ast.literal_eval(sample_query))
    else:
        query_emb = np.array(sample_query)

    if isinstance(sample_task, str):
        task_emb = np.array(ast.literal_eval(sample_task))
    else:
        task_emb = np.array(sample_task)

    validate_embedding_dimension(query_emb, dim, f"query_embedding (row 0)")
    validate_embedding_dimension(task_emb, dim, f"task_embedding (row 0)")
    logger.info(f"  Validated: embeddings are {dim}-dim as expected")

    # Rename embedding columns to standard names
    df = df.rename(columns={
        query_col: "query_embedding",
        task_col: "task_description_embedding",
    })

    # Load LLM performance data from router_csv
    logger.info("  Merging LLM performance data...")
    router_df = pd.read_csv(
        router_csv,
        usecols=lambda c: c not in ["query_embedding", "task_description_embedding"],
    )

    # Get columns from router_csv that aren't in our data
    extra_cols = [c for c in router_df.columns if c not in df.columns]
    logger.info(f"  Adding columns: {extra_cols}")

    # Merge on query + task_id
    df = df.merge(
        router_df[["query", "task_id"] + extra_cols],
        on=["query", "task_id"],
        how="inner",
    )

    logger.info(f"  Final shape: {df.shape}")
    df.to_parquet(output_parquet, index=False)
    logger.info(f"  Saved to {output_parquet}")


# =============================================================================
# LLM EMBEDDING GENERATION
# =============================================================================

def generate_llm_embeddings(
    llm_descriptions_path: Path,
    output_path: Path,
    embedding_dim: int,
) -> None:
    """Generate LLM embeddings using Gemini embedding API.

    Uses text-embedding-004 with Matryoshka output_dimensionality.
    VALIDATES output dimension matches request.
    """
    logger.info(f"Generating {embedding_dim}-dim LLM embeddings...")
    logger.info(f"  Model: {EMBEDDING_MODEL}")

    try:
        import google.generativeai as genai
    except ImportError:
        raise ImportError(
            "google-generativeai not installed.\n"
            "Run: pip install google-generativeai"
        )

    # Check for API key
    api_key = os.environ.get("GOOGLE_API_KEY")
    if not api_key:
        # Try to load from ~/.ember/config.yaml
        ember_config_path = Path.home() / ".ember" / "config.yaml"
        if ember_config_path.exists():
            with open(ember_config_path) as f:
                ember_config = yaml.safe_load(f) or {}
            google_section = ember_config.get("google", {})
            keys = google_section.get("google_keys", [])
            if keys:
                api_key = keys[0]
                logger.info("  Using API key from ~/.ember/config.yaml")

    if not api_key:
        raise ValueError(
            "GOOGLE_API_KEY not found.\n"
            "Set environment variable or add to ~/.ember/config.yaml:\n"
            "  google:\n"
            "    google_keys:\n"
            "      - your-api-key-here"
        )

    genai.configure(api_key=api_key)

    # Load LLM descriptions
    with open(llm_descriptions_path) as f:
        llm_data = json.load(f)

    llm_names = list(llm_data.keys())
    texts = [llm_data[name]["feature"] for name in llm_names]

    logger.info(f"  Embedding {len(texts)} LLM descriptions...")

    embeddings = []
    for i, text in enumerate(texts):
        result = genai.embed_content(
            model="models/text-embedding-004",
            content=text,
            output_dimensionality=embedding_dim,
        )
        emb = result["embedding"]

        # VALIDATE dimension
        validate_embedding_dimension(
            np.array(emb), embedding_dim, f"LLM '{llm_names[i]}'"
        )

        embeddings.append(emb)
        logger.info(f"    [{i+1}/{len(texts)}] {llm_names[i][:30]}...")

    embeddings_array = np.array(embeddings, dtype=np.float32)

    # Final validation
    validate_embedding_dimension(embeddings_array, embedding_dim, "LLM embeddings array")
    logger.info(f"  Final shape: {embeddings_array.shape}")

    with open(output_path, "wb") as f:
        pickle.dump(embeddings_array, f)
    logger.info(f"  Saved to {output_path}")


# =============================================================================
# CONFIG GENERATION
# =============================================================================

def create_config_file(
    template_path: Path,
    output_path: Path,
    variant_name: str,
    embedding_dim: int,
    data_path: str,
    llm_embedding_path: str,
) -> None:
    """Create a config file for a specific embedding variant.

    Includes explicit documentation of embedding source.
    """
    logger.info(f"Creating config for {variant_name}...")

    with open(template_path) as f:
        config = yaml.safe_load(f)

    # Update paths
    config["saved_router_data_path"] = data_path
    config["llm_embedding_path"] = llm_embedding_path
    config["model_path"] = f"model_path/v8_{variant_name}_best.pth"

    # Scale hidden dimension appropriately
    if embedding_dim == 256:
        config["embedding_dim"] = 256
    elif embedding_dim == 512:
        config["embedding_dim"] = 384
    elif embedding_dim == 768:
        config["embedding_dim"] = 512

    # Write with clear header
    with open(output_path, "w") as f:
        f.write(f"# v8 {variant_name}\n")
        f.write(f"#\n")
        f.write(f"# EMBEDDING SOURCE: {EMBEDDING_MODEL}\n")
        f.write(f"# EMBEDDING DIMENSION: {embedding_dim}\n")
        f.write(f"#\n")
        f.write(f"# All embeddings (query, task, LLM) use the same model and dimension.\n")
        f.write(f"# NO mixing of embedding models. NO dimension padding.\n")
        f.write(f"#\n")
        yaml.dump(config, f, default_flow_style=False, sort_keys=False)

    logger.info(f"  Saved to {output_path}")


# =============================================================================
# MAIN
# =============================================================================

def main():
    """Prepare all data and configs for Gemini embedding sweep."""
    source_parquet = _MODULE_DIR / "data" / "router_data_multi_embed.parquet"
    router_csv = _MODULE_DIR / "data" / "router_data.csv"
    llm_descriptions = _MODULE_DIR / "configs" / "LLM_Descriptions.json"
    template_config = _MODULE_DIR / "registry" / "v7_gat4_needle_weighted.yaml"

    if not source_parquet.exists():
        logger.error(f"Source parquet not found: {source_parquet}")
        logger.error("Run data_processing/compute_multi_embeddings.py first.")
        sys.exit(1)

    logger.info("=" * 70)
    logger.info("GEMINI EMBEDDING SWEEP PREPARATION")
    logger.info("=" * 70)
    logger.info(f"Embedding model: {EMBEDDING_MODEL}")
    logger.info(f"Variants: {[v['name'] for v in EMBEDDING_VARIANTS]}")
    logger.info("")

    # Step 1: Prepare data files
    logger.info("--- Step 1: Preparing data files ---")
    for variant in EMBEDDING_VARIANTS:
        output_path = _MODULE_DIR / "data" / f"router_data_{variant['name']}.parquet"
        if output_path.exists():
            logger.info(f"  {output_path.name} already exists, skipping")
            continue
        prepare_data_file(
            source_parquet=source_parquet,
            output_parquet=output_path,
            dim=variant["dim"],
            router_csv=router_csv,
        )

    # Step 2: Generate LLM embeddings
    logger.info("\n--- Step 2: Generating LLM embeddings ---")
    for variant in EMBEDDING_VARIANTS:
        output_path = _MODULE_DIR / "configs" / f"llm_embedding_{variant['name']}.pkl"
        if output_path.exists():
            # Validate existing file
            with open(output_path, "rb") as f:
                existing = pickle.load(f)
            validate_embedding_dimension(
                existing, variant["dim"], f"existing {output_path.name}"
            )
            logger.info(f"  {output_path.name} exists and validated ({existing.shape})")
            continue

        generate_llm_embeddings(
            llm_descriptions_path=llm_descriptions,
            output_path=output_path,
            embedding_dim=variant["dim"],
        )

    # Step 3: Create config files
    logger.info("\n--- Step 3: Creating config files ---")
    for variant in EMBEDDING_VARIANTS:
        config_path = _MODULE_DIR / "registry" / f"v8_{variant['name']}.yaml"
        create_config_file(
            template_path=template_config,
            output_path=config_path,
            variant_name=variant["name"],
            embedding_dim=variant["dim"],
            data_path=f"data/router_data_{variant['name']}.parquet",
            llm_embedding_path=f"configs/llm_embedding_{variant['name']}.pkl",
        )

    # Step 4: Final validation
    logger.info("\n--- Step 4: Final validation ---")
    for variant in EMBEDDING_VARIANTS:
        data_path = _MODULE_DIR / "data" / f"router_data_{variant['name']}.parquet"
        llm_path = _MODULE_DIR / "configs" / f"llm_embedding_{variant['name']}.pkl"

        # Check data file
        df = pd.read_parquet(data_path, columns=["query_embedding"])
        import ast
        sample = df["query_embedding"].iloc[0]
        if isinstance(sample, str):
            query_emb = np.array(ast.literal_eval(sample))
        else:
            query_emb = np.array(sample)

        # Check LLM embeddings
        with open(llm_path, "rb") as f:
            llm_emb = pickle.load(f)

        # Validate consistency
        validate_all_embeddings_consistent(
            query_dim=query_emb.shape[0],
            task_dim=variant["dim"],  # Already validated during prep
            llm_dim=llm_emb.shape[1],
            variant_dim=variant["dim"],
        )

        logger.info(f"  {variant['name']}: VALIDATED (all {variant['dim']}-dim)")

    logger.info("\n" + "=" * 70)
    logger.info("PREPARATION COMPLETE - ALL VALIDATIONS PASSED")
    logger.info("=" * 70)
    logger.info("\nCreated files:")
    for variant in EMBEDDING_VARIANTS:
        logger.info(f"  - data/router_data_{variant['name']}.parquet")
        logger.info(f"  - configs/llm_embedding_{variant['name']}.pkl")
        logger.info(f"  - registry/v8_{variant['name']}.yaml")

    logger.info("\nTo launch the sweep:")
    logger.info("  python scripts/launch_gemini_sweep.py")


if __name__ == "__main__":
    main()
