import asyncio
from typing import List, Dict, Optional
from dataclasses import dataclass
import numpy as np
from sentence_transformers import SentenceTransformer, CrossEncoder
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
import redis
import hashlib
import json

@dataclass
class RAGConfig:
    embedding_model: str = "all-mpnet-base-v2"
    rerank_model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"
    chunk_size: int = 512
    chunk_overlap: int = 128
    top_k: int = 5
    rerank_top_k: int = 3
    cache_ttl: int = 3600

class ProductionRAG:
    """Production-ready RAG system with caching and reranking"""
    
    def __init__(self, config: RAGConfig):
        self.config = config
        self.embedder = SentenceTransformer(config.embedding_model)
        self.reranker = CrossEncoder(config.rerank_model)
        self.vector_store = QdrantClient(host="localhost", port=6333)
        self.cache = redis.Redis(host='localhost', port=6379)
        self.collection_name = "rag_documents"
        self._init_collection()
    
    def _init_collection(self):
        """Initialize vector store collection"""
        self.vector_store.recreate_collection(
            collection_name=self.collection_name,
            vectors_config=VectorParams(
                size=self.embedder.get_sentence_embedding_dimension(),
                distance=Distance.COSINE
            )
        )
    
    async def ingest_documents(self, documents: List[Dict]) -> Dict:
        """Ingest and index documents with smart chunking"""
        from smart_chunker import SmartChunker
        chunker = SmartChunker(self.config.chunk_size, self.config.chunk_overlap)
        
        all_chunks = []
        for doc in documents:
            # Smart chunking
            chunks = chunker.chunk_by_structure(doc['content'])
            
            for i, chunk in enumerate(chunks):
                # Generate embedding
                embedding = self.embedder.encode(chunk)
                
                # Create point for vector store
                point = PointStruct(
                    id=hashlib.md5(f"{doc['id']}_{i}".encode()).hexdigest(),
                    vector=embedding.tolist(),
                    payload={
                        "content": chunk,
                        "doc_id": doc['id'],
                        "chunk_index": i,
                        "metadata": doc.get('metadata', {})
                    }
                )
                all_chunks.append(point)
        
        # Batch upsert to vector store
        self.vector_store.upsert(
            collection_name=self.collection_name,
            points=all_chunks
        )
        
        return {
            "documents_processed": len(documents),
            "chunks_created": len(all_chunks)
        }
    
    async def retrieve(self, query: str, use_cache: bool = True) -> List[Dict]:
        """Retrieve relevant documents with caching and reranking"""
        # Check cache
        if use_cache:
            cache_key = f"rag:{hashlib.md5(query.encode()).hexdigest()}"
            cached = self.cache.get(cache_key)
            if cached:
                return json.loads(cached)
        
        # Generate query embedding
        query_embedding = self.embedder.encode(query)
        
        # Vector search (get more than needed for reranking)
        results = self.vector_store.search(
            collection_name=self.collection_name,
            query_vector=query_embedding.tolist(),
            limit=self.config.top_k * 3
        )
        
        # Extract documents
        docs = [{
            'content': hit.payload['content'],
            'score': hit.score,
            'metadata': hit.payload.get('metadata', {})
        } for hit in results]
        
        # Rerank documents
        if len(docs) > 0:
            pairs = [[query, doc['content']] for doc in docs]
            rerank_scores = self.reranker.predict(pairs)
            
            # Sort by rerank scores and take top-k
            ranked_docs = sorted(
                zip(docs, rerank_scores),
                key=lambda x: x[1],
                reverse=True
            )[:self.config.rerank_top_k]
            
            final_docs = [
                {**doc, 'rerank_score': float(score)}
                for doc, score in ranked_docs
            ]
        else:
            final_docs = []
        
        # Cache results
        if use_cache and final_docs:
            self.cache.setex(
                cache_key,
                self.config.cache_ttl,
                json.dumps(final_docs)
            )
        
        return final_docs
    
    async def query(self, question: str, llm_client) -> Dict:
        """Complete RAG query pipeline"""
        # Retrieve relevant documents
        retrieved_docs = await self.retrieve(question)
        
        if not retrieved_docs:
            return {
                "answer": "No relevant documents found.",
                "sources": [],
                "confidence": 0.0
            }
        
        # Create context
        context = "\n\n".join([
            f"[Source {i+1}]\n{doc['content']}"
            for i, doc in enumerate(retrieved_docs)
        ])
        
        # Generate answer
        prompt = f"""Answer based on the context provided.

Context:
{context}

Question: {question}

Answer:"""
        
        answer = await llm_client.generate(prompt)
        
        # Calculate confidence
        avg_score = np.mean([d.get('rerank_score', 0) for d in retrieved_docs])
        
        return {
            "answer": answer,
            "sources": retrieved_docs,
            "confidence": min(avg_score, 1.0)
        }