#!/usr/bin/env python3
"""
Table Query - Data Access Layer
================================
Abstracts data sources (SQLite, CSV, etc.) for graph module.

Provides adapter pattern for different data sources while maintaining
consistent interface for plotting functions.

Features:
- Abstract DataAdapter interface
- SQLiteAdapter for database files
- CSVAdapter for CSV files
- Column utilities and labels
- Data alignment for drift comparison

Author: Kynsight
Version: 1.0.0
"""

from __future__ import annotations
from abc import ABC, abstractmethod
from typing import List, Dict, Tuple, Optional, Any
from dataclasses import dataclass
from pathlib import Path
import sqlite3

import numpy as np


# =============================================================================
# Data Structures
# =============================================================================

@dataclass
class SessionInfo:
    """Session metadata."""
    session_id: str
    created_at: str
    description: str
    run_count: int


@dataclass
class RunInfo:
    """Run metadata."""
    session_id: str
    run_name: str
    run_number: int
    sample_count: int
    start_time_ns: int
    end_time_ns: int
    duration_ms: float


@dataclass
class TelemetryData:
    """Decoded telemetry data (source-agnostic)."""
    session_id: str
    run_name: str
    run_no: int
    
    # Time axes
    t_ns: Optional[np.ndarray] = None
    time_ms: Optional[np.ndarray] = None
    
    # UART decoded data
    motor_current: Optional[np.ndarray] = None
    encoder_value: Optional[np.ndarray] = None
    relative_encoder_value: Optional[np.ndarray] = None
    v24_pec_diff: Optional[np.ndarray] = None
    pwm: Optional[np.ndarray] = None
    
    # I2C decoded data
    i2c_raw14: Optional[np.ndarray] = None
    i2c_zero_raw14: Optional[np.ndarray] = None
    i2c_delta_raw14: Optional[np.ndarray] = None
    i2c_angle_deg: Optional[np.ndarray] = None
    i2c_zero_angle_deg: Optional[np.ndarray] = None
    
    # Derived data
    angular_velocity: Optional[np.ndarray] = None


# =============================================================================
# Abstract Adapter
# =============================================================================

class DataAdapter(ABC):
    """
    Abstract data source adapter.
    
    Implement this interface for each data source type (SQLite, CSV, etc.)
    All adapters provide the same interface for loading telemetry data.
    """
    
    @abstractmethod
    def connect(self) -> bool:
        """
        Connect to data source.
        
        Returns:
            True on success, False on failure
        """
        pass
    
    @abstractmethod
    def close(self):
        """Close connection to data source."""
        pass
    
    @abstractmethod
    def get_sessions(self) -> List[SessionInfo]:
        """
        Get all sessions from data source.
        
        Returns:
            List of SessionInfo objects
        """
        pass
    
    @abstractmethod
    def get_runs(self, session_id: str) -> List[RunInfo]:
        """
        Get all runs for a specific session.
        
        Args:
            session_id: Session identifier
        
        Returns:
            List of RunInfo objects
        """
        pass
    
    @abstractmethod
    def load_run_data(self, session_id: str, run_no: int) -> Optional[TelemetryData]:
        """
        Load telemetry data for one run.
        
        Args:
            session_id: Session identifier
            run_no: Run number
        
        Returns:
            TelemetryData object or None on error
        """
        pass


# =============================================================================
# SQLite Adapter
# =============================================================================

