"""
Vector Service Client for interacting with vector database.

This module provides functionality for:
- Managing vector collections
- Storing and retrieving document embeddings (via the service)
- Handling semantic and agentic tools
- Chunking large documents
- Sanitizing metadata
"""

import json
from typing import Any, Dict, List, Optional, Union

import httpx
from guardianhub import get_logger

# Assuming 'settings' and 'get_logger' are available in the project environment
from guardianhub.config.settings import settings

logger = get_logger(__name__)

# Default collection names
DEFAULT_COLLECTION_DOCS = "document_templates"
BULLET_COLLECTION = "ace_context_bullets"


def _chunk_text(
        text: str,
        max_tokens: int = 500,
        approx_chars_per_token: int = 4
) -> List[str]:
    """Split text into chunks of approximately max_tokens length."""
    if not text:
        return []

    max_chars = max_tokens * approx_chars_per_token
    if len(text) <= max_chars:
        return [text]

    chunks = []
    start = 0

    while start < len(text):
        end = min(len(text), start + max_chars)
        slice_ = text[start:end]

        # Look for the last newline or space to cut cleanly
        if end < len(text):
            last_nl = slice_.rfind("\n")
            last_space = slice_.rfind(" ")
            cut = max(last_nl, last_space)
            if cut > 0:
                end = start + cut

        chunks.append(text[start:end].strip())
        start = end

    return [c for c in chunks if c]

def _sanitize_metadata_value(value: Any) -> Any:
    """Ensure metadata value is vector-DB safe."""
    if value is None or value == []:
        return None
    if isinstance(value, (str, int, float, bool)):
        return value
    try:
        # Complex types like dicts or lists should be JSON serialized
        return json.dumps(value)
    except Exception:
        return str(value)

def sanitize_metadata(metadata: Dict[str, Any]) -> Dict[str, Any]:
    """Clean metadata dictionary by dropping None values and normalizing types."""
    return {
        k: _sanitize_metadata_value(v)
        for k, v in metadata.items()
        if _sanitize_metadata_value(v) is not None
    }

