src/services/supabase_service.py

"""
Supabase integration service for direct database access.

This service provides a bridge between the backend and Supabase PostgreSQL,
allowing investigations to be stored centrally for frontend consumption.
"""

import os
from typing import Optional, List, Dict, Any
from datetime import datetime
from contextlib import asynccontextmanager

from asyncpg import Pool, create_pool, Connection
from pydantic import BaseModel, Field

from src.core import get_logger, settings
from src.core.exceptions import CidadaoAIError

logger = get_logger(__name__)


class SupabaseConfig(BaseModel):
    """Supabase connection configuration."""

    url: str = Field(..., description="Supabase PostgreSQL connection URL")
    anon_key: Optional[str] = Field(None, description="Supabase anon key (for Row Level Security)")
    service_role_key: Optional[str] = Field(None, description="Supabase service role key (bypasses RLS)")
    min_connections: int = Field(default=5, description="Minimum pool connections")
    max_connections: int = Field(default=20, description="Maximum pool connections")

    @classmethod
    def from_env(cls) -> "SupabaseConfig":
        """Load configuration from environment variables."""
        supabase_url = os.getenv("SUPABASE_DB_URL") or os.getenv("DATABASE_URL")

        if not supabase_url:
            raise ValueError(
                "SUPABASE_DB_URL or DATABASE_URL environment variable required. "
                "Get it from: Supabase Dashboard > Settings > Database > Connection string (URI)"
            )

        return cls(
            url=supabase_url,
            anon_key=os.getenv("SUPABASE_ANON_KEY"),
            service_role_key=os.getenv("SUPABASE_SERVICE_ROLE_KEY"),
            min_connections=int(os.getenv("SUPABASE_MIN_CONNECTIONS", "5")),
            max_connections=int(os.getenv("SUPABASE_MAX_CONNECTIONS", "20")),
        )