class SQLiteAdapter(DataAdapter):
    """
    SQLite database adapter.
    
    Works with decoded_telemetry table schema.
    """
    
    def __init__(self, db_path: str):
        """
        Initialize SQLite adapter.
        
        Args:
            db_path: Path to SQLite database file
        """
        self.db_path = db_path
        self._conn: Optional[sqlite3.Connection] = None
    
    def connect(self) -> bool:
        """Open database connection."""
        try:
            self._conn = sqlite3.connect(self.db_path)
            self._conn.row_factory = sqlite3.Row
            return True
        except Exception as e:
            print(f"[SQLite ERROR] Failed to connect: {e}")
            return False
    
    def close(self):
        """Close database connection."""
        if self._conn:
            self._conn.close()
            self._conn = None
    
    def get_sessions(self) -> List[SessionInfo]:
        """Get all sessions from database."""
        if not self._conn:
            return []
        
        try:
            cursor = self._conn.cursor()
            
            cursor.execute("""
                SELECT 
                    s.session_id,
                    s.created_at,
                    s.description,
                    COUNT(DISTINCT t.run_no) as run_count
                FROM sessions s
                LEFT JOIN decoded_telemetry t ON s.session_id = t.session_id
                GROUP BY s.session_id
                ORDER BY s.created_at DESC
            """)
            
            sessions = []
            for row in cursor.fetchall():
                sessions.append(SessionInfo(
                    session_id=row['session_id'],
                    created_at=row['created_at'] or '',
                    description=row['description'] or '',
                    run_count=row['run_count'] or 0
                ))
            
            return sessions
        
        except Exception as e:
            print(f"[SQLite ERROR] Failed to get sessions: {e}")
            return []
    
    def get_runs(self, session_id: str) -> List[RunInfo]:
        """Get all runs for a session."""
        if not self._conn:
            return []
        
        try:
            cursor = self._conn.cursor()
            
            cursor.execute("""
                SELECT 
                    session_id,
                    run_name,
                    run_no,
                    COUNT(*) as sample_count,
                    MIN(t_ns) as start_time_ns,
                    MAX(t_ns) as end_time_ns
                FROM decoded_telemetry
                WHERE session_id = ?
                GROUP BY session_id, run_no, run_name
                ORDER BY run_no
            """, (session_id,))
            
            runs = []
            for row in cursor.fetchall():
                duration_ns = row['end_time_ns'] - row['start_time_ns']
                duration_ms = duration_ns / 1_000_000.0
                
                runs.append(RunInfo(
                    session_id=row['session_id'],
                    run_name=row['run_name'],
                    run_number=row['run_no'],
                    sample_count=row['sample_count'],
                    start_time_ns=row['start_time_ns'],
                    end_time_ns=row['end_time_ns'],
                    duration_ms=duration_ms
                ))
            
            return runs
        
        except Exception as e:
            print(f"[SQLite ERROR] Failed to get runs: {e}")
            return []
    
    def load_run_data(self, session_id: str, run_no: int) -> Optional[TelemetryData]:
        """Load telemetry data for one run."""
        if not self._conn:
            return None
        
        try:
            cursor = self._conn.cursor()
            
            cursor.execute("""
                SELECT 
                    run_name,
                    t_ns,
                    time_ms,
                    motor_current,
                    encoder_value,
                    relative_encoder_value,
                    v24_pec_diff,
                    pwm,
                    i2c_raw14,
                    i2c_zero_raw14,
                    i2c_delta_raw14,
                    i2c_angle_deg,
                    i2c_zero_angle_deg,
                    angular_velocity
                FROM decoded_telemetry
                WHERE session_id = ? AND run_no = ?
                ORDER BY t_ns
            """, (session_id, run_no))
            
            rows = cursor.fetchall()
            
            if not rows:
                return None
            
            # Get run_name from first row
            run_name = rows[0]['run_name']
            
            # Extract columns
            data = TelemetryData(
                session_id=session_id,
                run_name=run_name,
                run_no=run_no,
                t_ns=self._extract_column(rows, 't_ns', dtype=np.int64),
                time_ms=self._extract_column(rows, 'time_ms', dtype=np.int64),
                motor_current=self._extract_column(rows, 'motor_current', dtype=np.float32),
                encoder_value=self._extract_column(rows, 'encoder_value', dtype=np.float32),
                relative_encoder_value=self._extract_column(rows, 'relative_encoder_value', dtype=np.float32),
                v24_pec_diff=self._extract_column(rows, 'v24_pec_diff', dtype=np.float32),
                pwm=self._extract_column(rows, 'pwm', dtype=np.float32),
                i2c_raw14=self._extract_column(rows, 'i2c_raw14', dtype=np.float32),
                i2c_zero_raw14=self._extract_column(rows, 'i2c_zero_raw14', dtype=np.float32),
                i2c_delta_raw14=self._extract_column(rows, 'i2c_delta_raw14', dtype=np.float32),
                i2c_angle_deg=self._extract_column(rows, 'i2c_angle_deg', dtype=np.float32),
                i2c_zero_angle_deg=self._extract_column(rows, 'i2c_zero_angle_deg', dtype=np.float32),
                angular_velocity=self._extract_column(rows, 'angular_velocity', dtype=np.float32)
            )
            
            return data
        
        except Exception as e:
            print(f"[SQLite ERROR] Failed to load run data: {e}")
            return None
    
    def _extract_column(self, rows: List, column: str, dtype=np.float32) -> Optional[np.ndarray]:
        """
        Extract column from rows, handling NULL values.
        
        Returns None if all values are NULL.
        """
        values = [row[column] for row in rows]
        
        # Check if any non-NULL values
        if all(v is None for v in values):
            return None
        
        # Replace None with NaN
        values = [v if v is not None else np.nan for v in values]
        
        return np.array(values, dtype=dtype)


