import faiss
import json
import asyncio
import aiohttp
import numpy as np
from typing import List, Dict, Any, Tuple, Optional
import os
import datetime
from tqdm import tqdm
import logging
import time

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class EmbeddingClient:
    """Centralized embedding client with rate limiting and error handling. Currently support OpenAI API."""
    
    def __init__(self, base_url: str, model: str = "Alibaba-NLP/gte-Qwen2-7B-instruct"):
        self.base_url = base_url
        self.model = model
        self.url = f"{base_url}/embeddings"
    
    async def get_embedding(self, session: aiohttp.ClientSession, text: str) -> List[float]:
        """Get embedding with improved error handling."""
        payload = {
            "model": self.model,
            "input": text,
        }
        
        try:
            async with session.post(self.url, json=payload) as response:
                if response.status == 429:  # Rate limit
                    retry_after = int(response.headers.get('Retry-After', 60))
                    logger.warning(f"Rate limited, waiting {retry_after}s")
                    await asyncio.sleep(retry_after)
                    return await self.get_embedding(session, text)
                
                if response.status != 200:
                    error_text = await response.text()
                    raise Exception(f"HTTP {response.status}: {error_text}")
                
                result = await response.json()
                return result["data"][0]["embedding"]
                
        except asyncio.TimeoutError:
            raise Exception("Request timed out")
        except Exception as e:
            raise Exception(f"Embedding request failed: {str(e)}")

