#!/usr/bin/env python3
"""Generate Qwen3 instruction-conditioned embeddings for training data.

Two-phase design to minimize memory pressure on 32GB Apple Silicon:
  Phase 1: Extract unique texts to JSON (separate script or inline)
  Phase 2: Load model and embed from JSON (this script)

Usage:
    # v2 (default):
    uv run python packages/embergraphrouter/scripts/generate_v2_embeddings.py

    # v3 (Gemini 3.1 Pro):
    uv run python packages/embergraphrouter/scripts/generate_v2_embeddings.py \
        --llm-descriptions configs/LLM_Descriptions_v3.json \
        --output-llm configs/llm_embedding_qwen3_v3_256_8llm.pkl

The script produces:
    - configs/query_embeddings_qwen3_instruct_router_v1_256.pkl
    - configs/task_embeddings_qwen3_instruct_router_v1_256.pkl
    - configs/llm_embedding_qwen3_v{N}_256_8llm.pkl
    - data/router_data_v{N}_8llm_qwen3.parquet (written in a separate step)
"""

import argparse
import gc
import json
import os
import pickle
import sys
import time
from pathlib import Path

import numpy as np
import torch
from sentence_transformers import SentenceTransformer

os.environ["KMP_DUPLICATE_LIB_OK"] = "True"

# Router-specific instruction (must match training exactly).
# v1: Original 5-category instruction (+5.36pp over retrieval default).
# v2: Adds science/logic domains, difficulty anchoring, answer format.
ROUTER_INSTRUCTION_V1 = (
    "Classify this query by task type (reasoning, mathematics, coding, "
    "factual recall, creative writing), cognitive difficulty level, and "
    "the language model capabilities required to answer it accurately."
)

ROUTER_INSTRUCTION = (
    "Classify this query by: task domain (reasoning, mathematics, coding, "
    "science, factual recall, creative writing, logic), difficulty level "
    "from trivial to expert, expected answer format (multiple choice, "
    "free-form, code, numeric, yes/no), and the specialized AI capabilities "
    "required to answer correctly."
)

# Instruction for LLM description embeddings. Mirrors the query instruction's
# semantic axes (task domains, difficulty, capabilities) but from the model
# perspective, so both embedding types attend to the same dimensions.
LLM_INSTRUCTION = (
    "Identify this AI model's strengths and weaknesses across task domains "
    "(reasoning, mathematics, coding, science, factual recall, creative "
    "writing, logic), the difficulty levels it handles reliably, and its "
    "specialized capabilities for different answer formats."
)

TARGET_DIM = 256
BATCH_SIZE = 2  # Very conservative for MPS on 32GB Apple Silicon
MODEL_NAME = "Qwen/Qwen3-Embedding-8B"

# Resolve paths relative to embergraphrouter root
EGR_ROOT = Path(__file__).resolve().parent.parent
INPUT_JSON = EGR_ROOT / "data" / ".v2_embed_input.json"
LLM_DESC_PATH_DEFAULT = EGR_ROOT / "configs" / "LLM_Descriptions_v2.json"
LLM_OUTPUT_DEFAULT = EGR_ROOT / "configs" / "llm_embedding_qwen3_v2_256_8llm.pkl"
CONFIGS_DIR = EGR_ROOT / "configs"
CHECKPOINT_PATH = EGR_ROOT / "data" / ".embed_checkpoint_v2.pkl"


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate Qwen3 instruction-conditioned embeddings."
    )
    parser.add_argument(
        "--llm-descriptions",
        type=Path,
        default=None,
        help=(
            "Path to LLM descriptions JSON (relative to embergraphrouter root "
            "or absolute). Defaults to configs/LLM_Descriptions_v2.json."
        ),
    )
    parser.add_argument(
        "--output-llm",
        type=Path,
        default=None,
        help=(
            "Output path for LLM embedding pickle (relative to "
            "embergraphrouter root or absolute). Defaults to "
            "configs/llm_embedding_qwen3_v2_256_8llm.pkl."
        ),
    )
    return parser.parse_args()