# =============================================================================
# CSV Adapter
# =============================================================================

class CSVAdapter(DataAdapter):
    """
    CSV file adapter.
    
    Expected CSV format:
    session_id,run_name,run_no,t_ns,time_ms,motor_current,encoder_value,...
    
    Single CSV file containing all sessions and runs.
    """
    
    def __init__(self, csv_path: str):
        """
        Initialize CSV adapter.
        
        Args:
            csv_path: Path to CSV file
        """
        self.csv_path = csv_path
        self._df = None  # pandas DataFrame
    
    def connect(self) -> bool:
        """Load CSV file."""
        try:
            import pandas as pd
            self._df = pd.read_csv(self.csv_path)
            
            # Validate required columns
            required = ['session_id', 'run_name', 'run_no']
            missing = [col for col in required if col not in self._df.columns]
            
            if missing:
                print(f"[CSV ERROR] Missing required columns: {missing}")
                return False
            
            return True
        
        except Exception as e:
            print(f"[CSV ERROR] Failed to load CSV: {e}")
            return False
    
    def close(self):
        """Release DataFrame."""
        self._df = None
    
    def get_sessions(self) -> List[SessionInfo]:
        """Get all sessions from CSV."""
        if self._df is None:
            return []
        
        try:
            # Group by session_id
            session_groups = self._df.groupby('session_id')
            
            sessions = []
            for session_id, group in session_groups:
                run_count = group['run_no'].nunique()
                
                # Try to get created_at if column exists
                created_at = ''
                if 'created_at' in group.columns and len(group) > 0:
                    created_at = str(group['created_at'].iloc[0])
                
                sessions.append(SessionInfo(
                    session_id=str(session_id),
                    created_at=created_at,
                    description='',  # Not in CSV
                    run_count=run_count
                ))
            
            return sessions
        
        except Exception as e:
            print(f"[CSV ERROR] Failed to get sessions: {e}")
            return []
    
    def get_runs(self, session_id: str) -> List[RunInfo]:
        """Get all runs for a session."""
        if self._df is None:
            return []
        
        try:
            # Filter by session
            session_df = self._df[self._df['session_id'] == session_id]
            
            # Group by run_no
            run_groups = session_df.groupby('run_no')
            
            runs = []
            for run_no, group in run_groups:
                sample_count = len(group)
                
                # Get run_name
                run_name = str(group['run_name'].iloc[0]) if 'run_name' in group.columns else ''
                
                # Get time info
                if 't_ns' in group.columns:
                    start_time_ns = int(group['t_ns'].min())
                    end_time_ns = int(group['t_ns'].max())
                    duration_ms = (end_time_ns - start_time_ns) / 1_000_000.0
                else:
                    start_time_ns = 0
                    end_time_ns = 0
                    duration_ms = 0.0
                
                runs.append(RunInfo(
                    session_id=session_id,
                    run_name=run_name,
                    run_number=int(run_no),
                    sample_count=sample_count,
                    start_time_ns=start_time_ns,
                    end_time_ns=end_time_ns,
                    duration_ms=duration_ms
                ))
            
            # Sort by run_number
            runs.sort(key=lambda r: r.run_number)
            
            return runs
        
        except Exception as e:
            print(f"[CSV ERROR] Failed to get runs: {e}")
            return []
    
    def load_run_data(self, session_id: str, run_no: int) -> Optional[TelemetryData]:
        """Load telemetry data for one run."""
        if self._df is None:
            return None
        
        try:
            # Filter by session and run
            run_df = self._df[
                (self._df['session_id'] == session_id) & 
                (self._df['run_no'] == run_no)
            ]
            
            if len(run_df) == 0:
                return None
            
            # Sort by timestamp
            if 't_ns' in run_df.columns:
                run_df = run_df.sort_values('t_ns')
            
            # Get run_name
            run_name = str(run_df['run_name'].iloc[0]) if 'run_name' in run_df.columns else ''
            
            # Extract columns
            data = TelemetryData(
                session_id=session_id,
                run_name=run_name,
                run_no=run_no,
                t_ns=self._extract_column_csv(run_df, 't_ns', dtype=np.int64),
                time_ms=self._extract_column_csv(run_df, 'time_ms', dtype=np.int64),
                motor_current=self._extract_column_csv(run_df, 'motor_current', dtype=np.float32),
                encoder_value=self._extract_column_csv(run_df, 'encoder_value', dtype=np.float32),
                relative_encoder_value=self._extract_column_csv(run_df, 'relative_encoder_value', dtype=np.float32),
                v24_pec_diff=self._extract_column_csv(run_df, 'v24_pec_diff', dtype=np.float32),
                pwm=self._extract_column_csv(run_df, 'pwm', dtype=np.float32),
                i2c_raw14=self._extract_column_csv(run_df, 'i2c_raw14', dtype=np.float32),
                i2c_zero_raw14=self._extract_column_csv(run_df, 'i2c_zero_raw14', dtype=np.float32),
                i2c_delta_raw14=self._extract_column_csv(run_df, 'i2c_delta_raw14', dtype=np.float32),
                i2c_angle_deg=self._extract_column_csv(run_df, 'i2c_angle_deg', dtype=np.float32),
                i2c_zero_angle_deg=self._extract_column_csv(run_df, 'i2c_zero_angle_deg', dtype=np.float32),
                angular_velocity=self._extract_column_csv(run_df, 'angular_velocity', dtype=np.float32)
            )
            
            return data
        
        except Exception as e:
            print(f"[CSV ERROR] Failed to load run data: {e}")
            return None
    
    def _extract_column_csv(self, df, column: str, dtype=np.float32) -> Optional[np.ndarray]:
        """
        Extract column from DataFrame, handling missing columns and NaN.
        
        Returns None if column doesn't exist or all values are NaN.
        """
        if column not in df.columns:
            return None
        
        values = df[column].values
        
        # Check if all NaN
        if np.all(np.isnan(values.astype(float, errors='ignore'))):
            return None
        
        return values.astype(dtype)


