#!/usr/bin/env python3
"""Generate embeddings for sweep using raw API calls (no ember dependency).

This script generates query/task embeddings for the multi-embedding sweep
using direct API calls to Google and OpenAI embedding endpoints.

Usage:
    python scripts/generate_sweep_embeddings.py [--sources gemini_256,openai_large_768]
"""

import argparse
import hashlib
import json
import logging
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from threading import Lock
from typing import Dict, List, Optional, Tuple

import numpy as np
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
_cache_lock = Lock()

# Embedding configurations for sweep
EMBEDDING_SOURCES = {
    "gemini_256": {"model": "text-embedding-004", "dimension": 256, "provider": "google"},
    "gemini_512": {"model": "text-embedding-004", "dimension": 512, "provider": "google"},
    "gemini_768": {"model": "text-embedding-004", "dimension": 768, "provider": "google"},
    "openai_large_256": {"model": "text-embedding-3-large", "dimension": 256, "provider": "openai"},
    "openai_large_512": {"model": "text-embedding-3-large", "dimension": 512, "provider": "openai"},
    "openai_large_768": {"model": "text-embedding-3-large", "dimension": 768, "provider": "openai"},
    "openai_large_1536": {"model": "text-embedding-3-large", "dimension": 1536, "provider": "openai"},
    "openai_large_3072": {"model": "text-embedding-3-large", "dimension": 3072, "provider": "openai"},
}


def load_api_keys() -> Dict[str, str]:
    """Load API keys from ~/.ember/config.yaml."""
    config_path = Path.home() / ".ember" / "config.yaml"
    if not config_path.exists():
        raise FileNotFoundError(f"Config not found: {config_path}")

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

    keys = {}
    providers = config.get("providers", {})
    if "google" in providers:
        keys["google"] = providers["google"].get("api_key")
    if "openai" in providers:
        keys["openai"] = providers["openai"].get("api_key")

    return keys


def _text_hash(text: str) -> str:
    """Generate unique hash for text."""
    return hashlib.sha256(text.encode()).hexdigest()[:16]


def embed_google(texts: List[str], model: str, dimension: int, api_key: str) -> List[List[float]]:
    """Embed texts using Google's embedding API."""
    import requests

    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:embedContent"
    embeddings = []

    for text in texts:
        payload = {
            "model": f"models/{model}",
            "content": {"parts": [{"text": text}]},
            "outputDimensionality": dimension,
        }
        headers = {"Content-Type": "application/json"}

        for attempt in range(5):
            try:
                resp = requests.post(
                    f"{url}?key={api_key}",
                    headers=headers,
                    json=payload,
                    timeout=30,
                )
                if resp.status_code == 429:
                    time.sleep(2 ** attempt)
                    continue
                resp.raise_for_status()
                data = resp.json()
                embeddings.append(data["embedding"]["values"])
                break
            except Exception as e:
                if attempt == 4:
                    logger.error(f"Failed to embed: {e}")
                    embeddings.append([0.0] * dimension)
                time.sleep(2 ** attempt)

    return embeddings


def embed_google_batch(texts: List[str], model: str, dimension: int, api_key: str, batch_size: int = 100) -> List[List[float]]:
    """Embed texts using Google's batch embedding API."""
    import requests

    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:batchEmbedContents"
    all_embeddings = []

    for i in range(0, len(texts), batch_size):
        batch = texts[i:i+batch_size]
        requests_payload = [
            {
                "model": f"models/{model}",
                "content": {"parts": [{"text": text}]},
                "outputDimensionality": dimension,
            }
            for text in batch
        ]

        for attempt in range(5):
            try:
                resp = requests.post(
                    f"{url}?key={api_key}",
                    headers={"Content-Type": "application/json"},
                    json={"requests": requests_payload},
                    timeout=60,
                )
                if resp.status_code == 429:
                    time.sleep(2 ** (attempt + 1))
                    continue
                resp.raise_for_status()
                data = resp.json()
                for emb in data.get("embeddings", []):
                    all_embeddings.append(emb["values"])
                break
            except Exception as e:
                if attempt == 4:
                    logger.error(f"Batch failed: {e}")
                    for _ in batch:
                        all_embeddings.append([0.0] * dimension)
                time.sleep(2 ** attempt)

    return all_embeddings


def embed_openai_batch(texts: List[str], model: str, dimension: int, api_key: str, batch_size: int = 100) -> List[List[float]]:
    """Embed texts using OpenAI's batch embedding API."""
    import requests

    url = "https://api.openai.com/v1/embeddings"
    all_embeddings = []

    for i in range(0, len(texts), batch_size):
        batch = texts[i:i+batch_size]

        for attempt in range(5):
            try:
                resp = requests.post(
                    url,
                    headers={
                        "Authorization": f"Bearer {api_key}",
                        "Content-Type": "application/json",
                    },
                    json={
                        "model": model,
                        "input": batch,
                        "dimensions": dimension,
                    },
                    timeout=60,
                )
                if resp.status_code == 429:
                    time.sleep(2 ** (attempt + 1))
                    continue
                resp.raise_for_status()
                data = resp.json()
                # Sort by index to maintain order
                sorted_data = sorted(data["data"], key=lambda x: x["index"])
                for item in sorted_data:
                    all_embeddings.append(item["embedding"])
                break
            except Exception as e:
                if attempt == 4:
                    logger.error(f"Batch failed: {e}")
                    for _ in batch:
                        all_embeddings.append([0.0] * dimension)
                time.sleep(2 ** attempt)

    return all_embeddings


