core/db_connector.py

"""
Enhanced Database Connector for CSH2 Web Dashboard
Smart table routing, connection health checks, pivot data support
"""
import psycopg2
from psycopg2.extras import RealDictCursor
import pandas as pd
from typing import List, Dict, Optional, Tuple
from datetime import datetime, timezone
import streamlit as st
from core import config


class DatabaseConnector:
    """Handles all database operations with smart table routing"""

    def __init__(self):
        self.conn = None
        self.tag_cache = {}
        self.index_to_tag = {}
        self._connect()
        self._load_tag_cache()

    def _connect(self):
        try:
            self.conn = psycopg2.connect(**config.DB_CONFIG)
            self.conn.autocommit = True
        except Exception as e:
            raise ConnectionError(f"Database connection failed: {e}")

    def _ensure_connection(self):
        """Reconnect if the connection has gone stale"""
        try:
            with self.conn.cursor() as cur:
                cur.execute("SELECT 1")
        except Exception:
            self._connect()

    def _load_tag_cache(self):
        try:
            query = f"SELECT tagindex, tagname FROM {config.TABLES['tag_metadata']}"
            df = pd.read_sql(query, self.conn)
            self.tag_cache = {row['tagname']: row['tagindex'] for _, row in df.iterrows()}
            self.index_to_tag = {row['tagindex']: row['tagname'] for _, row in df.iterrows()}
        except Exception as e:
            self.tag_cache = {}
            self.index_to_tag = {}

    def get_tag_index(self, tag_name: str) -> Optional[int]:
        return self.tag_cache.get(tag_name) or self.tag_cache.get(tag_name.upper())

    def get_tag_name(self, tag_index: int) -> Optional[str]:
        return self.index_to_tag.get(tag_index)

    def get_all_tag_names(self) -> List[str]:
        return sorted(self.tag_cache.keys())

    def _select_table(self, start_time: datetime, end_time: datetime) -> str:
        """Select the optimal table based on time range duration"""
        duration_seconds = (end_time - start_time).total_seconds()
        routing = config.TABLE_ROUTING

        if duration_seconds <= routing['raw_max_seconds']:
            return config.TABLES['raw']
        elif duration_seconds <= routing['agg_1sec_max_seconds']:
            return config.TABLES['agg_1sec']
        else:
            return config.TABLES['agg_15sec']

    def get_sensor_data(
        self,
        tag_names: List[str],
        start_time: datetime,
        end_time: datetime,
        table_override: Optional[str] = None,
        limit: Optional[int] = None,
    ) -> pd.DataFrame:
        """
        Fetch sensor data with smart table routing.

        Returns DataFrame with columns: timestamp, tag_name, value
        """
        self._ensure_connection()

        tag_indices = []
        valid_tags = []
        for tag in tag_names:
            idx = self.get_tag_index(tag)
            if idx is not None:
                tag_indices.append(idx)
                valid_tags.append(tag)

        if not tag_indices:
            return pd.DataFrame(columns=['timestamp', 'tag_name', 'value'])

        table = table_override or self._select_table(start_time, end_time)
        tag_list = ', '.join(map(str, tag_indices))
        limit_clause = f"LIMIT {limit}" if limit else f"LIMIT {config.MAX_ROWS_RETURN}"

        query = f"""
        SELECT
            utc_full_timestamp as timestamp,
            tagindex,
            val as value
        FROM {table}
        WHERE tagindex IN ({tag_list})
          AND utc_full_timestamp >= %s
          AND utc_full_timestamp <= %s
        ORDER BY utc_full_timestamp, tagindex
        {limit_clause}
        """

        try:
            df = pd.read_sql(query, self.conn, params=(start_time, end_time))
            df['tag_name'] = df['tagindex'].map(self.index_to_tag)
            df = df[['timestamp', 'tag_name', 'value']]
            return df
        except Exception as e:
            st.error(f"Query failed: {e}")
            return pd.DataFrame(columns=['timestamp', 'tag_name', 'value'])

    def get_pivot_data(
        self,
        tag_names: List[str],
        start_time: datetime,
        end_time: datetime,
    ) -> pd.DataFrame:
        """Fetch data and return in wide (pivoted) format"""
        df = self.get_sensor_data(tag_names, start_time, end_time)
        if df.empty:
            return df

        pivot_df = df.pivot_table(
            index='timestamp',
            columns='tag_name',
            values='value',
        ).reset_index()

        pivot_df.columns.name = None
        return pivot_df

    def get_available_months(self) -> List[str]:
        """Get distinct months where data exists (using 15sec table for speed)"""
        self._ensure_connection()
        query = f"""
        SELECT DISTINCT date_trunc('month', utc_full_timestamp)::date as month
        FROM {config.TABLES['agg_15sec']}
        ORDER BY month
        """
        try:
            df = pd.read_sql(query, self.conn)
            return [row['month'].strftime('%Y-%m') for _, row in df.iterrows()]
        except Exception:
            # Fallback: use raw table with a limit
            try:
                query = f"""
                SELECT DISTINCT date_trunc('month', utc_full_timestamp)::date as month
                FROM {config.TABLES['raw']}
                ORDER BY month
                """
                df = pd.read_sql(query, self.conn)
                return [row['month'].strftime('%Y-%m') for _, row in df.iterrows()]
            except Exception:
                return []

    def get_time_range(self) -> Tuple[Optional[datetime], Optional[datetime]]:
        """Get overall time range of available data"""
        self._ensure_connection()
        query = f"""
        SELECT MIN(utc_full_timestamp) as earliest, MAX(utc_full_timestamp) as latest
        FROM {config.TABLES['agg_15sec']}
        """
        try:
            df = pd.read_sql(query, self.conn)
            return df['earliest'][0], df['latest'][0]
        except Exception:
            return None, None

    def get_data_summary(
        self,
        tag_name: str,
        start_time: datetime,
        end_time: datetime,
    ) -> Dict:
        """Get statistical summary for a single tag"""
        self._ensure_connection()
        tag_index = self.get_tag_index(tag_name)
        if tag_index is None:
            return {'error': f'Tag {tag_name} not found'}

        table = self._select_table(start_time, end_time)
        query = f"""
        SELECT COUNT(*) as count, MIN(val) as min_value, MAX(val) as max_value,
               AVG(val) as avg_value, STDDEV(val) as std_value
        FROM {table}
        WHERE tagindex = %s AND utc_full_timestamp >= %s AND utc_full_timestamp <= %s
        """
        try:
            with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute(query, (tag_index, start_time, end_time))
                result = cur.fetchone()
                return dict(result)
        except Exception as e:
            return {'error': str(e)}

    def check_table_exists(self, table_name: str) -> bool:
        """Check if a table exists in the database"""
        self._ensure_connection()
        query = """
        SELECT EXISTS (
            SELECT FROM information_schema.tables
            WHERE table_schema = 'public' AND table_name = %s
        )
        """
        try:
            with self.conn.cursor() as cur:
                cur.execute(query, (table_name,))
                return cur.fetchone()[0]
        except Exception:
            return False

    def get_cooldown_periods(self, start_time: datetime, end_time: datetime) -> pd.DataFrame:
        """Query cooldown_periods table if it exists"""
        self._ensure_connection()
        if not self.check_table_exists('cooldown_periods'):
            return pd.DataFrame()

        query = """
        SELECT cooldown_start, cooldown_end, start_temperature, end_temperature, duration_sec
        FROM cooldown_periods
        WHERE cooldown_start >= %s AND cooldown_start <= %s
        ORDER BY cooldown_start
        """
        try:
            return pd.read_sql(query, self.conn, params=(start_time, end_time))
        except Exception:
            return pd.DataFrame()

    def resolve_motor_speed_tag(self) -> Optional[str]:
        """Find which motor speed tag exists in the database"""
        for tag in config.MOTOR_SPEED_TAGS:
            if tag in self.tag_cache:
                return tag
        # Search for any tag containing 'speed' or 'rpm'
        for tag_name in self.tag_cache:
            if 'speed' in tag_name.lower() or 'rpm' in tag_name.lower():
                return tag_name
        return None

    def get_cycle_periods(self, month_str: str = None) -> List[Dict]:
        """Fetch pre-computed cycles from the cycle_periods table.

        Args:
            month_str: Optional "YYYY-MM" filter. None returns all cycles.

        Returns:
            List of cycle dicts in the format expected by downstream pages
            (cycle_id, start_time, end_time, duration_minutes, peak_speed,
            peak_pressure, avg_flow, etc.), plus new fields from the DB.
        """
        self._ensure_connection()

        query = """
        SELECT id, cycle_start, cycle_end, duration_sec,
               max_motor_speed, max_flow_rate, avg_flow_rate,
               total_kg_dispensed, min_discharge_pressure, max_discharge_pressure,
               min_fill_temperature, max_fill_temperature, avg_ambient_temperature,
               cryotank_start_pressure, cryotank_end_pressure, total_pump_strokes
        FROM cycle_periods
        """
        params = []
        if month_str:
            year, month = map(int, month_str.split('-'))
            if month == 12:
                next_year, next_month = year + 1, 1
            else:
                next_year, next_month = year, month + 1
            query += " WHERE cycle_start >= %s AND cycle_start < %s"
            params = [
                datetime(year, month, 1, tzinfo=timezone.utc),
                datetime(next_year, next_month, 1, tzinfo=timezone.utc),
            ]
        query += " ORDER BY cycle_start"

        try:
            df = pd.read_sql(query, self.conn, params=params or None)
        except Exception:
            return []

        if df.empty:
            return []

        from core.timezone import cycle_periods_to_utc, cycle_periods_to_eastern

        speed_mult = getattr(config, 'MOTOR_SPEED_DISPLAY_MULTIPLIER', 1.0)

        cycles = []
        for _, row in df.iterrows():
            raw_start = row['cycle_start'].to_pydatetime() if hasattr(row['cycle_start'], 'to_pydatetime') else row['cycle_start']
            raw_end = row['cycle_end'].to_pydatetime() if hasattr(row['cycle_end'], 'to_pydatetime') else row['cycle_end']

            cycles.append({
                # Times corrected to real UTC — use these for all DB queries
                'start_time': cycle_periods_to_utc(raw_start),
                'end_time': cycle_periods_to_utc(raw_end),
                # Eastern display times — use these for all UI rendering
                'start_time_et': cycle_periods_to_eastern(raw_start),
                'end_time_et': cycle_periods_to_eastern(raw_end),
                # Standard fields
                'cycle_id': int(row['id']),
                'duration_minutes': float(row['duration_sec']) / 60.0,
                'peak_speed': float(row['max_motor_speed'] or 0) * speed_mult,
                'peak_pressure': float(row['max_discharge_pressure'] or 0),
                'initial_pressure': float(row['min_discharge_pressure'] or 0),
                'avg_flow': float(row['avg_flow_rate'] or 0),
                'peak_TT110': float(row['max_fill_temperature'] or 0),
                # Additional fields from cycle_periods table
                'max_flow_rate': float(row['max_flow_rate'] or 0),
                'total_kg_dispensed': float(row['total_kg_dispensed'] or 0),
                'min_fill_temperature': float(row['min_fill_temperature'] or 0),
                'avg_ambient_temperature': float(row['avg_ambient_temperature'] or 0),
                'cryotank_start_pressure': float(row['cryotank_start_pressure'] or 0),
                'cryotank_end_pressure': float(row['cryotank_end_pressure'] or 0),
                'total_pump_strokes': float(row['total_pump_strokes'] or 0),
            })
        return cycles

    def get_cycle_months(self) -> List[str]:
        """Get distinct months where cycles exist in cycle_periods table."""
        self._ensure_connection()
        query = """
        SELECT DISTINCT date_trunc('month', cycle_start)::date as month
        FROM cycle_periods
        ORDER BY month
        """
        try:
            df = pd.read_sql(query, self.conn)
            return [row['month'].strftime('%Y-%m') for _, row in df.iterrows()]
        except Exception:
            return []

    # ── Text Search Methods ──

    def search_memo_log(
        self,
        query: str,
        limit: int = 20,
        engineer: Optional[str] = None,
        severity: Optional[str] = None,
    ) -> pd.DataFrame:
        """Full-text ranked search across engineer field notes in memo_log.

        Uses PostgreSQL tsvector/tsquery with ts_rank for BM25-style relevance.
        Searches across summary, issues_found, raw_transcript, maintenance_done,
        and system_performance columns.
        """
        self._ensure_connection()
        sql = """
        SELECT id, logged_at, engineer, severity, activity_type,
               summary, issues_found, maintenance_done, system_performance,
               components_affected,
               ts_rank(
                   to_tsvector('english',
                       coalesce(summary, '') || ' ' ||
                       coalesce(issues_found, '') || ' ' ||
                       coalesce(raw_transcript, '') || ' ' ||
                       coalesce(maintenance_done, '') || ' ' ||
                       coalesce(system_performance, '')
                   ),
                   plainto_tsquery('english', %(query)s)
               ) AS rank
        FROM memo_log
        WHERE to_tsvector('english',
                  coalesce(summary, '') || ' ' ||
                  coalesce(issues_found, '') || ' ' ||
                  coalesce(raw_transcript, '') || ' ' ||
                  coalesce(maintenance_done, '') || ' ' ||
                  coalesce(system_performance, '')
              ) @@ plainto_tsquery('english', %(query)s)
        """
        params: Dict = {'query': query}
        if engineer:
            sql += " AND engineer = %(engineer)s"
            params['engineer'] = engineer
        if severity:
            sql += " AND severity = %(severity)s"
            params['severity'] = severity
        sql += " ORDER BY rank DESC LIMIT %(limit)s"
        params['limit'] = limit

        try:
            return pd.read_sql(sql, self.conn, params=params)
        except Exception as e:
            st.error(f"Memo search failed: {e}")
            return pd.DataFrame()

    def search_events(
        self,
        query: str,
        condition: Optional[str] = None,
        limit: int = 50,
    ) -> pd.DataFrame:
        """Search PLC alarm events using both full-text and pattern matching.

        Alarm messages contain sensor tags in bracket notation (e.g. [AI_PT110_Alarm]),
        so we use ILIKE alongside tsvector to catch tag name patterns.
        """
        self._ensure_connection()
        # Use ILIKE for pattern matching (catches sensor names in brackets)
        # plus tsvector for any natural text content
        sql = """
        SELECT eventtimestamp, conditionname, message, inputvalue, limitvalue
        FROM allevent
        WHERE (
            message ILIKE %(pattern)s
            OR conditionname ILIKE %(pattern)s
        )
        """
        params: Dict = {'query': query, 'pattern': f'%{query}%', 'limit': limit}
        if condition:
            sql += " AND conditionname = %(condition)s"
            params['condition'] = condition
        sql += " ORDER BY eventtimestamp DESC LIMIT %(limit)s"
        try:
            return pd.read_sql(sql, self.conn, params=params)
        except Exception as e:
            st.error(f"Event search failed: {e}")
            return pd.DataFrame()

    def get_event_conditions(self) -> List[str]:
        """Get distinct condition types from allevent for filter dropdown."""
        self._ensure_connection()
        try:
            with self.conn.cursor() as cur:
                cur.execute(
                    "SELECT conditionname, COUNT(*) as cnt FROM allevent "
                    "GROUP BY conditionname ORDER BY cnt DESC"
                )
                return [row[0] for row in cur.fetchall() if row[0]]
        except Exception:
            return []

    def search_tags_fuzzy(self, query: str, limit: int = 15) -> pd.DataFrame:
        """Fuzzy tag/sensor search using pg_trgm similarity on tagimport."""
        self._ensure_connection()
        sql = """
        SELECT tagname, tagdesc, tagunit,
               greatest(
                   similarity(lower(tagname), lower(%(query)s)),
                   similarity(lower(coalesce(tagdesc, '')), lower(%(query)s))
               ) AS sim
        FROM tagimport
        WHERE lower(tagname) %% lower(%(query)s)
           OR lower(coalesce(tagdesc, '')) %% lower(%(query)s)
        ORDER BY sim DESC
        LIMIT %(limit)s
        """
        try:
            return pd.read_sql(sql, self.conn, params={'query': query, 'limit': limit})
        except Exception as e:
            st.error(f"Tag search failed: {e}")
            return pd.DataFrame()

    def get_memo_engineers(self) -> List[str]:
        """Get distinct engineers from memo_log for filter dropdown."""
        self._ensure_connection()
        try:
            with self.conn.cursor() as cur:
                cur.execute("SELECT DISTINCT engineer FROM memo_log WHERE engineer IS NOT NULL ORDER BY engineer")
                return [row[0] for row in cur.fetchall()]
        except Exception:
            return []

    def get_memo_severities(self) -> List[str]:
        """Get distinct severity levels from memo_log for filter dropdown."""
        self._ensure_connection()
        try:
            with self.conn.cursor() as cur:
                cur.execute("SELECT DISTINCT severity FROM memo_log WHERE severity IS NOT NULL ORDER BY severity")
                return [row[0] for row in cur.fetchall()]
        except Exception:
            return []

    def close(self):
        if self.conn:
            self.conn.close()


@st.cache_resource
def get_db_connector() -> DatabaseConnector:
    """Create a singleton database connector"""
    return DatabaseConnector()