# =============================================================================
# Column Utilities
# =============================================================================

def get_available_columns() -> List[str]:
    """
    Get list of all plottable columns.
    
    Returns:
        List of column names
    """
    return [
        't_ns',
        'time_ms',
        'motor_current',
        'encoder_value',
        'relative_encoder_value',
        'v24_pec_diff',
        'pwm',
        'i2c_raw14',
        'i2c_zero_raw14',
        'i2c_delta_raw14',
        'i2c_angle_deg',
        'i2c_zero_angle_deg',
        'angular_velocity'
    ]


def get_column_label(column: str) -> str:
    """
    Get human-readable label for column.
    
    Args:
        column: Column name
    
    Returns:
        Human-readable label
    """
    labels = {
        't_ns': 'Time (ns)',
        'time_ms': 'Time (ms)',
        'motor_current': 'Motor Current',
        'encoder_value': 'Encoder Value',
        'relative_encoder_value': 'Relative Encoder',
        'v24_pec_diff': 'V24 PEC Diff',
        'pwm': 'PWM',
        'i2c_raw14': 'I2C Raw (14-bit)',
        'i2c_zero_raw14': 'I2C Zero Raw',
        'i2c_delta_raw14': 'I2C Delta Raw',
        'i2c_angle_deg': 'Angle (degrees)',
        'i2c_zero_angle_deg': 'Zero Angle (degrees)',
        'angular_velocity': 'Angular Velocity'
    }
    return labels.get(column, column)


