src/datalake/catalog.py

"""
Data lake catalog for discovering structure and metadata using AWS Athena/Glue.

Provides methods to explore the data lake organization using Athena metadata:
- List devices, messages, and dates from table structure
- Get schemas for message/rule tables
- Understand data availability
"""

from typing import List, Dict, Optional
import re
from .athena import AthenaQuery
from .config import DataLakeConfig
from .logger import setup_logger

logger = setup_logger(__name__)


class DataLakeCatalog:
    """
    Catalog for exploring data lake structure using Athena/Glue.
    
    Assumes Athena database contains tables organized by device and message.
    Table naming convention: {device_id}_{message_rule} or similar
    """
    
    def __init__(self, athena_query: AthenaQuery, config: DataLakeConfig):
        """
        Initialize catalog.
        
        Args:
            athena_query: AthenaQuery instance
            config: DataLakeConfig instance
        """
        self.athena = athena_query
        self.config = config
        self._cache: Dict[str, Dict] = {}
        logger.info(f"Initialized catalog for database: {config.database_name}")
    
    def list_tables(self) -> List[str]:
        """
        List all tables in the database.
        
        Returns:
            Sorted list of table names
        """
        tables = self.athena.list_tables()
        logger.info(f"Found {len(tables)} tables in database")
        return sorted(tables)
    
    def list_devices(self, device_filter: Optional[str] = None) -> List[str]:
        """
        List all device IDs by extracting from table names.
        
        Args:
            device_filter: Optional regex pattern to filter devices
        
        Returns:
            Sorted list of device IDs
        
        Note:
            Extracts device IDs from table names. Assumes table naming like:
            - {device_id}_{message_rule}
            - {device_id}__{message_rule}
            - Or similar pattern
        """
        tables = self.list_tables()
        devices = set()
        
        for table in tables:
            # Try common patterns: device_message, device__message, device.message
            parts = re.split(r'_', table, maxsplit=2)
            if len(parts) >= 1:
                device = parts[1]
                if device == 'aggregations': # skip aggregations table
                    continue
                if device_filter is None or re.search(device_filter, device):
                    devices.add(device)
        
        result = sorted(devices)
        logger.info(f"Found {len(result)} device(s)")
        return result
    
    def list_messages(self, device_id: str, message_filter: Optional[str] = None) -> List[str]:
        """
        List all message/rule names for a device.
        
        Args:
            device_id: Device identifier
            message_filter: Optional regex pattern to filter messages
        
        Returns:
            Sorted list of message/rule names
        
        Note:
            Extracts message names from table names. Assumes table naming like:
            - prefix_{device_id}_{message_rule}
            - Or {device_id}_{message_rule}
        """
        tables = self.list_tables()
        messages = set()
        
        for table in tables:
            # Split table name by underscore (consistent with list_devices)
            parts = re.split(r'_', table, maxsplit=2)
            
            # Try pattern: prefix_device_message
            if len(parts) >= 3:
                table_device = parts[1]
                if table_device == device_id:
                    message = parts[2]
                    if message_filter is None or re.search(message_filter, message):
                        messages.add(message)
            # Try pattern: device_message (no prefix)
            elif len(parts) >= 2:
                table_device = parts[0]
                if table_device == device_id:
                    message = parts[1]
                    if message_filter is None or re.search(message_filter, message):
                        messages.add(message)
        
        result = sorted(messages)
        logger.info(f"Found {len(result)} messages for device {device_id}")
        return result
    
    def get_table_name(self, device_id: str, message: str) -> str:
        """
        Get table name for device/message combination.
        
        Args:
            device_id: Device identifier
            message: Message/rule name
        
        Returns:
            Table name (tries common patterns)
        
        Raises:
            ValueError: If table not found
        """
        tables = self.list_tables()
        
        # Try patterns consistent with list_devices/list_messages
        # Pattern 1: prefix_device_message
        for table in tables:
            parts = re.split(r'_', table, maxsplit=2)
            if len(parts) >= 3:
                if parts[1] == device_id and parts[2] == message:
                    return table
        
        # Pattern 2: device_message (no prefix)
        for table in tables:
            parts = re.split(r'_', table, maxsplit=1)
            if len(parts) >= 2:
                if parts[0] == device_id and parts[1] == message:
                    return table
        
        # Fallback: try exact matches
        patterns = [
            f"{device_id}_{message}",
            f"{device_id}__{message}",
            f"{device_id}_{message}".lower(),
            f"{device_id}__{message}".lower(),
        ]
        
        for pattern in patterns:
            if pattern in tables:
                return pattern
        
        raise ValueError(
            f"Table not found for {device_id}/{message}. "
            f"Available tables: {tables[:10]}..."
        )
    
    def get_schema(self, device_id: str, message: str) -> Optional[Dict[str, str]]:
        """
        Get schema for a message table.
        
        Args:
            device_id: Device identifier
            message: Message/rule name
        
        Returns:
            Dict mapping column names to data types, or None if not found
        """
        cache_key = f"{device_id}/{message}"
        if cache_key in self._cache:
            logger.debug(f"Using cached schema for {cache_key}")
            return self._cache[cache_key]
        
        try:
            table_name = self.get_table_name(device_id, message)
            schema_df = self.athena.describe_table(table_name)
            
            if schema_df.empty:
                logger.warning(f"No schema found for {device_id}/{message}")
                return None
            
            schema_dict = {
                row['column_name']: row['data_type']
                for _, row in schema_df.iterrows()
            }
            
            self._cache[cache_key] = schema_dict
            logger.info(f"Schema for {cache_key}: {len(schema_dict)} columns")
            return schema_dict
        except Exception as e:
            logger.error(f"Failed to get schema for {device_id}/{message}: {e}")
            return None
    
    def list_partitions(self, device_id: str, message: str) -> List[str]:
        """
        List partition values (dates) for a table.
        
        Args:
            device_id: Device identifier
            message: Message/rule name
        
        Returns:
            List of partition values (dates) in YYYY-MM-DD format
        
        Note:
            Handles hierarchical partitioning format: year=YYYY/month=MM/day=DD
            Data structure: {device_id}/{message}/{year}/{month}/{day}/file.parquet
        """
        try:
            table_name = self.get_table_name(device_id, message)
            
            # Query partition information
            # query = f"SHOW PARTITIONS {self.config.database_name}.{table_name}"
            query = f"""
                WITH files AS (
                SELECT DISTINCT "$path" AS p
                FROM {self.config.database_name}.{table_name}
                WHERE "$path" LIKE '%.parquet'
                ),
                parts AS (
                SELECT
                    try_cast(element_at(split(p, '/'), -4) AS INTEGER) AS year,
                    try_cast(element_at(split(p, '/'), -3) AS INTEGER) AS month,
                    try_cast(element_at(split(p, '/'), -2) AS INTEGER) AS day
                FROM files
                )
                SELECT DISTINCT year, month, day
                FROM parts
                WHERE year IS NOT NULL AND month IS NOT NULL AND day IS NOT NULL
                ORDER BY year DESC, month DESC, day DESC
                """
            df = self.athena.query_to_dataframe(query)
            
            if df.empty:
                logger.warning(f"No partitions found for {table_name}")
                return []
            
            # Extract date from partition string
            # Format: YYYY-MM-DD
            dates = []
            for _, row in df.iterrows():
                dates.append(f'{row.iloc[0]}-{row.iloc[1]}-{row.iloc[2]:02d}')
            logger.info(f"Found {len(dates)} partitions for {table_name}")
            return sorted(set(dates))
        except Exception as e:
            logger.warning(f"Could not list partitions for {device_id}/{message}: {e}")
            # Table might not be partitioned or query might have failed
            return []
    
    def clear_cache(self) -> None:
        """Clear schema cache."""
        self._cache.clear()
        logger.debug("Schema cache cleared")