def load_cache(cache_path: Path) -> Dict:
    """Load embedding cache from disk."""
    if cache_path.exists():
        try:
            with open(cache_path) as f:
                return json.load(f)
        except Exception:
            return {}
    return {}


def save_cache(cache: Dict, cache_path: Path) -> None:
    """Save embedding cache atomically."""
    temp_path = cache_path.with_suffix(".tmp")
    with _cache_lock:
        with open(temp_path, "w") as f:
            json.dump(cache, f)
        temp_path.rename(cache_path)


def compute_embeddings_for_source(
    texts: List[str],
    source_name: str,
    config: Dict,
    api_keys: Dict[str, str],
    cache_dir: Path,
) -> Dict[str, List[float]]:
    """Compute embeddings for a single source."""
    cache_path = cache_dir / f"embed_cache_{source_name}.json"
    cache = load_cache(cache_path)

    # Find texts not in cache
    to_compute = []
    hashes = []
    for text in texts:
        h = _text_hash(text)
        if h not in cache:
            to_compute.append(text)
            hashes.append(h)

    if not to_compute:
        logger.info(f"  {source_name}: All {len(texts)} texts cached")
        return cache

    logger.info(f"  {source_name}: Computing {len(to_compute)}/{len(texts)} texts ({len(texts) - len(to_compute)} cached)")

    provider = config["provider"]
    model = config["model"]
    dimension = config["dimension"]
    api_key = api_keys.get(provider)

    if not api_key:
        raise ValueError(f"No API key for {provider}")

    # Compute embeddings
    if provider == "google":
        embeddings = embed_google_batch(to_compute, model, dimension, api_key, batch_size=100)
    else:  # openai
        embeddings = embed_openai_batch(to_compute, model, dimension, api_key, batch_size=100)

    # Update cache
    for h, emb in zip(hashes, embeddings):
        cache[h] = emb

    save_cache(cache, cache_path)
    return cache


def main():
    parser = argparse.ArgumentParser(description="Generate sweep embeddings")
    parser.add_argument("--sources", type=str, default="all",
                       help="Comma-separated sources or 'all'")
    parser.add_argument("--input", type=str, default=None,
                       help="Input CSV path")
    parser.add_argument("--output", type=str, default=None,
                       help="Output parquet path")
    args = parser.parse_args()

    # Load API keys
    api_keys = load_api_keys()
    logger.info(f"Loaded API keys: {list(api_keys.keys())}")

    # Parse sources
    if args.sources == "all":
        sources = list(EMBEDDING_SOURCES.keys())
    else:
        sources = [s.strip() for s in args.sources.split(",")]

    logger.info(f"Sources: {sources}")

    # Load data
    input_path = Path(args.input) if args.input else _MODULE_DIR / "data" / "unified_qa_data.csv"
    output_path = Path(args.output) if args.output else _MODULE_DIR / "data" / "router_data_multi_embed.parquet"
    cache_dir = _MODULE_DIR / "data" / ".embed_cache"
    cache_dir.mkdir(parents=True, exist_ok=True)

    logger.info(f"Loading data from {input_path}")
    df = pd.read_csv(input_path, low_memory=False)
    logger.info(f"Loaded {len(df)} rows")

    # Get unique texts
    all_queries = df["query"].drop_duplicates().tolist()
    all_task_descs = df["task_description"].drop_duplicates().tolist()
    all_texts = list(set(all_queries + all_task_descs))
    logger.info(f"Unique texts: {len(all_queries)} queries, {len(all_task_descs)} tasks, {len(all_texts)} total")

    # Compute embeddings for each source
    for source_name in sources:
        if source_name not in EMBEDDING_SOURCES:
            logger.warning(f"Unknown source: {source_name}")
            continue

        config = EMBEDDING_SOURCES[source_name]
        logger.info(f"\nProcessing {source_name} ({config['dimension']}-dim, {config['provider']})...")

        cache = compute_embeddings_for_source(all_texts, source_name, config, api_keys, cache_dir)

        # Add columns to dataframe
        query_col = f"query_embedding_{source_name}"
        task_col = f"task_description_embedding_{source_name}"

        df[query_col] = df["query"].apply(lambda q: str(cache.get(_text_hash(q), [])))
        df[task_col] = df["task_description"].apply(lambda t: str(cache.get(_text_hash(t), [])))

        logger.info(f"  Added columns: {query_col}, {task_col}")

    # Save output
    logger.info(f"\nSaving to {output_path}")
    df.to_parquet(output_path, index=False)

    csv_path = output_path.with_suffix(".csv")
    df.to_csv(csv_path, index=False)
    logger.info(f"Saved CSV to {csv_path}")

    print("\n" + "=" * 60)
    print("EMBEDDING GENERATION COMPLETE")
    print("=" * 60)
    print(f"Rows: {len(df):,}")
    print(f"Sources: {sources}")
    print(f"Output: {output_path}")


if __name__ == "__main__":
    main()