def get_column_groups() -> Dict[str, List[str]]:
    """
    Get columns grouped by category.
    
    Returns:
        Dictionary of category -> list of columns
    """
    return {
        'Time': ['t_ns', 'time_ms'],
        'UART': ['motor_current', 'encoder_value', 'relative_encoder_value', 
                 'v24_pec_diff', 'pwm'],
        'I2C': ['i2c_raw14', 'i2c_zero_raw14', 'i2c_delta_raw14', 
                'i2c_angle_deg', 'i2c_zero_angle_deg'],
        'Derived': ['angular_velocity']
    }


def get_default_columns() -> List[str]:
    """
    Get default selected columns.
    
    Returns:
        List of column names to check by default
    """
    return [
        'time_ms',
        'motor_current',
        'encoder_value',
        'relative_encoder_value',
        'pwm',
        'i2c_angle_deg',
        'angular_velocity'
    ]


# =============================================================================
# Data Alignment
# =============================================================================

def align_data_to_reference(
    data_list: List[TelemetryData],
    x_column: str,
    reference_index: int = 0
) -> List[TelemetryData]:
    """
    Align multiple runs to reference run's X-axis.
    
    Useful for drift comparison when runs have different timestamps
    or sampling rates.
    
    Args:
        data_list: List of TelemetryData objects
        x_column: Column name to use as X-axis
        reference_index: Index of reference run (default: 0)
    
    Returns:
        List of TelemetryData objects aligned to reference
    """
    if reference_index >= len(data_list):
        reference_index = 0
    
    reference = data_list[reference_index]
    ref_x = getattr(reference, x_column, None)
    
    if ref_x is None:
        return data_list
    
    aligned_list = [reference]  # Reference stays as-is
    
    for idx, data in enumerate(data_list):
        if idx == reference_index:
            continue
        
        x_data = getattr(data, x_column, None)
        
        if x_data is None:
            aligned_list.append(data)
            continue
        
        # Create aligned data object
        aligned = TelemetryData(
            session_id=data.session_id,
            run_name=data.run_name,
            run_no=data.run_no
        )
        
        # Interpolate all columns to reference X
        for col in get_available_columns():
            if col == x_column:
                # X column uses reference
                setattr(aligned, col, ref_x)
            else:
                y_data = getattr(data, col, None)
                if y_data is not None and len(y_data) == len(x_data):
                    # Interpolate
                    y_aligned = np.interp(ref_x, x_data, y_data)
                    setattr(aligned, col, y_aligned)
        
        aligned_list.append(aligned)
    
    return aligned_list


# =============================================================================
# Demo
# =============================================================================

if __name__ == "__main__":
    print("Table Query - Data Access Layer")
    print("=" * 60)
    print()
    print("Available Adapters:")
    print("  • SQLiteAdapter - for .db files")
    print("  • CSVAdapter - for .csv files")
    print()
    print("Column Utilities:")
    print(f"  • {len(get_available_columns())} plottable columns")
    print(f"  • {len(get_column_groups())} categories")
    print(f"  • {len(get_default_columns())} default selections")
    print()
    print("Usage:")
    print("  from table_query import SQLiteAdapter, CSVAdapter")
    print("  adapter = SQLiteAdapter('./database/ehinge.db')")
    print("  adapter.connect()")
    print("  sessions = adapter.get_sessions()")
    print("  data = adapter.load_run_data('Session_A', 1)")