"""
Pinecone Vector Database Client for CognitiveRAG.

Handles storage and retrieval of embeddings using Pinecone.
"""

from __future__ import annotations

import logging
from typing import Dict, List, Optional, Any

from pinecone import Pinecone, ServerlessSpec

# handle both package import and direct import 
try:
    from . import config
except ImportError:
    import config

# configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class PineconeClient:
    """Client for interacting with Pinecone vector database."""

    def __init__(self, api_key: str = None, index_name: str = None):
        """
        Initialize the Pinecone client.

        args:
            api_key: pinecone API key. (check config for api key or ask me)
            index_name: Name of the pinecone index. (currently rag-hybrid-index-v1) 
        """
        self.api_key = api_key or config.PINECONE_API_KEY
        self.index_name = index_name or config.PINECONE_INDEX_NAME
        
        # initialize pinecone
        self.pc = Pinecone(api_key=self.api_key)
        self._index = None

    def _ensure_index_exists(self) -> None:
        """create the index if it doesn't exist."""
        existing_indexes = [idx.name for idx in self.pc.list_indexes()]
        
        if self.index_name not in existing_indexes:
            logger.info(f"Creating new Pinecone index: {self.index_name}")
            self.pc.create_index(
                name=self.index_name,
                dimension=config.EMBEDDING_DIMENSION,
                metric=config.SIMILARITY_METRIC,    #cosine for now but can be changed in config
                spec=ServerlessSpec(
                    cloud=config.PINECONE_CLOUD,    #set to aws (for aws track maybe?)
                    region=config.PINECONE_REGION    #set to us-east-1 
                )
            )
            logger.info(f"Index '{self.index_name}' created successfully")
        else:
            logger.info(f"Index '{self.index_name}' already exists")

    def get_index(self):
        """get or create the pinecone index."""
        if self._index is None:
            self._ensure_index_exists()
            self._index = self.pc.Index(self.index_name)
        return self._index

    def store_embeddings(
        self, 
        chunks: List[Dict[str, Any]], 
        batch_size: int = 100,
        namespace: str = ""
    ) -> Dict[str, Any]:
        """
        store chunk embeddings in pinecone.use this if adding a few documents, otherwise use batch_upsert_from_file (can crash laptop)

        args:
            chunks: list of chunk dictionaries containing:
                - chunk_id: unique identifier for the chunk
                - embedding: vector embedding (list of floats)
                - text: original text content
                - doc_id: document identifier
                - section: section name (optional)
                - page: page number (optional)
            batch_size: number of vectors to upsert per batch
            namespace: pinecone namespace for organizing vectors

        returns:
            dictionary with storage results:
                - stored_count: number of vectors stored
                - errors: list of any errors encountered
        """
        index = self.get_index()
        stored_count = 0
        errors = []

        # Process in batches for efficiency
        for i in range(0, len(chunks), batch_size):
            batch = chunks[i:i + batch_size]
            vectors = []

            for chunk in batch:
                try:
                    # validate required fields
                    if "chunk_id" not in chunk:
                        errors.append(f"Missing chunk_id in chunk at index {i}")
                        continue
                    if "embedding" not in chunk:
                        errors.append(f"Missing embedding for chunk {chunk.get('chunk_id', 'unknown')}")
                        continue

                    # build metadata (pinecone stores this alongside the vector)
                    # Handle both 'text' and 'content' field names
                    text_content = chunk.get("text") or chunk.get("content", "")
                    page_val = chunk.get("page")
                    
                    # Extract related_chunks for relational retrieval
                    related = chunk.get("related_chunks", {})
                    
                    metadata = {
                        "text": text_content,
                        "doc_id": chunk.get("doc_id", ""),
                        "section": chunk.get("section", ""),
                        "page": page_val if page_val is not None else "",
                        # Store prev/next for relational chunking
                        "prev_chunk": related.get("prev") or "",
                        "next_chunk": related.get("next") or "",
                    }

                    vectors.append({
                        "id": chunk["chunk_id"],
                        "values": chunk["embedding"],
                        "metadata": metadata
                    })

                except Exception as e:
                    errors.append(f"Error processing chunk {chunk.get('chunk_id', 'unknown')}: {str(e)}")

            # upsert batch to pinecone 
            if vectors:
                try:
                    index.upsert(vectors=vectors, namespace=namespace)
                    stored_count += len(vectors)
                    logger.info(f"Stored batch of {len(vectors)} vectors (total: {stored_count})")
                except Exception as e:
                    errors.append(f"Error upserting batch: {str(e)}")

        return {
            "stored_count": stored_count,
            "total_chunks": len(chunks),
            "errors": errors if errors else None
        }

    def retrieve_similar(
        self, 
        query_embedding: List[float], 
        top_k: int = None,
        namespace: str = "",
        filter: Dict = None
    ) -> Dict[str, Any]:
        """
        retrieve the most similar chunks for a query embedding.

        args:
            query_embedding: the query vector (from ajwaad's embedding model)
            top_k: number of results to return. can be edited in config
            namespace: pinecone namespace to search in
            filter: optional metadata filter for the query from earlier

        Returns:
            dictionary with retrieved context:
                - query_embedding_length: length of the query vector
                - retrieved_chunks: list of matching chunks with:
                    - chunk_id: The chunk identifier
                    - text: The original text content
                    - score: Similarity score
                    - doc_id: Document identifier
                    - section: Section name
                    - page: Page number 
        """
        index = self.get_index()
        top_k = top_k or config.TOP_K

        try:
            # query Pinecone for similar vectors
            results = index.query(
                vector=query_embedding,
                top_k=top_k,
                include_metadata=True,
                namespace=namespace,
                filter=filter
            )

            # format results
            retrieved_chunks = []
            for match in results.get("matches", []):
                metadata = match.get("metadata", {})
                retrieved_chunks.append({
                    "chunk_id": match["id"],
                    "text": metadata.get("text", ""),
                    "score": match["score"],
                    "doc_id": metadata.get("doc_id", ""),
                    "section": metadata.get("section", ""),
                    "page": metadata.get("page"),
                    # Include relational chunk IDs
                    "prev_chunk": metadata.get("prev_chunk", ""),
                    "next_chunk": metadata.get("next_chunk", ""),
                })

            return {
                "query_embedding_length": len(query_embedding),
                "retrieved_chunks": retrieved_chunks
            }

        except Exception as e:
            logger.error(f"Error retrieving similar chunks: {str(e)}")
            return {
                "query_embedding_length": len(query_embedding),
                "retrieved_chunks": [],
                "error": str(e)
            }

    def delete_vectors(
        self, 
        ids: List[str] = None, 
        delete_all: bool = False,
        namespace: str = ""
    ) -> Dict[str, Any]:
        """
        delete vectors from the index (i think it can be done manually on the pinecone website).

        args:
            ids: list of vector ids to delete
            delete_all: if true, delete all vectors in the namespace
            namespace: pinecone namespace

        returns:
            dictionary with deletion results
        """
        index = self.get_index()
        
        try:
            if delete_all:
                index.delete(delete_all=True, namespace=namespace)
                return {"deleted": "all", "namespace": namespace}
            elif ids:
                index.delete(ids=ids, namespace=namespace)
                return {"deleted_count": len(ids), "ids": ids}
            else:
                return {"error": "Must provide ids or set delete_all=True"}
        except Exception as e:
            logger.error(f"Error deleting vectors: {str(e)}")
            return {"error": str(e)}

    def batch_upsert_from_file(
        self,
        file_path: str,
        batch_size: int = 100,
        namespace: str = "",
        show_progress: bool = True
    ) -> Dict[str, Any]:
        """
        batch upsert embeddings from a JSONL file. (for mass uploads to pinecone server)
        
        use this for uploading the initial corpus. each line in the file
        should be a JSON object with chunk_id, embedding, text, doc_id, etc.

        Args:
        file_path: path to JSONL file containing embedded chunks
        batch_size: Number of vectors to upsert per batch (default: 100)
        namespace: Pinecone namespace
        show_progress: Whether to print progress updates
        Returns:
        Dictionary with upload results:
            stored_count: Total vectors stored
            total_lines: Total lines processed
            errors: Any errors encountered
            batches_processed: Number of batches
        """
        import json
        
        index = self.get_index()
        stored_count = 0
        total_lines = 0
        errors = []
        batches_processed = 0
        
        current_batch = []
        
        with open(file_path, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, 1):
                total_lines += 1
                
                try:
                    chunk = json.loads(line.strip())
                    
                    # validate required fields
                    if "chunk_id" not in chunk or "embedding" not in chunk:
                        errors.append(f"Line {line_num}: Missing required fields")
                        continue
                    
                    # build vector for Pinecone
                    # Handle both 'text' and 'content' field names
                    text_content = chunk.get("text") or chunk.get("content", "")
                    # Pinecone doesn't allow null values in metadata, use empty string
                    page_val = chunk.get("page")
                    
                    # Extract related_chunks for relational retrieval
                    related = chunk.get("related_chunks", {})
                    
                    vector = {
                        "id": chunk["chunk_id"],
                        "values": chunk["embedding"],
                        "metadata": {
                            "text": text_content,
                            "doc_id": chunk.get("doc_id", ""),
                            "section": chunk.get("section", ""),
                            "page": page_val if page_val is not None else "",
                            # Store prev/next for relational chunking
                            "prev_chunk": related.get("prev") or "",
                            "next_chunk": related.get("next") or "",
                        }
                    }
                    current_batch.append(vector)
                    
                    # upsert when batch is full
                    if len(current_batch) >= batch_size:
                        try:
                            index.upsert(vectors=current_batch, namespace=namespace)
                            stored_count += len(current_batch)
                            batches_processed += 1
                            
                            if show_progress:
                                logger.info(f"Batch {batches_processed}: Stored {stored_count} vectors")
                            
                            current_batch = []
                        except Exception as e:
                            errors.append(f"Batch {batches_processed + 1} error: {str(e)}")
                            current_batch = []
                            
                except json.JSONDecodeError as e:
                    errors.append(f"Line {line_num}: Invalid JSON - {str(e)}")
                except Exception as e:
                    errors.append(f"Line {line_num}: {str(e)}")
        
        # upsert remaining vectors
        if current_batch:
            try:
                index.upsert(vectors=current_batch, namespace=namespace)
                stored_count += len(current_batch)
                batches_processed += 1
                
                if show_progress:
                    logger.info(f"Final batch: Stored {stored_count} total vectors")
            except Exception as e:
                errors.append(f"Final batch error: {str(e)}")
        
        return {
            "stored_count": stored_count,
            "total_lines": total_lines,
            "batches_processed": batches_processed,
            "errors": errors if errors else None
        }

    def get_index_stats(self) -> Dict[str, Any]:
        """Get statistics about the index.""" #i included this for debugging but it could be useful... (use after batch upsert to confirm it works)
        index = self.get_index()
        try:
            stats = index.describe_index_stats()
            return {
                "dimension": stats.get("dimension"),
                "total_vector_count": stats.get("total_vector_count"),
                "namespaces": stats.get("namespaces", {})
            }
        except Exception as e:
            logger.error(f"Error getting index stats: {str(e)}")
            return {"error": str(e)}


