#!/usr/bin/env python3
"""
SMART Logger V3 - Uses PostgreSQL devices table for metadata.

Fetches SMART data from Scrutiny's InfluxDB, looks up device metadata from
PostgreSQL devices table, calculates BackBlaze failure probability, and logs
the enhanced data to PostgreSQL.
"""

import os
import sys
from datetime import datetime, timezone
from typing import Dict, List, Optional
import json

from influxdb_client import InfluxDBClient
import psycopg
from psycopg.rows import dict_row
from dotenv import load_dotenv

from backblaze_tables import calculate_afr, afr_to_failure_probability

load_dotenv()


class SmartLogger:
    def __init__(self):
        # InfluxDB configuration
        self.influx_url = os.getenv('INFLUXDB_URL')
        self.influx_token = os.getenv('INFLUXDB_TOKEN')
        self.influx_org = os.getenv('INFLUXDB_ORG', 'scrutiny')
        self.influx_bucket = os.getenv('INFLUXDB_BUCKET', 'metrics')

        # PostgreSQL configuration
        self.pg_host = os.getenv('POSTGRES_HOST', 'localhost')
        self.pg_port = int(os.getenv('POSTGRES_PORT', '5432'))
        self.pg_db = os.getenv('POSTGRES_DB', 'smart_monitoring')
        self.pg_user = os.getenv('POSTGRES_USER', 'postgres')
        self.pg_password = os.getenv('POSTGRES_PASSWORD')

        self.influx_client = None
        self.pg_conn = None
        self.device_metadata = {}  # WWN -> device info mapping

    def connect_influxdb(self):
        """Connect to InfluxDB."""
        print(f"Connecting to InfluxDB at {self.influx_url}...")
        self.influx_client = InfluxDBClient(
            url=self.influx_url,
            token=self.influx_token,
            org=self.influx_org
        )
        print("✓ Connected to InfluxDB")

    def connect_postgres(self):
        """Connect to PostgreSQL."""
        print(f"Connecting to PostgreSQL at {self.pg_host}:{self.pg_port}...")
        self.pg_conn = psycopg.connect(
            host=self.pg_host,
            port=self.pg_port,
            dbname=self.pg_db,
            user=self.pg_user,
            password=self.pg_password,
            row_factory=dict_row
        )
        print("✓ Connected to PostgreSQL")

    def load_device_metadata(self):
        """Load device metadata from PostgreSQL devices table."""
        print("Loading device metadata from database...")

        query = """
        SELECT device_wwn, device_path, serial_number, model, manufacturer,
               capacity_bytes, size_tb, disk_role, notes
        FROM devices
        """

        with self.pg_conn.cursor() as cur:
            cur.execute(query)
            for row in cur.fetchall():
                wwn = row['device_wwn']
                self.device_metadata[wwn] = dict(row)

        print(f"✓ Loaded metadata for {len(self.device_metadata)} device(s)")

        if not self.device_metadata:
            print("\n⚠ WARNING: No devices found in the 'devices' table!")
            print("  Please populate the devices table using populate_devices.sql")
            print("  Run: psql -h <host> -U <user> -d <db> -f populate_devices.sql\n")

    def fetch_smart_data_from_influx(self) -> Dict[str, Dict]:
        """
        Fetch latest SMART data from InfluxDB for all devices.

        Returns:
            Dict mapping WWN to device data
        """
        query_api = self.influx_client.query_api()

        # Query to get latest data for each device
        # Include attribute 193 for SnapRAID pre-v13.0 compatibility
        query = f'''
        from(bucket: "{self.influx_bucket}")
            |> range(start: -1h)
            |> filter(fn: (r) => r._measurement == "smart")
            |> filter(fn: (r) =>
                r._field == "temp" or
                r._field == "power_on_hours" or
                r._field =~ /failure_rate$/ or
                r._field == "attr.5.raw_value" or
                r._field == "attr.187.raw_value" or
                r._field == "attr.188.raw_value" or
                r._field == "attr.193.raw_value" or
                r._field == "attr.197.raw_value" or
                r._field == "attr.198.raw_value" or
                r._field == "attr.9.raw_value"
            )
            |> last()
        '''

        print("Querying InfluxDB for SMART data...")
        result = query_api.query(query)

        devices_data = {}

        for table in result:
            for record in table.records:
                wwn = record.values.get('device_wwn')
                if not wwn:
                    continue

                if wwn not in devices_data:
                    devices_data[wwn] = {
                        'timestamp': record.get_time(),
                        'wwn': wwn,
                        'protocol': record.values.get('device_protocol'),
                        'smart_attributes': {},
                        'failure_rates': {},  # Scrutiny's pre-calculated rates
                        'temp': None,
                        'power_on_hours': None
                    }

                field = record.get_field()
                value = record.get_value()

                # Map fields to device data
                if field == 'temp':
                    devices_data[wwn]['temp'] = int(value) if value else None
                elif field == 'power_on_hours':
                    devices_data[wwn]['power_on_hours'] = int(value) if value else None
                elif field.startswith('attr.'):
                    # Extract attribute ID from field name like "attr.5.raw_value"
                    parts = field.split('.')
                    if len(parts) >= 2:
                        # Try to parse as integer, skip if not numeric (e.g., NVMe attributes)
                        try:
                            attr_id = int(parts[1])
                        except ValueError:
                            continue  # Skip non-numeric attribute IDs (NVMe, etc.)

                        # Collect raw_value for logging
                        if parts[-1] == 'raw_value':
                            raw_val = int(value) if value else 0
                            # Apply correct bit mask per attribute (like SnapRAID does)
                            # 16-bit mask for: 187 (Uncorrectable), 188 (Timeout)
                            # 32-bit mask for: 5 (Reallocated), 193 (Load Cycle), 197 (Pending), 198 (Offline)
                            if attr_id in [187, 188]:
                                devices_data[wwn]['smart_attributes'][attr_id] = raw_val & 0xFFFF
                            else:
                                devices_data[wwn]['smart_attributes'][attr_id] = raw_val & 0xFFFFFFFF
                        # Collect Scrutiny's failure_rate for each attribute
                        elif parts[-1] == 'failure_rate':
                            devices_data[wwn]['failure_rates'][attr_id] = float(value) if value else 0.0

        print(f"✓ Found data for {len(devices_data)} device(s)")
        return devices_data

    def process_device_data(self, devices_data: Dict[str, Dict]) -> List[Dict]:
        """
        Process device data and calculate failure probabilities.

        Args:
            devices_data: Dict mapping WWN to device data from InfluxDB

        Returns:
            Processed device data with calculated AFR
        """
        processed = []

        for wwn, data in devices_data.items():
            # Get device metadata from database
            metadata = self.device_metadata.get(wwn, {})

            if not metadata:
                print(f"\n⚠ Warning: No metadata found for device {wwn}")
                print(f"  Add this device to the devices table using populate_devices.sql")
                continue

            # Get SMART attributes
            smart_attrs = data.get('smart_attributes', {})

            # Determine if Seagate (for attribute 188 handling)
            model = metadata.get('model', '').lower()
            manufacturer = metadata.get('manufacturer', '').lower()
            is_seagate = 'seagate' in model or 'seagate' in manufacturer

            # Calculate failure probability BOTH ways for trending
            # v12 (with attr 193) - matches your current SnapRAID
            afr_v12 = calculate_afr(smart_attrs, is_seagate=is_seagate, include_193=True)
            prob_v12 = afr_to_failure_probability(afr_v12)
            failure_probability_v12_pct = prob_v12 * 100

            # v13 (without attr 193) - modern SnapRAID v13.0+
            afr_v13 = calculate_afr(smart_attrs, is_seagate=is_seagate, include_193=False)
            prob_v13 = afr_to_failure_probability(afr_v13)
            failure_probability_v13_pct = prob_v13 * 100

            # Calculate power-on days
            power_on_hours = data.get('power_on_hours')
            power_on_days = power_on_hours // 24 if power_on_hours else None

            processed_device = {
                'timestamp': data.get('timestamp', datetime.now(timezone.utc)),
                'device_wwn': wwn,
                'temperature_celsius': data.get('temp'),
                'power_on_days': power_on_days,
                'error_count': smart_attrs.get(199, smart_attrs.get(187, 0)),
                'failure_probability_v12_pct': round(failure_probability_v12_pct, 2),
                'failure_probability_v13_pct': round(failure_probability_v13_pct, 2),
                'smart_attributes': json.dumps(smart_attrs)
            }

            processed.append(processed_device)

            # Print summary with metadata
            print(f"\n{metadata.get('device_path', wwn)}")
            print(f"  WWN: {wwn}")
            print(f"  Serial: {metadata.get('serial_number', 'unknown')}")
            print(f"  Model: {metadata.get('model', 'unknown')}")
            print(f"  Role: {metadata.get('disk_role', '-')}")
            print(f"  Temp: {processed_device['temperature_celsius']}°C")
            print(f"  Power-On: {processed_device['power_on_days']} days")
            print(f"  Size: {metadata.get('size_tb', '?')} TB")
            print(f"  SMART Attrs: {list(smart_attrs.keys())}")
            print(f"  Failure Probability (v12 w/193): {processed_device['failure_probability_v12_pct']:.2f}%")
            print(f"  Failure Probability (v13 no193): {processed_device['failure_probability_v13_pct']:.2f}%")

        return processed

    def save_to_postgres(self, devices: List[Dict]):
        """
        Save processed device data to PostgreSQL.

        Args:
            devices: Processed device data
        """
        if not devices:
            print("\nNo devices to save")
            return

        print(f"\nSaving {len(devices)} device record(s) to PostgreSQL...")

        insert_query = '''
        INSERT INTO smart_metrics (
            timestamp, device_wwn,
            temperature_celsius, power_on_days, error_count,
            failure_probability_v12_pct, failure_probability_v13_pct, smart_attributes
        ) VALUES (
            %(timestamp)s, %(device_wwn)s,
            %(temperature_celsius)s, %(power_on_days)s, %(error_count)s,
            %(failure_probability_v12_pct)s, %(failure_probability_v13_pct)s, %(smart_attributes)s::jsonb
        )
        ON CONFLICT (device_wwn, timestamp) DO UPDATE SET
            temperature_celsius = EXCLUDED.temperature_celsius,
            power_on_days = EXCLUDED.power_on_days,
            error_count = EXCLUDED.error_count,
            failure_probability_v12_pct = EXCLUDED.failure_probability_v12_pct,
            failure_probability_v13_pct = EXCLUDED.failure_probability_v13_pct,
            smart_attributes = EXCLUDED.smart_attributes
        '''

        with self.pg_conn.cursor() as cur:
            for device in devices:
                cur.execute(insert_query, device)

        self.pg_conn.commit()
        print("✓ Data saved to PostgreSQL")

    def print_summary(self):
        """Print a summary with both v12 and v13 failure calculations."""
        print("\n" + "=" * 95)
        print("SMART Summary (v12=with attr 193, v13=without attr 193)")
        print("=" * 95)
        print(f"{'Temp':>5} {'Power':>6} {'Error':>7} {'v12':>4} {'v13':>4} {'Size':>5}")
        print(f"{'C':>5} {'OnDays':>6} {'Count':>7} {'FP':>4} {'FP':>4} {'TB':>5}  "
              f"{'Serial':<16} {'Device':<16} {'Disk':<10}")
        print("-" * 95)

        query = """
        SELECT
            device_path,
            serial_number,
            temperature_celsius,
            power_on_days,
            error_count,
            failure_probability_v12_pct,
            failure_probability_v13_pct,
            size_tb,
            disk_role
        FROM smart_summary
        ORDER BY
            CASE disk_role
                WHEN 'parity' THEN 1
                WHEN '2-parity' THEN 2
                ELSE 3
            END,
            device_path
        """

        with self.pg_conn.cursor() as cur:
            cur.execute(query)
            for row in cur.fetchall():
                temp = row['temperature_celsius'] or 0
                days = row['power_on_days'] or 0
                errors = row['error_count'] or 0
                fp_v12 = int(row['failure_probability_v12_pct'] or 0)
                fp_v13 = int(row['failure_probability_v13_pct'] or 0)
                size = row['size_tb'] or 0
                serial = (row['serial_number'] or '')[:16]
                device = (row['device_path'] or '')[:16]
                role = row['disk_role'] or '-'

                print(f"{temp:5d} {days:6d} {errors:7d} {fp_v12:3d}% {fp_v13:3d}% {size:5.1f}  "
                      f"{serial:<16} {device:<16} {role:<10}")

        print()

    def run(self):
        """Main execution flow."""
        try:
            self.connect_influxdb()
            self.connect_postgres()

            # Load device metadata from PostgreSQL
            self.load_device_metadata()

            # Fetch SMART data from InfluxDB
            devices_data = self.fetch_smart_data_from_influx()

            # Process and calculate failure probabilities
            processed = self.process_device_data(devices_data)

            # Save to PostgreSQL
            self.save_to_postgres(processed)

            # Print SnapRAID-style summary
            self.print_summary()

        except Exception as e:
            print(f"ERROR: {e}")
            import traceback
            traceback.print_exc()
            sys.exit(1)

        finally:
            if self.influx_client:
                self.influx_client.close()
            if self.pg_conn:
                self.pg_conn.close()


def main():
    # Validate environment variables
    required_vars = ['INFLUXDB_URL', 'INFLUXDB_TOKEN', 'POSTGRES_PASSWORD']
    missing = [var for var in required_vars if not os.getenv(var)]

    if missing:
        print(f"ERROR: Missing required environment variables: {', '.join(missing)}")
        print("Please create a .env file based on .env.example")
        sys.exit(1)

    logger = SmartLogger()
    logger.run()


if __name__ == '__main__':
    main()