"""
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")
    )
    
    # 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:
                query_embedding = await self.embedding_func(query)
            else:
                raise ValueError("query_embedding is required when embedding_func is not set")
        
        # 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 [
                {
                    "id": row["id"],
                    "content": row["content"],
                    "metadata": row.get("metadata", {}),
                    "score": row.get("similarity", 0.0),
                }
                for row in response.data
            ]
        except Exception as e:
            # Log and return empty results
            print(f"Vector search error: {e}")
            return []
    
    async def upsert(self, data: dict) -> None:
        """
        Insert or update vector embedding.
        
        Args:
            data: Dict containing:
                - id: Unique identifier
                - content: Text content
                - embedding: Vector embedding
                - metadata: Optional metadata dict
        """
        if not self._initialized:
            await self.initialize()
        
        doc_id = data.get("id")
        content = data.get("content", "")
        embedding = data.get("embedding", [])
        metadata = data.get("metadata", {})
        
        record = {
            "content": content,
            "embedding": embedding,
            "metadata": json.dumps(metadata) if isinstance(metadata, dict) else metadata,
            "namespace": self.namespace or "default",
        }
        
        if doc_id:
            record["id"] = doc_id
        
        try:
            await self._client.table(self.table_name).upsert(
                record,
                on_conflict="id",
            ).execute()
        except Exception as e:
            print(f"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)}
