server/services/api_key_service.py

"""
API Key Service
Manages API keys for programmatic access to the platform.
"""

from __future__ import annotations

import hashlib
import secrets
import logging
from datetime import datetime, timezone
from typing import Dict, List, Optional, Any

logger = logging.getLogger(__name__)

# API pricing in credits
API_COSTS = {
    "hazardguard": 10,
    "weatherwise": 10,
    "geovision": 15,
    "data_layers": 5,
    "chatbot": 10,
    "timelapse": 5,
}


class APIKeyService:
    """Service for managing API keys using Supabase."""

    def __init__(self, supabase_client):
        """
        Initialize the API Key service.
        
        Args:
            supabase_client: Initialized Supabase client with service role key
        """
        self._client = supabase_client

    @staticmethod
    def _generate_key() -> tuple[str, str, str]:
        """
        Generate a new API key.
        
        Returns:
            Tuple of (full_key, key_hash, key_prefix)
        """
        # Generate a secure random key with prefix
        random_part = secrets.token_urlsafe(32)
        full_key = f"sk_live_{random_part}"
        
        # Hash the key for storage
        key_hash = hashlib.sha256(full_key.encode()).hexdigest()
        
        # Store prefix for identification (first 12 chars)
        key_prefix = full_key[:12]
        
        return full_key, key_hash, key_prefix

    @staticmethod
    def _hash_key(api_key: str) -> str:
        """Hash an API key for comparison."""
        return hashlib.sha256(api_key.encode()).hexdigest()

    def create_key(
        self,
        user_id: str,
        name: str,
        permissions: Optional[List[str]] = None,
        expires_at: Optional[datetime] = None,
    ) -> Dict[str, Any]:
        """
        Create a new API key for a user.
        
        Args:
            user_id: The user's UUID
            name: A friendly name for the key
            permissions: List of allowed scopes (defaults to all)
            expires_at: Optional expiration datetime
            
        Returns:
            Dict with key details (includes the full key ONLY on creation)
        """
        try:
            if permissions is None:
                permissions = ["models", "data_layers", "chatbot", "timelapse"]

            full_key, key_hash, key_prefix = self._generate_key()

            data = {
                "user_id": user_id,
                "name": name,
                "key_hash": key_hash,
                "key_prefix": key_prefix,
                "permissions": permissions,
                "is_active": True,
            }

            if expires_at:
                data["expires_at"] = expires_at.isoformat()

            result = self._client.table("api_keys").insert(data).execute()

            if result.data and len(result.data) > 0:
                key_record = result.data[0]
                return {
                    "success": True,
                    "api_key": full_key,  # Only returned on creation!
                    "key_id": key_record["id"],
                    "name": key_record["name"],
                    "key_prefix": key_record["key_prefix"],
                    "permissions": key_record["permissions"],
                    "created_at": key_record["created_at"],
                    "expires_at": key_record.get("expires_at"),
                }

            return {"success": False, "error": "Failed to create API key"}

        except Exception as e:
            logger.error(f"Error creating API key: {e}")
            return {"success": False, "error": str(e)}

    def list_keys(self, user_id: str) -> Dict[str, Any]:
        """
        List all API keys for a user (without exposing the actual keys).
        
        Args:
            user_id: The user's UUID
            
        Returns:
            Dict with list of key metadata
        """
        try:
            result = (
                self._client.table("api_keys")
                .select("id, name, key_prefix, permissions, is_active, last_used_at, usage_count, credits_consumed, created_at, expires_at")
                .eq("user_id", user_id)
                .order("created_at", desc=True)
                .execute()
            )

            return {
                "success": True,
                "keys": result.data or [],
            }

        except Exception as e:
            logger.error(f"Error listing API keys: {e}")
            return {"success": False, "error": str(e), "keys": []}

    def revoke_key(self, user_id: str, key_id: str) -> Dict[str, Any]:
        """
        Revoke (soft delete) an API key.
        
        Args:
            user_id: The user's UUID
            key_id: The key's UUID
            
        Returns:
            Dict with success status
        """
        try:
            result = (
                self._client.table("api_keys")
                .update({"is_active": False})
                .eq("id", key_id)
                .eq("user_id", user_id)
                .execute()
            )

            if result.data and len(result.data) > 0:
                return {"success": True, "message": "API key revoked successfully"}

            return {"success": False, "error": "API key not found"}

        except Exception as e:
            logger.error(f"Error revoking API key: {e}")
            return {"success": False, "error": str(e)}

    def delete_key(self, user_id: str, key_id: str) -> Dict[str, Any]:
        """
        Permanently delete an API key.
        
        Args:
            user_id: The user's UUID
            key_id: The key's UUID
            
        Returns:
            Dict with success status
        """
        try:
            result = (
                self._client.table("api_keys")
                .delete()
                .eq("id", key_id)
                .eq("user_id", user_id)
                .execute()
            )

            return {"success": True, "message": "API key deleted successfully"}

        except Exception as e:
            logger.error(f"Error deleting API key: {e}")
            return {"success": False, "error": str(e)}

    def validate_key(self, api_key: str) -> Dict[str, Any]:
        """
        Validate an API key and return associated user info.
        
        Args:
            api_key: The full API key to validate
            
        Returns:
            Dict with validation result and user info
        """
        try:
            if not api_key or not api_key.startswith("sk_live_"):
                return {"valid": False, "error": "Invalid API key format"}

            key_hash = self._hash_key(api_key)

            result = (
                self._client.table("api_keys")
                .select("id, user_id, name, permissions, is_active, expires_at")
                .eq("key_hash", key_hash)
                .eq("is_active", True)
                .execute()
            )

            if not result.data or len(result.data) == 0:
                return {"valid": False, "error": "API key not found or inactive"}

            key_record = result.data[0]

            # Check expiration
            if key_record.get("expires_at"):
                expires = datetime.fromisoformat(key_record["expires_at"].replace("Z", "+00:00"))
                if expires < datetime.now(timezone.utc):
                    return {"valid": False, "error": "API key has expired"}

            return {
                "valid": True,
                "key_id": key_record["id"],
                "user_id": key_record["user_id"],
                "name": key_record["name"],
                "permissions": key_record["permissions"],
            }

        except Exception as e:
            logger.error(f"Error validating API key: {e}")
            return {"valid": False, "error": str(e)}

    def record_usage(
        self,
        key_id: str,
        user_id: str,
        endpoint: str,
        method: str,
        credits_charged: int,
        status_code: int,
        response_time_ms: int,
        ip_address: Optional[str] = None,
        user_agent: Optional[str] = None,
        request_metadata: Optional[Dict] = None,
    ) -> bool:
        """
        Record an API usage event and update key statistics.
        
        Args:
            key_id: The API key's UUID
            user_id: The user's UUID
            endpoint: The API endpoint called
            method: HTTP method
            credits_charged: Credits deducted
            status_code: HTTP response code
            response_time_ms: Response time in milliseconds
            ip_address: Client IP
            user_agent: Client user agent
            request_metadata: Additional metadata
            
        Returns:
            True if recorded successfully
        """
        try:
            # Insert usage log
            log_data = {
                "api_key_id": key_id,
                "user_id": user_id,
                "endpoint": endpoint,
                "method": method,
                "credits_charged": credits_charged,
                "status_code": status_code,
                "response_time_ms": response_time_ms,
                "ip_address": ip_address,
                "user_agent": user_agent,
                "request_metadata": request_metadata or {},
            }

            self._client.table("api_usage_logs").insert(log_data).execute()

            # Update key statistics
            self._client.rpc(
                "increment_api_key_usage",
                {"p_key_id": key_id, "p_credits": credits_charged}
            ).execute()

            return True

        except Exception as e:
            logger.error(f"Error recording API usage: {e}")
            # Don't fail the request if logging fails
            return False

    def get_usage_stats(self, user_id: str, key_id: Optional[str] = None, days: int = 30) -> Dict[str, Any]:
        """
        Get usage statistics for a user's API keys.
        
        Args:
            user_id: The user's UUID
            key_id: Optional specific key to filter by
            days: Number of days to look back
            
        Returns:
            Dict with usage statistics
        """
        try:
            query = (
                self._client.table("api_usage_logs")
                .select("endpoint, credits_charged, created_at, status_code")
                .eq("user_id", user_id)
                .gte("created_at", f"now() - interval '{days} days'")
                .order("created_at", desc=True)
                .limit(1000)
            )

            if key_id:
                query = query.eq("api_key_id", key_id)

            result = query.execute()

            logs = result.data or []

            # Aggregate stats
            total_calls = len(logs)
            total_credits = sum(log.get("credits_charged", 0) for log in logs)
            
            endpoint_stats = {}
            for log in logs:
                ep = log.get("endpoint", "unknown")
                if ep not in endpoint_stats:
                    endpoint_stats[ep] = {"calls": 0, "credits": 0}
                endpoint_stats[ep]["calls"] += 1
                endpoint_stats[ep]["credits"] += log.get("credits_charged", 0)

            return {
                "success": True,
                "total_calls": total_calls,
                "total_credits": total_credits,
                "endpoint_breakdown": endpoint_stats,
                "recent_logs": logs[:50],  # Last 50 calls
            }

        except Exception as e:
            logger.error(f"Error getting usage stats: {e}")
            return {"success": False, "error": str(e)}

    @staticmethod
    def get_api_costs() -> Dict[str, int]:
        """Return the API cost structure."""
        return API_COSTS.copy()

    def increment_credits_consumed(self, key_id: str, credits: int) -> bool:
        """
        Increment the credits_consumed counter for an API key.
        
        Args:
            key_id: The API key's UUID
            credits: Number of credits to add to the counter
            
        Returns:
            True if updated successfully
        """
        try:
            # Use direct SQL update to increment
            self._client.rpc(
                "increment_api_key_credits",
                {"p_key_id": key_id, "p_credits": credits}
            ).execute()
            return True
        except Exception as e:
            # Fallback: try direct update
            try:
                result = (
                    self._client.table("api_keys")
                    .select("credits_consumed")
                    .eq("id", key_id)
                    .execute()
                )
                if result.data and len(result.data) > 0:
                    current = result.data[0].get("credits_consumed", 0) or 0
                    self._client.table("api_keys").update({
                        "credits_consumed": current + credits
                    }).eq("id", key_id).execute()
                    return True
            except Exception as e2:
                logger.error(f"Error incrementing credits_consumed: {e2}")
            return False