class SupabaseService:
    """
    Service for interacting with Supabase PostgreSQL.

    Provides connection pooling and CRUD operations for investigations.
    """

    def __init__(self, config: Optional[SupabaseConfig] = None):
        """
        Initialize Supabase service.

        Args:
            config: Supabase configuration (loads from env if None)
        """
        self.config = config or SupabaseConfig.from_env()
        self._pool: Optional[Pool] = None
        self._initialized = False

    async def initialize(self) -> None:
        """Initialize connection pool."""
        if self._initialized:
            logger.warning("Supabase service already initialized")
            return

        try:
            logger.info("Initializing Supabase connection pool")

            self._pool = await create_pool(
                dsn=self.config.url,
                min_size=self.config.min_connections,
                max_size=self.config.max_connections,
                command_timeout=30,
                server_settings={
                    'application_name': 'cidadao-ai-backend',
                    'timezone': 'UTC',
                }
            )

            # Test connection
            async with self._pool.acquire() as conn:
                version = await conn.fetchval("SELECT version()")
                logger.info(f"Connected to Supabase PostgreSQL: {version[:50]}...")

            self._initialized = True
            logger.info("Supabase service initialized successfully")

        except Exception as e:
            logger.error(f"Failed to initialize Supabase service: {e}", exc_info=True)
            raise CidadaoAIError(f"Supabase initialization failed: {e}")

    async def close(self) -> None:
        """Close connection pool."""
        if self._pool:
            await self._pool.close()
            self._initialized = False
            logger.info("Supabase connection pool closed")

    @asynccontextmanager
    async def get_connection(self):
        """
        Get a database connection from the pool.

        Yields:
            Connection instance
        """
        if not self._initialized:
            await self.initialize()

        async with self._pool.acquire() as conn:
            yield conn

    async def create_investigation(
        self,
        user_id: str,
        query: str,
        data_source: str,
        filters: Optional[Dict[str, Any]] = None,
        anomaly_types: Optional[List[str]] = None,
        session_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Create a new investigation in Supabase.

        Args:
            user_id: User ID
            query: Investigation query
            data_source: Data source to investigate
            filters: Query filters
            anomaly_types: Types of anomalies to detect
            session_id: Optional session ID

        Returns:
            Created investigation as dict
        """
        async with self.get_connection() as conn:
            import json

            row = await conn.fetchrow(
                """
                INSERT INTO investigations (
                    user_id, session_id, query, data_source,
                    status, filters, anomaly_types, progress,
                    created_at, updated_at
                )
                VALUES ($1, $2, $3, $4, $5, $6::jsonb, $7::jsonb, $8, $9, $10)
                RETURNING *
                """,
                user_id,
                session_id,
                query,
                data_source,
                "pending",
                json.dumps(filters or {}),
                json.dumps(anomaly_types or []),
                0.0,
                datetime.utcnow(),
                datetime.utcnow(),
            )

            logger.info(f"Created investigation {row['id']} in Supabase")
            return dict(row)

    async def get_investigation(self, investigation_id: str) -> Optional[Dict[str, Any]]:
        """
        Get investigation by ID.

        Args:
            investigation_id: Investigation UUID

        Returns:
            Investigation dict or None
        """
        async with self.get_connection() as conn:
            row = await conn.fetchrow(
                "SELECT * FROM investigations WHERE id = $1",
                investigation_id
            )

            return dict(row) if row else None

    async def update_investigation(
        self,
        investigation_id: str,
        **updates
    ) -> Dict[str, Any]:
        """
        Update investigation fields.

        Args:
            investigation_id: Investigation UUID
            **updates: Fields to update

        Returns:
            Updated investigation dict
        """
        import json

        # JSONB fields that need special handling
        jsonb_fields = {'results', 'filters', 'anomaly_types'}

        # Build dynamic UPDATE query
        set_clauses = []
        values = []
        param_index = 1

        for key, value in updates.items():
            if key in jsonb_fields and isinstance(value, (dict, list)):
                set_clauses.append(f"{key} = ${param_index}::jsonb")
                values.append(json.dumps(value))
            else:
                set_clauses.append(f"{key} = ${param_index}")
                values.append(value)
            param_index += 1

        # Always update updated_at
        set_clauses.append(f"updated_at = ${param_index}")
        values.append(datetime.utcnow())
        param_index += 1

        # Add investigation_id as last parameter
        values.append(investigation_id)

        query = f"""
            UPDATE investigations
            SET {', '.join(set_clauses)}
            WHERE id = ${param_index}
            RETURNING *
        """

        async with self.get_connection() as conn:
            row = await conn.fetchrow(query, *values)

            if not row:
                raise ValueError(f"Investigation {investigation_id} not found")

            logger.debug(f"Updated investigation {investigation_id}")
            return dict(row)

    async def update_progress(
        self,
        investigation_id: str,
        progress: float,
        current_phase: str,
        records_processed: Optional[int] = None,
        anomalies_found: Optional[int] = None,
    ) -> Dict[str, Any]:
        """
        Update investigation progress.

        Args:
            investigation_id: Investigation UUID
            progress: Progress percentage (0.0 to 1.0)
            current_phase: Current processing phase
            records_processed: Number of records processed
            anomalies_found: Number of anomalies detected

        Returns:
            Updated investigation dict
        """
        updates = {
            "progress": progress,
            "current_phase": current_phase,
        }

        if records_processed is not None:
            updates["total_records_analyzed"] = records_processed

        if anomalies_found is not None:
            updates["anomalies_found"] = anomalies_found

        return await self.update_investigation(investigation_id, **updates)

    async def complete_investigation(
        self,
        investigation_id: str,
        results: List[Dict[str, Any]],
        summary: str,
        confidence_score: float,
        total_records: int,
        anomalies_found: int,
    ) -> Dict[str, Any]:
        """
        Mark investigation as completed with results.

        Args:
            investigation_id: Investigation UUID
            results: List of anomaly results
            summary: Investigation summary
            confidence_score: Overall confidence
            total_records: Total records analyzed
            anomalies_found: Total anomalies found

        Returns:
            Updated investigation dict
        """
        return await self.update_investigation(
            investigation_id,
            status="completed",
            progress=1.0,
            current_phase="completed",
            results=results,
            summary=summary,
            confidence_score=confidence_score,
            total_records_analyzed=total_records,
            anomalies_found=anomalies_found,
            completed_at=datetime.utcnow(),
        )

    async def fail_investigation(
        self,
        investigation_id: str,
        error_message: str,
    ) -> Dict[str, Any]:
        """
        Mark investigation as failed.

        Args:
            investigation_id: Investigation UUID
            error_message: Error description

        Returns:
            Updated investigation dict
        """
        return await self.update_investigation(
            investigation_id,
            status="failed",
            current_phase="failed",
            error_message=error_message,
            completed_at=datetime.utcnow(),
        )

    async def list_user_investigations(
        self,
        user_id: str,
        limit: int = 20,
        offset: int = 0,
        status: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """
        List investigations for a user.

        Args:
            user_id: User ID
            limit: Maximum results
            offset: Pagination offset
            status: Filter by status

        Returns:
            List of investigation dicts
        """
        async with self.get_connection() as conn:
            query = """
                SELECT * FROM investigations
                WHERE user_id = $1
            """
            params = [user_id]

            if status:
                query += " AND status = $2"
                params.append(status)

            query += " ORDER BY created_at DESC LIMIT $" + str(len(params) + 1)
            params.append(limit)

            query += " OFFSET $" + str(len(params) + 1)
            params.append(offset)

            rows = await conn.fetch(query, *params)

            return [dict(row) for row in rows]

    async def delete_investigation(
        self,
        investigation_id: str,
        user_id: str,
    ) -> bool:
        """
        Delete an investigation (soft delete by marking as cancelled).

        Args:
            investigation_id: Investigation UUID
            user_id: User ID (for authorization)

        Returns:
            True if deleted, False if not found
        """
        async with self.get_connection() as conn:
            result = await conn.execute(
                """
                UPDATE investigations
                SET status = 'cancelled', completed_at = $1
                WHERE id = $2 AND user_id = $3
                """,
                datetime.utcnow(),
                investigation_id,
                user_id,
            )

            # Extract number of rows affected
            rows_affected = int(result.split()[-1])

            if rows_affected > 0:
                logger.info(f"Cancelled investigation {investigation_id}")
                return True

            return False

    async def health_check(self) -> Dict[str, Any]:
        """
        Check Supabase connection health.

        Returns:
            Health status dict
        """
        try:
            async with self.get_connection() as conn:
                # Simple query to test connection
                result = await conn.fetchval("SELECT 1")
                pool_size = self._pool.get_size()
                pool_free = self._pool.get_idle_size()

                return {
                    "status": "healthy",
                    "connected": True,
                    "pool_size": pool_size,
                    "pool_free": pool_free,
                    "pool_used": pool_size - pool_free,
                }
        except Exception as e:
            logger.error(f"Supabase health check failed: {e}")
            return {
                "status": "unhealthy",
                "connected": False,
                "error": str(e),
            }


# Global service instance
supabase_service = SupabaseService()


async def get_supabase_service() -> SupabaseService:
    """Get the global Supabase service instance."""
    if not supabase_service._initialized:
        await supabase_service.initialize()
    return supabase_service