"""
Pinecone Retriever Client for NEXUS OS v2.1
Uses llama-text-embed-v2-index with nexus-repos namespace.
Also integrates Pinecone Assistant "pineosman2" for chat-based retrieval.

API key: loaded from env PINECONE_API_KEY
"""
import os
from typing import List, Dict, Optional, Any
from dataclasses import dataclass

try:
    from pinecone import Pinecone
    from pinecone_plugins.assistant.models.chat import Message
    PINECONE_AVAILABLE = True
except ImportError:
    PINECONE_AVAILABLE = False


@dataclass
class RetrievalResult:
    text: str
    score: float
    metadata: Dict[str, Any]
    source: str  # "vector" or "assistant"


class PineconeRetriever:
    """
    Dual-mode Pinecone retriever:
    1. Vector search: llama-text-embed-v2-index (dense embeddings)
    2. Assistant chat: pineosman2 (conversational RAG)
    """

    INDEX_NAME = "llama-text-embed-v2-index"
    NAMESPACE = "nexus-repos"
    ASSISTANT_NAME = "pineosman2"

    def __init__(self, api_key: Optional[str] = None, top_k: int = 10):
        if not PINECONE_AVAILABLE:
            raise ImportError("pinecone SDK not installed. Run: pip install pinecone")
        self.api_key = api_key or os.environ.get("PINECONE_API_KEY")
        if not self.api_key:
            raise ValueError("PINECONE_API_KEY required (env var or arg)")
        self.pc = Pinecone(api_key=self.api_key)
        self.index = self.pc.Index(self.INDEX_NAME)
        self.assistant = self.pc.assistant.Assistant(assistant_name=self.ASSISTANT_NAME)
        self.top_k = top_k

    def vector_search(self, query: str, top_k: Optional[int] = None) -> List[RetrievalResult]:
        """Dense vector search over nexus-repos namespace."""
        k = top_k or self.top_k
        results = self.index.search(
            namespace=self.NAMESPACE,
            vector={"input": query},
            top_k=k,
            include_metadata=True,
        )
        return [
            RetrievalResult(
                text=hit.get("metadata", {}).get("text", ""),
                score=hit.get("score", 0.0),
                metadata=hit.get("metadata", {}),
                source="vector",
            )
            for hit in results.get("matches", [])
        ]

    def assistant_chat(self, query: str) -> str:
        """Chat with Pinecone Assistant for conversational retrieval."""
        msg = Message(content=query)
        resp = self.assistant.chat(messages=[msg])
        return resp.get("message", {}).get("content", "")

    def assistant_chat_stream(self, query: str):
        """Streaming chat with Pinecone Assistant."""
        msg = Message(content=query)
        chunks = self.assistant.chat(messages=[msg], stream=True)
        for chunk in chunks:
            if chunk:
                yield chunk

    def hybrid_retrieve(self, query: str) -> Dict[str, Any]:
        """
        Hybrid retrieval: vector results + assistant summary.
        Returns structured evidence for TWAVE mu_ret calculation.
        """
        vector_results = self.vector_search(query)
        assistant_answer = self.assistant_chat(query)

        return {
            "query": query,
            "vector_results": [
                {"text": r.text, "score": r.score, "metadata": r.metadata}
                for r in vector_results
            ],
            "assistant_summary": assistant_answer,
            "top_score": vector_results[0].score if vector_results else 0.0,
            "avg_score": sum(r.score for r in vector_results) / len(vector_results) if vector_results else 0.0,
        }

    def get_evidence_for_ckplug(self, query: str) -> List[Dict[str, Any]]:
        """
        Format retrieval evidence for CK-PLUG token-level coupling.
        Returns list of evidence chunks with relevance scores.
        """
        results = self.vector_search(query)
        return [
            {
                "text": r.text,
                "relevance": r.score,
                "type": r.metadata.get("type", "unknown"),
                "owner": r.metadata.get("owner", "unknown"),
            }
            for r in results
        ]


# Mock retriever for offline testing
class MockPineconeRetriever:
    """Offline mock of PineconeRetriever for development/testing."""

    def __init__(self, top_k: int = 5):
        self.top_k = top_k
        self._mock_data = [
            {"text": "NEXUS OS is a hybrid inference operating system with thermodynamic control.", "score": 0.95, "type": "repo", "owner": "specimba"},
            {"text": "TWAVE uses Landau-Ginzburg free energy to detect hallucination bifurcations.", "score": 0.88, "type": "repo", "owner": "specimba"},
            {"text": "CK-PLUG enables token-level confidence gain for retrieval coupling.", "score": 0.82, "type": "paper", "owner": "CAS"},
            {"text": "Bose-Einstein condensate analogy provides stable reasoning at T≈0.8Tc.", "score": 0.79, "type": "research", "owner": "specimba"},
            {"text": "QWAVE allocates inference budget across local and cloud tiers.", "score": 0.76, "type": "repo", "owner": "specimba"},
        ]

    def vector_search(self, query: str, top_k: Optional[int] = None) -> List[RetrievalResult]:
        k = top_k or self.top_k
        matches = []
        for item in self._mock_data:
            score = item["score"] * (0.5 + 0.5 * (len(set(query.lower().split()) & set(item["text"].lower().split())) / max(1, len(query.split()))))
            matches.append(RetrievalResult(text=item["text"], score=score, metadata=item, source="vector"))
        matches.sort(key=lambda x: x.score, reverse=True)
        return matches[:k]

    def hybrid_retrieve(self, query: str) -> Dict[str, Any]:
        vector_results = self.vector_search(query)
        return {
            "query": query,
            "vector_results": [{"text": r.text, "score": r.score, "metadata": r.metadata} for r in vector_results],
            "assistant_summary": f"Mock summary for: {query}",
            "top_score": vector_results[0].score if vector_results else 0.0,
            "avg_score": sum(r.score for r in vector_results) / len(vector_results) if vector_results else 0.0,
        }

    def get_evidence_for_ckplug(self, query: str) -> List[Dict[str, Any]]:
        results = self.vector_search(query)
        return [{"text": r.text, "relevance": r.score, "type": r.metadata.get("type"), "owner": r.metadata.get("owner")} for r in results]