class SparseClient:
    """Client for Pinecone sparse vectors using pinecone-sparse-english-v0."""

    def __init__(self, api_key: str = None, index_name: str = None):
        self.api_key = api_key or config.PINECONE_API_KEY
        self.index_name = index_name or config.SPARSE_INDEX_NAME
        self.pc = Pinecone(api_key=self.api_key)
        self._index = None

    def get_index(self):
        """Get the sparse index."""
        if self._index is None:
            self._index = self.pc.Index(self.index_name)
        return self._index

    def embed_sparse(self, texts: List[str]) -> List[Dict]:
        """Generate sparse embeddings using Pinecone inference."""
        embeddings = self.pc.inference.embed(
            model=config.SPARSE_EMBEDDING_MODEL,
            inputs=texts,
            parameters={"input_type": "passage"}
        )
        return [
            {"indices": e.to_dict()["sparse_indices"], "values": e.to_dict()["sparse_values"]}
            for e in embeddings.data
        ]

    def embed_query_sparse(self, query: str) -> Dict:
        """Generate sparse embedding for a query."""
        embeddings = self.pc.inference.embed(
            model=config.SPARSE_EMBEDDING_MODEL,
            inputs=[query],
            parameters={"input_type": "query"}
        )
        d = embeddings.data[0].to_dict()
        return {"indices": d["sparse_indices"], "values": d["sparse_values"]}

    def store_sparse(
        self,
        chunks: List[Dict[str, Any]],
        batch_size: int = 96,
        namespace: str = ""
    ) -> Dict[str, Any]:
        """Store sparse embeddings in the sparse index."""
        index = self.get_index()
        stored_count = 0
        errors = []

        for i in range(0, len(chunks), batch_size):
            batch_chunks = chunks[i:i + batch_size]
            
            # Embed this batch only (avoid 413 errors)
            texts = [c.get("text") or c.get("content", "") for c in batch_chunks]
            try:
                sparse_embeddings = self.embed_sparse(texts)
            except Exception as e:
                errors.append(f"Embed error batch {i//batch_size}: {str(e)}")
                continue
            
            vectors = []
            for chunk, sparse in zip(batch_chunks, sparse_embeddings):
                if "chunk_id" not in chunk:
                    errors.append(f"Missing chunk_id")
                    continue

                text_content = chunk.get("text") or chunk.get("content", "")
                page_val = chunk.get("page")
                related = chunk.get("related_chunks", {})

                vectors.append({
                    "id": chunk["chunk_id"],
                    "sparse_values": sparse,
                    "metadata": {
                        "text": text_content,
                        "doc_id": chunk.get("doc_id", ""),
                        "section": chunk.get("section", ""),
                        "page": page_val if page_val is not None else "",
                        "prev_chunk": related.get("prev") or "",
                        "next_chunk": related.get("next") or "",
                    }
                })

            if vectors:
                try:
                    index.upsert(vectors=vectors, namespace=namespace)
                    stored_count += len(vectors)
                    logger.info(f"Stored {stored_count} sparse vectors")
                except Exception as e:
                    errors.append(f"Error upserting batch: {str(e)}")

        return {
            "stored_count": stored_count,
            "total_chunks": len(chunks),
            "errors": errors if errors else None
        }

    def retrieve_sparse(
        self,
        query: str,
        top_k: int = None,
        namespace: str = ""
    ) -> Dict[str, Any]:
        """Retrieve using sparse embeddings."""
        index = self.get_index()
        top_k = top_k or config.TOP_K

        try:
            sparse_query = self.embed_query_sparse(query)
            results = index.query(
                sparse_vector=sparse_query,
                top_k=top_k,
                include_metadata=True,
                namespace=namespace
            )

            retrieved_chunks = []
            for match in results.get("matches", []):
                metadata = match.get("metadata", {})
                retrieved_chunks.append({
                    "chunk_id": match["id"],
                    "text": metadata.get("text", ""),
                    "score": match["score"],
                    "doc_id": metadata.get("doc_id", ""),
                    "section": metadata.get("section", ""),
                    "page": metadata.get("page"),
                    "prev_chunk": metadata.get("prev_chunk", ""),
                    "next_chunk": metadata.get("next_chunk", ""),
                })

            return {"retrieved_chunks": retrieved_chunks}

        except Exception as e:
            logger.error(f"Error retrieving sparse: {str(e)}")
            return {"retrieved_chunks": [], "error": str(e)}

    def get_index_stats(self) -> Dict[str, Any]:
        """Get statistics about the sparse index."""
        index = self.get_index()
        try:
            stats = index.describe_index_stats()
            return {
                "total_vector_count": stats.get("total_vector_count"),
                "namespaces": stats.get("namespaces", {})
            }
        except Exception as e:
            logger.error(f"Error getting sparse index stats: {str(e)}")
            return {"error": str(e)}


# convenience functions for direct usage
_client = None
_sparse_client = None

def get_client() -> PineconeClient:
    """Get or create a single Pinecone Client instance."""
    global _client
    if _client is None:
        _client = PineconeClient()
    return _client


def get_sparse_client() -> SparseClient:
    """Get or create a single SparseClient instance."""
    global _sparse_client
    if _sparse_client is None:
        _sparse_client = SparseClient()
    return _sparse_client


def store_embeddings(chunks: List[Dict[str, Any]], **kwargs) -> Dict[str, Any]:
    """Store embeddings using the default client."""
    return get_client().store_embeddings(chunks, **kwargs)


def retrieve_similar(query_embedding: List[float], **kwargs) -> Dict[str, Any]:
    """Retrieve similar chunks using the default client."""
    return get_client().retrieve_similar(query_embedding, **kwargs)


def batch_upsert_from_file(file_path: str, **kwargs) -> Dict[str, Any]:
    """Batch upsert embeddings from a JSONL file using the default client."""
    return get_client().batch_upsert_from_file(file_path, **kwargs)