#!/usr/bin/env python3
"""
Build a production-ready semantic search index from documents.
"""

import argparse
import json
import pickle
from pathlib import Path
from typing import List
import numpy as np
from pygeai import Client


def load_documents(file_path: Path) -> List[str]:
    """Load documents from various formats."""
    if file_path.suffix == ".json":
        with open(file_path) as f:
            data = json.load(f)
            return [doc["text"] for doc in data]
    elif file_path.suffix == ".txt":
        with open(file_path) as f:
            return [line.strip() for line in f if line.strip()]
    else:
        raise ValueError(f"Unsupported format: {file_path.suffix}")


def create_embeddings(client: Client, texts: List[str], model: str, batch_size: int = 100):
    """Generate embeddings in batches."""
    all_embeddings = []
    
    print(f"Processing {len(texts)} documents...")
    
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        
        response = client.embeddings.create(input=batch, model=model)
        embeddings = [emb.embedding for emb in response.data]
        all_embeddings.extend(embeddings)
        
        print(f"  {min(i + batch_size, len(texts))}/{len(texts)} documents processed")
    
    return all_embeddings


def save_index(documents: List[str], embeddings: List[List[float]], output_path: Path):
    """Save search index to disk."""
    index_data = {
        "documents": documents,
        "embeddings": np.array(embeddings).astype('float32'),
        "dimension": len(embeddings[0]),
        "count": len(documents)
    }
    
    with open(output_path, "wb") as f:
        pickle.dump(index_data, f)
    
    print(f"\n✅ Index saved to {output_path}")
    print(f"   Documents: {len(documents)}")
    print(f"   Dimensions: {len(embeddings[0])}")
    print(f"   Size: {output_path.stat().st_size / 1024 / 1024:.2f} MB")


def main():
    parser = argparse.ArgumentParser(description="Build semantic search index")
    parser.add_argument("input", type=Path, help="Input file (.json or .txt)")
    parser.add_argument("output", type=Path, help="Output index file (.pkl)")
    parser.add_argument("--model", default="text-embedding-3-small", help="Embedding model")
    parser.add_argument("--batch-size", type=int, default=100, help="Batch size")
    
    args = parser.parse_args()
    
    print("🔨 Building Search Index\n")
    print("=" * 60)
    
    # Load documents
    print(f"\n📄 Loading documents from {args.input}")
    documents = load_documents(args.input)
    print(f"   Loaded {len(documents)} documents")
    
    # Create embeddings
    print(f"\n🧠 Generating embeddings with {args.model}")
    client = Client()
    embeddings = create_embeddings(client, documents, args.model, args.batch_size)
    
    # Save index
    print(f"\n💾 Saving index to {args.output}")
    save_index(documents, embeddings, args.output)


if __name__ == "__main__":
    main()