class VectorClient:
    """Client for interacting with the vector database service."""

    def __init__(
            self,
            base_url: str = settings.endpoints.VECTOR_SERVICE_URL,
            collections: Union[str, List[str]] = None,
            http_timeout: float = 30.0,
            **collection_kwargs
    ):
        """Initialize the vector service client.

        Args:
            base_url: Base URL for the vector service
            collections: Name or list of names of document collections to use.
                       If None, uses DEFAULT_COLLECTION_DOCS
            http_timeout: Timeout for HTTP requests in seconds
            **collection_kwargs: Additional collection configuration
        """
        self.base_url = base_url.rstrip("/")

        # Handle collections parameter - convert single string to list for consistency
        if collections is None:
            self.collections = [DEFAULT_COLLECTION_DOCS]
        elif isinstance(collections, str):
            self.collections = [collections]
        else:
            self.collections = list(collections)

        self.collection_config = collection_kwargs
        self._client = httpx.AsyncClient(
            base_url=self.base_url,
            timeout=http_timeout
        )
        self.initialized = False

        """Initialize the vector service client."""
        self.base_url = base_url.rstrip("/")
        self.collections = collections
        self.collection_config = collection_kwargs
        self._client = httpx.AsyncClient(
            base_url=self.base_url,
            timeout=http_timeout
        )
        self.initialized = False

    async def initialize(self):
        """Initialize the client and check connection asynchronously"""
        try:
            await self._check_connection()
            # Ensure the default collection exists
            collection_names = [self.collections, BULLET_COLLECTION,"episodes","lessons"]
            logger.info(f"Initializing Collections.. {collection_names}")
            await self.ensure_collection_exists(
                collection_names,
                metadata={
                    "description": "Collection being managed by sutram-orchestrator",
                    "created_by": settings.service.name
                }
            )

            self.initialized = True
            return self.initialized
        except Exception as e:
            logger.error(f"Vector client initialization failed: {str(e)}")
            self.initialized = False
            return False

    async def ensure_collection_exists(self, collection_names: List[str] = None, **collection_kwargs) -> bool:
        """Ensure the specified collection exists, create it if it doesn't.

        Args:
            collection_name: Name of the collection to check/create
            **collection_kwargs: Additional arguments for collection creation

        Returns:
            bool: True if collection exists or was created, False otherwise
        """
        for collection_name in collection_names:
            try:
                # Try to create the collection - will succeed if it doesn't exist
                await self.ensure_collection(collection_name, **collection_kwargs)
            except httpx.HTTPStatusError as e:
                if e.response.status_code == 409:  # Collection already exists
                    logger.debug(f"Collection {collection_name} already exists")
                logger.error(f"Error ensuring collection {collection_name} exists: {str(e)}")


    async def ensure_collection(self, collection_name: str, **collection_kwargs):
        """Create a new collection with optional arguments."""
        response = await self._client.post(
            f"/collections/{collection_name}/ensure",
            json=collection_kwargs
        )
        response.raise_for_status()
        logger.info(response.json())
        # logger.info(f"Collection : {collection_name} is ensured to exist")

    async def _check_connection(self):
        """Check if the vector service is available."""
        try:
            response = await self._client.get("/health")
            response.raise_for_status()
            self.initialized = True
        except Exception as e:
            logger.error(f"Vector client health check failed: {str(e)}")
            self.initialized = False
            raise

    async def close(self):
        """Close the HTTP client."""
        await self._client.aclose()

    async def upsert_documents(
            self,
            ids: List[str],
            documents: List[str],
            metadatas: Optional[List[Dict[str, Any]]],
            collection: Optional[str] = DEFAULT_COLLECTION_DOCS,
            embeddings: Optional[List[List[float]]] = None,
    ) -> Any:
        """Upsert documents into the vector database in a single batch call."""
        if not ids or not documents or len(ids) != len(documents):
            raise ValueError("ids and documents must be non-empty and same length")

        if metadatas and len(metadatas) != len(ids):
            raise ValueError("metadatas must be same length as ids if provided")

        # Ensure metadata is safe for the vector store
        sanitized_metas = [
            sanitize_metadata(m) if m else {}
            for m in (metadatas or [{}] * len(ids))
        ]

        payload = {
            # Note: We do not pass embeddings here; the service will generate them.
            "documents": documents,
            "ids": ids,
            "metadatas": sanitized_metas
        }

        response = await self._client.post(
            f"/collections/{collection}/add",
            json=payload
        )
        response.raise_for_status()
        return response.json()

    async def upsert_document_from_text(
            self,
            document_content: str,
            doc_id: str,
            metadata: Dict[str, Any],
            collection: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Chunks a single document and upserts all chunks into the vector database.

        The ingestion activity (store_in_vector_db_activity) should call this method.
        It returns the ID of the first chunk, which acts as the 'vector_id' for the document.
        """
        collection_name = collection or self.collections

        if not document_content:
            logger.warning(f"Attempted to upsert empty content for document ID: {doc_id}")
            return {"status": "skipped", "message": "Empty document content"}

        # 1. Chunk the document
        chunks = _chunk_text(document_content)

        if not chunks:
             logger.warning(f"Chunking failed for document ID: {doc_id}")
             return {"status": "skipped", "message": "Chunking failed to produce content"}

        # 2. Prepare IDs and Metadata for upsert
        chunk_ids = [f"{doc_id}-{i}" for i in range(len(chunks))]

        # All chunks share the same base metadata, plus chunk-specific indices
        base_metadata = sanitize_metadata(metadata.copy())

        chunk_metadatas = []
        for i in range(len(chunks)):
            meta = base_metadata.copy()
            # Store the primary document ID in the metadata for easy filtering/retrieval
            meta["original_doc_id"] = doc_id
            meta["chunk_index"] = i
            meta["chunk_total"] = len(chunks)
            chunk_metadatas.append(meta)

        # 3. Upsert the documents (chunks)
        response = await self.upsert_documents(
            ids=chunk_ids,
            documents=chunks,
            metadatas=chunk_metadatas,
            collection=collection_name,
        )

        # Return the ID of the first chunk
        return {"id": chunk_ids[0], "status": "success", "chunk_count": len(chunks), "service_response": response}


    async def delete_document(self, doc_id: str, collection: str) -> None:
        """Delete a document (and all its chunks) from the vector database."""
        # For a complete document deletion, we must delete based on the original_doc_id in metadata
        response = await self._client.post(
            f"/collections/{collection}/delete",
            json={
                "where": {
                    "original_doc_id": doc_id
                }
            }
        )
        response.raise_for_status()
        logger.info(f"Deleted document chunks for original ID {doc_id} from {collection}")

    # Removed the previous embed_text method as embedding is handled by the service.

    async def query(
            self,
            query: str,
            collection: str,
            n_results: int = 5,
            where: Optional[Dict[str, Any]] = None
    ) -> List[Dict[str, Any]]:
        """Query the vector database using a text query."""
        if not self.initialized:
            await self.initialize()

        payload = {
            "query_texts": [query],
            "n_results": n_results
        }
        # Where clause must be sanitized if needed, but we rely on the service to interpret Chroma syntax
        if where:
            payload["where"] = where

        response = await self._client.post(
            f"/collections/{collection}/query",
            json=payload
        )
        response.raise_for_status()
        return response.json()

    async def delete_collection(self, name: str) -> None:
        """Delete a collection."""
        response = await self._client.delete(f"/collections/{name}")
        response.raise_for_status()
        logger.info(f"Deleted collection {name}")



    async def retrieve_relevant_context(
            self,
            query: str,
            collection: Optional[str] = None,
            n_results: int = 5,
            where: Optional[Dict[str, Any]] = None
    ) -> List[Dict[str, Any]]:
        """
        Retrieves context chunks relevant to the query, formatted for the ContextManager.

        Args:
            query: The search query string.
            collection: Optional collection name override.
            n_results: Number of results to retrieve.
            where: Metadata filter dictionary.

        Returns:
            List of dictionaries with keys: 'text', 'source', 'score', 'metadata'.
        """
        target_collection = collection or self.collections

        try:
            # Use the existing low-level query method
            # Note: The service returns: {'results': {'documents': [[...]], 'metadatas': [[...]], 'distances': [[...]]}}
            response_data = await self.query(
                query=query,
                collection=target_collection,
                n_results=n_results,
                where=where
            )

            # Extract the inner results dictionary (the service wraps it in 'results')
            results = response_data.get('results', {})

            if not results or not results.get('documents'):
                return []

            # Process the parallel arrays from Chroma
            documents = results['documents'][0]
            metadatas = results['metadatas'][0]
            distances = results['distances'][0]

            context_chunks = []

            for i, doc_text in enumerate(documents):
                # Convert distance to similarity score (1 - distance)
                # Assuming cosine distance where 0 is identical
                score = 1.0 - distances[i]
                meta = metadatas[i]

                context_chunks.append({
                    "text": doc_text,
                    "source": meta.get("source", f"Doc-{meta.get('original_doc_id', 'unknown')}"),
                    "score": score,
                    "metadata": meta
                })

            return context_chunks

        except Exception as e:
            logger.error(f"Failed to retrieve relevant context from vector DB: {str(e)}")
            return []

        # services/clients/vector_client.py (Add this method)

    async def query_context_bullets(
            self,
            query: str,
            template_id: str,
            limit: int = 5
    ) -> List[Dict[str, Any]]:
        """
        Retrieves specialized ContextBullet items for ACE Playbook assembly.
        Assumes bullets are stored in a dedicated collection or filtered by type.
        """
        try:
            # Ensure the client is initialized
            if not self.initialized:
                await self.initialize()

            # Prepare the query payload
            payload = {
                "query_texts": [query],
                "n_results": limit,
                "where": {"template_id": template_id} if template_id else {}
            }

            # Make the request directly to avoid any coroutine serialization issues
            response = await self._client.post(
                f"/collections/{BULLET_COLLECTION}/query",
                json=payload
            )

            # Handle the response
            response.raise_for_status()
            response_data = response.json()

            # Add debug logging to inspect the response structure
            logger.debug(f"Vector query response: {response_data}")

            # Handle empty or unexpected response structure
            if not response_data or 'documents' not in response_data or not response_data['documents']:
                logger.warning(f"No documents found in vector query response for template {template_id}")
                return []

            # Process the results
            documents = response_data.get('documents', [[]])[0]
            metadatas = response_data.get('metadatas', [[]])[0]
            distances = response_data.get('distances', [[]])[0]
            ids = response_data.get('ids', [[]])[0]  # Handle missing ids gracefully

            results = []
            for i, doc_text in enumerate(documents):
                if i >= len(metadatas) or i >= len(distances):
                    break

                score = 1.0 - distances[i]  # Convert distance to similarity score
                meta = metadatas[i] if i < len(metadatas) else {}

                result = {
                    "id": ids[i] if i < len(ids) else f"bullet-{i}",  # Changed from bullet_id to id to match expected format
                    "content": doc_text,
                    "score": score,
                    "metadata": meta
                }
                results.append(result)

            return results

        except Exception as e:
            logger.error(f"Error querying context bullets: {str(e)}", exc_info=True)
            return []

    async def get_recent_episodes(
            self,
            query: str,
            collection: str = "episodes",
            limit: int = 10,
            similarity_threshold: float = 0.7
    ) -> List[Dict[str, Any]]:
        """
        Retrieve recent episodes that are semantically similar to the query.
        Returns empty list if the collection doesn't exist or on error.
        """
        try:
            response = await self._client.post(
                f"/collections/{collection}/query",
                json={
                    "query_texts": [query],
                    "n_results": limit
                }
            )
            response.raise_for_status()

            results = response.json()
            episodes = []

            if not (results and
                    isinstance(results, dict) and
                    results.get("results", {}).get("documents") and
                    results["results"].get("ids")):
                return []

            documents = results["results"]["documents"][0]
            distances = results["results"].get("distances", [[1.0] * len(documents)])[0]
            ids = results["results"]["ids"][0]
            metadatas = (results["results"].get("metadatas", [[{}] * len(ids)])[0]
                         if results["results"].get("metadatas") else [{}] * len(ids))

            for i, (doc, distance, metadata) in enumerate(zip(documents, distances, metadatas)):
                try:
                    # Skip if distance exceeds threshold (lower distance = more similar)
                    if distance > (1.0 - similarity_threshold):
                        continue

                    episode = {
                        "id": ids[i],
                        "content": doc,
                        "score": float(1.0 - distance),  # Convert to similarity score (0-1)
                        "metadata": metadata,
                        "type": metadata.get("type", "episode"),
                        "timestamp": metadata.get("timestamp")
                    }
                    episodes.append(episode)

                except Exception as e:
                    logger.warning(f"Error processing episode {i}: {e}")
                    continue

            # Sort by score (descending) and timestamp (ascending)
            return sorted(episodes,
                          key=lambda x: (-x["score"], x.get("timestamp", "")))[:limit]

        except Exception as e:
            logger.error(f"Error in get_recent_episodes: {e}")
            return []
