"""
SupabaseVectorStorage: Vector storage adapter using Supabase pgvector.

This module implements LightRAG's BaseVectorStorage interface using
Supabase's PostgreSQL with pgvector extension for optimal vector similarity search.
"""

from __future__ import annotations

import os
import json
from datetime import datetime, timezone
from dataclasses import dataclass, field
from typing import Any

from dotenv import load_dotenv

load_dotenv(dotenv_path=".env", override=False)

# Import base class from lightrag
try:
    from lightrag.base import BaseVectorStorage
except ImportError:
    from dataclasses import dataclass as base_dataclass
    
    @base_dataclass
    class BaseVectorStorage:
        namespace: str = ""
        workspace: str = ""
        global_config: dict = field(default_factory=dict)
        embedding_func: Any = None


@dataclass
class SupabaseVectorStorage(BaseVectorStorage):
    """
    Vector storage adapter using Supabase's pgvector extension.
    
    This implementation provides:
    - Fast vector similarity search
    - Metadata filtering
    - Automatic distance calculation
    - Hybrid search capabilities
    
    Configuration via environment variables:
    - SUPABASE_URL: Supabase project URL
    - SUPABASE_KEY: Supabase API key
    - SUPABASE_TABLE_NAME: Table name for vectors (default: embeddings)
    """
    
    # Supabase configuration
    supabase_url: str = field(
        default_factory=lambda: os.getenv("SUPABASE_URL", "")
    )
    supabase_key: str = field(
        default_factory=lambda: os.getenv("SUPABASE_KEY", "")
    )
    table_name: str = field(
        default_factory=lambda: os.getenv("SUPABASE_TABLE_NAME", "embeddings")
    )
    
    # Meta fields for filtering, compatible with LightRAG interface
    meta_fields: set[str] = field(default_factory=set)
    
    # Embedding configuration
    embedding_dim: int = field(
        default_factory=lambda: int(os.getenv("EMBEDDING_DIM", "1536"))
    )
    
    # Client instance
    _client: Any = None
    _initialized: bool = False
    
    async def initialize(self) -> None:
        """Initialize Supabase client."""
        if self._initialized:
            return
            
        if not self.supabase_url or not self.supabase_key:
            raise ValueError(
                "SUPABASE_URL and SUPABASE_KEY environment variables are required"
            )
        
        try:
            from supabase import create_async_client
        except ImportError:
            raise ImportError(
                "supabase is required. Install with: pip install supabase"
            )
        
        self._client = await create_async_client(
            self.supabase_url,
            self.supabase_key,
        )
        
        # Ensure table exists with proper schema
        await self._ensure_table_exists()
        self._initialized = True
    
    async def _ensure_table_exists(self) -> None:
        """
        Ensure the embeddings table exists with proper schema.
        
        Expected schema (should be created via Supabase SQL editor):
        ```sql
        create table if not exists embeddings (
            id uuid primary key default gen_random_uuid(),
            content text not null,
            embedding vector(1536),
            metadata jsonb default '{}',
            namespace text default 'default',
            created_at timestamptz default now()
        );
        
        create index on embeddings using ivfflat (embedding vector_cosine_ops)
        with (lists = 100);
        ```
        """
        # Note: Table creation should be done via Supabase dashboard or migrations
        # This is just a check
        pass
    
    async def finalize(self) -> None:
        """Close Supabase client."""
        self._client = None
        self._initialized = False
    
    async def index_done_callback(self) -> None:
        """Called after batch indexing is complete."""
        pass
    
    async def query(
        self,
        query: str,
        top_k: int = 10,
        query_embedding: list[float] | None = None,
        **kwargs,
    ) -> list[dict]:
        """
        Search for similar vectors using cosine similarity.
        
        Args:
            query: Query text (used for logging/debugging).
            top_k: Number of results to return.
            query_embedding: Pre-computed query embedding.
            **kwargs: Additional filter parameters.
            
        Returns:
            List of matching documents with similarity scores.
        """
        if not self._initialized:
            await self.initialize()
        
        if query_embedding is None:
            # Use the embedding function from global config
            if self.embedding_func:
                # embedding_func expects a list of strings, not a bare string
                result = await self.embedding_func([query])
                # result is np.array of shape (1, dim) or list of lists
                if hasattr(result, "tolist"):
                    result = result.tolist()
                query_embedding = result[0] if isinstance(result[0], list) else result
            else:
                raise ValueError("query_embedding is required when embedding_func is not set")
                
        # Ensure query_embedding is a flat list of floats
        if hasattr(query_embedding, "tolist"):
            query_embedding = query_embedding.tolist()
        
        # If it's a list of lists (e.g. [[0.1, ...]]), take the first one
        if isinstance(query_embedding, list) and len(query_embedding) > 0 and isinstance(query_embedding[0], list):
            query_embedding = query_embedding[0]
        
        # Call Supabase RPC function for vector similarity search
        # This function should be created in Supabase:
        # ```sql
        # create or replace function match_embeddings(
        #     query_embedding vector(1536),
        #     match_count int default 10,
        #     filter_namespace text default null
        # )
        # returns table (
        #     id uuid,
        #     content text,
        #     metadata jsonb,
        #     similarity float
        # )
        # language plpgsql
        # as $$
        # begin
        #     return query
        #     select
        #         id,
        #         content,
        #         metadata,
        #         1 - (embedding <=> query_embedding) as similarity
        #     from embeddings
        #     where filter_namespace is null or namespace = filter_namespace
        #     order by embedding <=> query_embedding
        #     limit match_count;
        # end;
        # $$;
        # ```
        
        try:
            response = await self._client.rpc(
                "match_embeddings",
                {
                    "query_embedding": query_embedding,
                    "match_count": top_k,
                    "filter_namespace": self.namespace or None,
                },
            ).execute()
            
            return [
                {
                    **row,
                    **(row.get("metadata", {}) or {}),
                    "score": row.get("similarity", 0.0),
                }
                for row in response.data
            ]
        except Exception as e:
            from lightrag.utils import logger
            logger.error(f"Vector search error: {e}")
            return []
    async def delete_entity(self, entity_name: str) -> None:
        """Delete a single entity by its name."""
        try:
            from lightrag.utils import compute_mdhash_id
            entity_id = compute_mdhash_id(entity_name, prefix="ent-")
            await self.delete([entity_id])
        except Exception as e:
            print(f"Delete entity error: {e}")

    async def delete_entity_relation(self, entity_name: str) -> None:
        """Delete relations for a given entity."""
        if not self._initialized:
            await self.initialize()
            
        try:
            # Search for relations in metadata where src_id or tgt_id matches
            # This requires querying by metadata fields
            response = await self._client.table(self.table_name).select("id").or_(
                f"metadata->>src_id.eq.{entity_name},metadata->>tgt_id.eq.{entity_name}"
            ).execute()
            
            ids_to_delete = [row["id"] for row in response.data]
            if ids_to_delete:
                await self.delete(ids_to_delete)
        except Exception as e:
            print(f"Delete entity relation error: {e}")

    async def get_by_id(self, id: str) -> dict[str, Any] | None:
        """Get vector data by its ID."""
        if not self._initialized:
            await self.initialize()
            
        try:
            response = await self._client.table(self.table_name).select("*").eq("id", id).limit(1).execute()
            if response.data:
                row = response.data[0]
                metadata = row.get("metadata", {})
                if isinstance(metadata, str):
                    metadata = json.loads(metadata)
                return {
                    "id": row["id"],
                    "content": row["content"],
                    "metadata": metadata,
                }
            return None
        except Exception as e:
            print(f"Get by ID error: {e}")
            return None

    async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
        """Get multiple vector data by their IDs."""
        if not ids:
            return []
        if not self._initialized:
            await self.initialize()
            
        try:
            response = await self._client.table(self.table_name).select("*").in_("id", ids).execute()
            results = []
            for row in response.data:
                metadata = row.get("metadata", {})
                if isinstance(metadata, str):
                    metadata = json.loads(metadata)
                results.append({
                    "id": row["id"],
                    "content": row["content"],
                    "metadata": metadata,
                })
            return results
        except Exception as e:
            print(f"Get by IDs error: {e}")
            return []

    async def get_vectors_by_ids(self, ids: list[str]) -> dict[str, list[float]]:
        """Get vectors by their IDs."""
        if not ids:
            return {}
        if not self._initialized:
            await self.initialize()
            
        try:
            response = await self._client.table(self.table_name).select("id, embedding").in_("id", ids).execute()
            # Handle list return and potential dict format from different Supabase versions
            results = {}
            for row in response.data:
                emb = row["embedding"]
                if isinstance(emb, str): # Handle text-represented vectors if any
                    emb = json.loads(emb)
                results[row["id"]] = emb
            return results
        except Exception as e:
            from lightrag.utils import logger
            logger.error(f"Get vectors by IDs error: {e}")
            return {}
    
    async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
        """
        Batch insert or update vector embeddings.
        
        This implementation follows the LightRAG BaseVectorStorage interface:
        1. It takes a dictionary of {id: {content, embedding, ...}}
        2. It automatically generates embeddings for records that are missing them
        3. It filters metadata based on self.meta_fields
        4. It performs a batch upsert to Supabase
        
        Args:
            data: Dictionary mapping IDs to their data (content, embedding, metadata)
        """
        if not data:
            return
            
        if not self._initialized:
            await self.initialize()
            
        from lightrag.utils import logger
        import numpy as np
        import asyncio

        # 1. Identify records that need embeddings
        records_to_embed = []
        ids_to_embed = []
        
        for doc_id, doc_data in data.items():
            embedding = doc_data.get("embedding")
            # If embedding is missing, empty, or None, we need to generate it
            needs_embedding = False
            if embedding is None:
                needs_embedding = True
            elif isinstance(embedding, (list, np.ndarray)):
                if len(embedding) == 0:
                    needs_embedding = True
            
            if needs_embedding:
                content = doc_data.get("content")
                if content:
                    ids_to_embed.append(doc_id)
                    records_to_embed.append(content)
                else:
                    logger.warning(f"Upsert record {doc_id} has no content and no embedding, skipping.")

        # 2. Generate embeddings in batches if needed
        if records_to_embed and self.embedding_func:
            # Respect batch size from global_config or default to 10
            batch_size = self.global_config.get("embedding_batch_num", 10)
            all_embeddings = []
            
            # Use chunks for parallel processing
            for i in range(0, len(records_to_embed), batch_size):
                batch = records_to_embed[i : i + batch_size]
                batch_embeddings = await self.embedding_func(batch)
                
                # Handle numpy or list returns
                if hasattr(batch_embeddings, "tolist"):
                    batch_embeddings = batch_embeddings.tolist()
                
                # EmbeddingFunc might return a single vector if batch size is 1?
                # Usually it returns a list of lists.
                if len(batch) == 1 and len(batch_embeddings) > 0 and not isinstance(batch_embeddings[0], list):
                     all_embeddings.append(batch_embeddings)
                else:
                     all_embeddings.extend(batch_embeddings)
            
            # Update data with generated embeddings
            for doc_id, emb in zip(ids_to_embed, all_embeddings):
                data[doc_id]["embedding"] = emb

        # 3. Prepare records for Supabase
        supabase_records = []
        for doc_id, doc_data in data.items():
            content = doc_data.get("content", "")
            embedding = doc_data.get("embedding")
            
            # Ensure embedding is a list (not numpy array) for Supabase
            if hasattr(embedding, "tolist"):
                embedding = embedding.tolist()
                
            if embedding is None or (isinstance(embedding, list) and not embedding):
                 logger.warning(f"Upsert skipping doc {doc_id}: Missing embedding after processing.")
                 continue

            # Filter metadata based on meta_fields
            metadata = {k: v for k, v in doc_data.items() if k in self.meta_fields}
            
            supabase_records.append({
                "id": doc_id,
                "content": content,
                "embedding": embedding,
                "metadata": metadata,
                "namespace": self.namespace or "default",
            })

        if not supabase_records:
            return

        # 4. Perform Supabase Upsert
        try:
            # supabase-py handles batch upsert
            await self._client.table(self.table_name).upsert(
                supabase_records,
                on_conflict="id",
            ).execute()
            logger.debug(f"Successfully upserted {len(supabase_records)} records to Supabase table {self.table_name}")
        except Exception as e:
            logger.error(f"Supabase batch upsert error: {e}")
    
    async def delete(self, ids: list[str]) -> None:
        """Delete embeddings by IDs."""
        if not self._initialized:
            await self.initialize()
        
        try:
            await self._client.table(self.table_name).delete().in_(
                "id", ids
            ).execute()
        except Exception as e:
            print(f"Delete error: {e}")
    
    async def drop(self) -> dict[str, str]:
        """Drop all data for this namespace."""
        if not self._initialized:
            await self.initialize()
        
        try:
            await self._client.table(self.table_name).delete().eq(
                "namespace", self.namespace or "default"
            ).execute()
            return {"status": "success"}
        except Exception as e:
            return {"status": "error", "message": str(e)}