def truncate_and_normalize(embeddings: np.ndarray, dim: int) -> np.ndarray:
    """Apply Matryoshka truncation and re-normalize."""
    if dim < embeddings.shape[1]:
        embeddings = embeddings[:, :dim]
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    return embeddings / np.maximum(norms, 1e-12)


def embed_batched(
    model: SentenceTransformer,
    texts: list[str],
    batch_size: int,
    label: str,
) -> np.ndarray:
    """Batch-encode texts with progress reporting."""
    embeddings = []
    start = time.time()
    total = len(texts)
    for i in range(0, total, batch_size):
        batch = texts[i : i + batch_size]
        batch_embs = model.encode(
            batch, convert_to_numpy=True, normalize_embeddings=True
        )
        embeddings.append(batch_embs)
        done = i + len(batch)
        if done <= 10 or done % (batch_size * 50) == 0 or done == total:
            elapsed = time.time() - start
            rate = done / elapsed if elapsed > 0 else 0
            eta = (total - done) / rate if rate > 0 else 0
            print(
                f"  [{label}] {done}/{total} "
                f"({rate:.1f} texts/s, ETA {eta:.0f}s)",
                flush=True,
            )
    return np.vstack(embeddings)


def main() -> None:
    args = _parse_args()

    # Resolve LLM description and output paths
    llm_desc_path = args.llm_descriptions or LLM_DESC_PATH_DEFAULT
    if not llm_desc_path.is_absolute():
        llm_desc_path = EGR_ROOT / llm_desc_path
    llm_output_path = args.output_llm or LLM_OUTPUT_DEFAULT
    if not llm_output_path.is_absolute():
        llm_output_path = EGR_ROOT / llm_output_path

    print("=" * 70, flush=True)
    print("Qwen3 Instruction-Conditioned Embedding Generation", flush=True)
    print(f"  LLM descriptions: {llm_desc_path.name}", flush=True)
    print(f"  LLM output:       {llm_output_path.name}", flush=True)
    print("=" * 70, flush=True)

    # Load lightweight JSON input (no pandas DataFrame needed)
    print(f"\nLoading input from {INPUT_JSON}...", flush=True)
    with open(INPUT_JSON) as f:
        data = json.load(f)
    queries = data["queries"]
    task_ids = data["task_ids"]
    task_descriptions = data["task_descriptions"]
    llm_names = data["llm_names"]
    print(f"  Queries: {len(queries):,}", flush=True)
    print(f"  Tasks: {len(task_ids)}", flush=True)
    print(f"  LLMs: {len(llm_names)}", flush=True)

    # Check for checkpoint (resume from query embeddings)
    if CHECKPOINT_PATH.exists():
        print(f"\nResuming from checkpoint: {CHECKPOINT_PATH}", flush=True)
        with open(CHECKPOINT_PATH, "rb") as f:
            checkpoint = pickle.load(f)
        query_map = checkpoint["query_map"]
        print(f"  Loaded {len(query_map)} query embeddings from checkpoint", flush=True)
        # Still need model for task + LLM embeddings
        need_queries = False
    else:
        query_map = None
        need_queries = True

    # Detect device
    if torch.backends.mps.is_available():
        device = "mps"
    elif torch.cuda.is_available():
        device = "cuda"
    else:
        device = "cpu"
    print(f"\nDevice: {device}", flush=True)

    # Load model
    print(f"Loading {MODEL_NAME}...", flush=True)
    start = time.time()
    model = SentenceTransformer(
        MODEL_NAME,
        device=device,
        trust_remote_code=True,
        model_kwargs={"torch_dtype": torch.float16},
    )
    print(f"Model loaded in {time.time() - start:.1f}s", flush=True)
    print(f"Native dimension: {model.get_sentence_embedding_dimension()}", flush=True)

    if need_queries:
        # Generate query embeddings with router instruction
        print(f"\nGenerating query embeddings ({len(queries):,} queries)...", flush=True)
        formatted = [
            f"Instruct: {ROUTER_INSTRUCTION}\nQuery: {q}" for q in queries
        ]
        query_embs_full = embed_batched(model, formatted, BATCH_SIZE, "Queries")
        query_emb_256 = truncate_and_normalize(query_embs_full, TARGET_DIM)
        del query_embs_full
        gc.collect()
        print(f"Query embeddings: {query_emb_256.shape}", flush=True)

        # Save checkpoint
        query_map = {q: emb for q, emb in zip(queries, query_emb_256)}
        with open(CHECKPOINT_PATH, "wb") as f:
            pickle.dump({"query_map": query_map, "timestamp": time.time()}, f)
        print(f"Checkpoint saved: {CHECKPOINT_PATH.name}", flush=True)
    else:
        query_emb_256 = np.array(
            [query_map[q] for q in queries], dtype=np.float32
        )

    # Generate task description embeddings (document mode, no instruction)
    print(f"\nGenerating task description embeddings ({len(task_descriptions)} tasks)...", flush=True)
    task_embs_full = embed_batched(model, task_descriptions, BATCH_SIZE, "Tasks")
    task_emb_256 = truncate_and_normalize(task_embs_full, TARGET_DIM)
    del task_embs_full
    print(f"Task embeddings: {task_emb_256.shape}", flush=True)

    # Generate LLM description embeddings (instruction-conditioned).
    # Uses LLM_INSTRUCTION to project descriptions into the same semantic
    # subspace as queries (task domains, difficulty, capabilities).
    print(f"\nGenerating LLM description embeddings from {llm_desc_path.name}...", flush=True)
    with open(llm_desc_path) as f:
        llm_desc_data = json.load(f)
    llm_texts = [
        f"Instruct: {LLM_INSTRUCTION}\nQuery: "
        + llm_desc_data.get(name, {}).get("feature", name)
        for name in llm_names
    ]
    llm_embs_full = embed_batched(model, llm_texts, BATCH_SIZE, "LLMs")
    llm_emb_256 = truncate_and_normalize(llm_embs_full, TARGET_DIM)
    del llm_embs_full
    print(f"LLM embeddings: {llm_emb_256.shape}", flush=True)
    print(f"LLM names: {llm_names}", flush=True)

    # Free model memory before saving
    del model
    gc.collect()
    if torch.backends.mps.is_available():
        torch.mps.empty_cache()

    # Save pickle files
    suffix = "qwen3_instruct_router_v2_256"

    query_pkl = CONFIGS_DIR / f"query_embeddings_{suffix}.pkl"
    with open(query_pkl, "wb") as f:
        pickle.dump(query_map, f)
    print(f"\nSaved: {query_pkl.name} ({len(query_map)} entries)", flush=True)

    task_map = {tid: task_emb_256[i] for i, tid in enumerate(task_ids)}
    task_pkl = CONFIGS_DIR / f"task_embeddings_{suffix}.pkl"
    with open(task_pkl, "wb") as f:
        pickle.dump(task_map, f)
    print(f"Saved: {task_pkl.name} ({len(task_map)} entries)", flush=True)

    llm_pkl = llm_output_path
    with open(llm_pkl, "wb") as f:
        pickle.dump(llm_emb_256, f)
    print(f"Saved: {llm_pkl.name} ({llm_emb_256.shape})", flush=True)

    # Cleanup checkpoint
    if CHECKPOINT_PATH.exists():
        CHECKPOINT_PATH.unlink()

    print("\n" + "=" * 70, flush=True)
    print("EMBEDDING GENERATION COMPLETE", flush=True)
    print("=" * 70, flush=True)
    print(f"\nNow run the parquet writer:", flush=True)
    print(f"  uv run python packages/embergraphrouter/scripts/write_v2_parquet.py", flush=True)


if __name__ == "__main__":
    main()
