"""
pinecone.py
-----------
Pinecone implementation of VectorBackend.
"""

from typing import Any, Dict, List, Optional, Sequence, Tuple
import numpy as np
import os

try:
    from pinecone import Pinecone, ServerlessSpec
except ImportError:
    Pinecone = None

from hieromem_core.backends.interface import VectorBackend


class PineconeBackend(VectorBackend):
    def __init__(
        self, api_key: str = None, environment: str = None, index_name: str = "hieromem"
    ):
        if not Pinecone:
            raise ImportError("pinecone-client is required for PineconeBackend")

        self.pc = Pinecone(api_key=api_key or os.getenv("PINECONE_API_KEY"))
        self.index_name = index_name

        # Check if index exists, else create (if serverless spec provided? for now assume exists or standard create)
        existing = [i.name for i in self.pc.list_indexes()]
        if index_name not in existing:
            # Basic creation for starter/serverless
            self.pc.create_index(
                name=index_name,
                dimension=768,
                metric="cosine",
                spec=ServerlessSpec(cloud="aws", region=environment or "us-east-1"),
            )

        self.index = self.pc.Index(index_name)

    def add(
        self,
        ids: Sequence[str],
        embeddings: np.ndarray,
        metadata: Optional[List[Dict[str, Any]]] = None,
    ) -> None:
        # Pinecone upsert expects (id, values, metadata)
        vectors = []
        for i, eid in enumerate(ids):
            meta = metadata[i] if metadata else {}
            vectors.append((eid, embeddings[i].tolist(), meta))

        # Batching is recommended for large upserts, doing simple impl for now
        batch_size = 100
        for i in range(0, len(vectors), batch_size):
            self.index.upsert(vectors[i:i + batch_size])

    def search(
        self,
        query_embedding: np.ndarray,
        k: int = 10,
        filter: Optional[Dict[str, Any]] = None,
    ) -> List[Tuple[str, float, Dict[str, Any]]]:
        res = self.index.query(
            vector=query_embedding.tolist(),
            top_k=k,
            include_metadata=True,
            filter=filter,
        )

        return [(match.id, match.score, match.metadata or {}) for match in res.matches]

    def get(
        self, ids: Sequence[str]
    ) -> List[Optional[Tuple[str, np.ndarray, Dict[str, Any]]]]:
        res = self.index.fetch(ids=list(ids))
        results = []
        for eid in ids:
            if eid in res.vectors:
                vec = res.vectors[eid]
                results.append((eid, np.array(vec.values), vec.metadata or {}))
            else:
                results.append(None)
        return results

    def delete(self, ids: Sequence[str]) -> None:
        self.index.delete(ids=list(ids))

    def count(self) -> int:
        return self.index.describe_index_stats().total_vector_count

    def clear(self) -> None:
        self.index.delete(delete_all=True)