class FAISSIndexBuilder:
    """FAISS index builder with improved configuration and monitoring."""
    
    def __init__(self, embedding_client: EmbeddingClient):
        self.embedding_client = embedding_client
        self.stats = {
            'total_processed': 0,
            'successful': 0,
            'failed': 0,
            'retries': 0,
            'start_time': None
        }
    
    async def build_index(
        self,
        corpus: List[Dict[str, Any]],
        hnsw_m: int = 64,
        ef_construction: int = 400,
        ef_search: int = 200,
        save_path: Optional[str] = None,
        batch_size: int = 50,
        batch_delay: float = 1.0,
        max_retries: int = 3
    ) -> faiss.Index:
        """Build FAISS index with improved monitoring and error handling."""
        
        logger.info(f"🚀 Building FAISS index for {len(corpus)} documents")
        logger.info(f"📦 Batch size: {batch_size}, delay: {batch_delay}s")
        
        self.stats['start_time'] = time.time()
        
        # Configure session with optimized settings
        connector = aiohttp.TCPConnector(
            limit=batch_size * 2,
            limit_per_host=batch_size,
            ttl_dns_cache=300,
            use_dns_cache=True,
            keepalive_timeout=60
        )
        
        timeout = aiohttp.ClientTimeout(
            total=3600,
            connect=60,
            sock_read=300
        )
        
        async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
            embeddings = await self._process_documents_in_batches(
                session, corpus, batch_size, batch_delay, max_retries
            )
            
            if self.stats['successful'] == 0:
                raise Exception("No embeddings were successfully generated!")
            
            return self._build_faiss_index(
                embeddings, hnsw_m, ef_construction, ef_search, save_path, corpus
            )
    
    async def _process_documents_in_batches(
        self,
        session: aiohttp.ClientSession,
        corpus: List[Dict[str, Any]],
        batch_size: int,
        batch_delay: float,
        max_retries: int
    ) -> List[Optional[List[float]]]:
        """Process documents in batches with progress tracking."""
        
        embeddings = [None] * len(corpus)
        total_batches = (len(corpus) + batch_size - 1) // batch_size
        
        with tqdm(total=len(corpus), desc="Processing documents", unit="docs") as pbar:
            for batch_idx in range(0, len(corpus), batch_size):
                batch_docs = corpus[batch_idx:batch_idx + batch_size]
                current_batch_num = (batch_idx // batch_size) + 1
                
                # Process batch
                await self._process_batch(session, batch_docs, embeddings, max_retries)
                
                pbar.update(len(batch_docs))
                
                # Progress logging
                elapsed = time.time() - self.stats['start_time']
                docs_per_sec = self.stats['total_processed'] / elapsed if elapsed > 0 else 0
                
                logger.info(
                    f"📊 Batch {current_batch_num}/{total_batches}: "
                    f"{self.stats['successful']} success, {self.stats['failed']} failed, "
                    f"{docs_per_sec:.2f} docs/sec"
                )
                
                # Wait before next batch
                if batch_idx + batch_size < len(corpus) and batch_delay > 0:
                    await asyncio.sleep(batch_delay)
        
        return embeddings
    
    async def _process_batch(
        self,
        session: aiohttp.ClientSession,
        batch_docs: List[Dict[str, Any]],
        embeddings: List[Optional[List[float]]],
        max_retries: int
    ):
        """Process a single batch of documents."""
        
        tasks = [
            self._process_document_with_retry(session, doc, max_retries)
            for doc in batch_docs
        ]
        
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        for result in results:
            self.stats['total_processed'] += 1
            
            if isinstance(result, Exception):
                logger.error(f"❌ Batch exception: {result}")
                self.stats['failed'] += 1
                continue
            
            chunk_id, embedding = result
            embeddings[chunk_id] = embedding
            
            if embedding is not None:
                self.stats['successful'] += 1
            else:
                self.stats['failed'] += 1
    
    async def _process_document_with_retry(
        self,
        session: aiohttp.ClientSession,
        doc: Dict[str, Any],
        max_retries: int
    ) -> Tuple[int, Optional[List[float]]]:
        """Process document with exponential backoff retry."""
        
        text = f"(Part of {doc['title']}) {doc['text']}"
        
        for attempt in range(max_retries):
            try:
                embedding = await self.embedding_client.get_embedding(session, text)
                return (doc["chunk_id"], embedding)
            
            except Exception as e:
                if attempt < max_retries - 1:
                    retry_delay = 2 ** attempt  # Exponential backoff
                    self.stats['retries'] += 1
                    logger.warning(
                        f"🔄 Retry {attempt + 1}/{max_retries} for doc {doc['chunk_id']}: {e}"
                    )
                    await asyncio.sleep(retry_delay)
                else:
                    logger.error(f"❌ Final failure for doc {doc['chunk_id']}: {e}")
                    return (doc["chunk_id"], None)
    
    def _build_faiss_index(
        self,
        embeddings: List[Optional[List[float]]],
        hnsw_m: int,
        ef_construction: int,
        ef_search: int,
        save_path: Optional[str],
        corpus: List[Dict[str, Any]]
    ) -> faiss.Index:
        """Build the actual FAISS index."""
        
        logger.info("🔄 Converting embeddings to numpy array...")
        
        # Filter valid embeddings
        valid_embeddings = []
        valid_indices = []
        for i, emb in enumerate(embeddings):
            if emb is not None:
                valid_embeddings.append(emb)
                valid_indices.append(i)
        
        embeddings_array = np.array(valid_embeddings, dtype=np.float32)
        logger.info(f"📊 Valid embeddings: {len(valid_embeddings)}/{len(embeddings)}")
        
        # Create and configure index
        dimension = embeddings_array.shape[1]
        index = faiss.IndexHNSWFlat(dimension, hnsw_m)
        index.hnsw.efConstruction = ef_construction
        index.hnsw.efSearch = ef_search
        
        # Add vectors in batches
        logger.info("📊 Adding vectors to FAISS index...")
        faiss_batch_size = 1000
        for i in tqdm(range(0, len(embeddings_array), faiss_batch_size), 
                     desc="Adding to index", unit="batch"):
            batch = embeddings_array[i:i + faiss_batch_size]
            index.add(batch)
        
        if save_path:
            self._save_index_and_metadata(
                index, save_path, corpus, hnsw_m, ef_construction, 
                ef_search, valid_indices, dimension
            )
        
        return index
    
    def _save_index_and_metadata(self, index, save_path, corpus, hnsw_m, 
                                ef_construction, ef_search, valid_indices, dimension):
        """Save index and metadata to files."""
        
        logger.info("💾 Saving index and metadata...")
        
        # Prepare paths
        os.makedirs(os.path.dirname(save_path) if os.path.dirname(save_path) else '.', 
                   exist_ok=True)
        
        if not save_path.endswith('.faiss'):
            faiss_path = save_path + '.faiss'
            base_path = save_path
        else:
            faiss_path = save_path
            base_path = save_path[:-6]
        
        # Save index
        faiss.write_index(index, faiss_path)
        
        # Save metadata
        elapsed_time = time.time() - self.stats['start_time']
        metadata = {
            'created_at': datetime.datetime.now().isoformat(),
            'processing_time_seconds': elapsed_time,
            'corpus_size': len(corpus),
            'successful_embeddings': self.stats['successful'],
            'failed_embeddings': self.stats['failed'],
            'retry_count': self.stats['retries'],
            'dimension': dimension,
            'hnsw_m': hnsw_m,
            'ef_construction': ef_construction,
            'ef_search': ef_search,
            'total_vectors': index.ntotal,
            'valid_indices': valid_indices,
            'stats': self.stats
        }
        
        metadata_path = f"{base_path}_metadata.json"
        with open(metadata_path, 'w') as f:
            json.dump(metadata, f, indent=2)
        
        logger.info(f"✅ FAISS index saved to: {faiss_path}")
        logger.info(f"📋 Metadata saved to: {metadata_path}")
        logger.info(f"📊 Processing completed in {elapsed_time:.2f}s")

# # Usage example
# async def main():
#     """Main execution function."""
    
#     # Configuration
#     corpus_path = "datasets/qasper/qasper_corpus_chunked.jsonl"
#     index_save_path = "datasets/qasper/qasper_corpus_index.faiss"
#     base_url = "http://localhost:30000/v1"
    
#     # Load corpus
#     logger.info(f"📖 Loading corpus from {corpus_path}")
#     with open(corpus_path, "r") as f:
#         corpus = [json.loads(line) for line in f]
    
#     # Build index
#     embedding_client = EmbeddingClient(base_url)
#     builder = FAISSIndexBuilder(embedding_client)
    
#     index = await builder.build_index(
#         corpus=corpus,
#         hnsw_m=64,
#         ef_construction=400,
#         ef_search=200,
#         save_path=index_save_path,
#         batch_size=50,
#         batch_delay=2.0,
#         max_retries=3
#     )
    
#     logger.info("🎉 Index building completed!")

# if __name__ == "__main__":
#     asyncio.run(main())