#!/usr/bin/env python3
"""
System Resource Monitor for Featrix Sphere
Tracks GPU, CPU, and memory usage every 15-30 seconds
Correlates resource usage with job completion events
"""

import json
import logging
import os
import sqlite3
import time
import re
import threading
import linecache
import hashlib
import socket
import urllib.request
import urllib.parse
from datetime import datetime, timezone, timedelta
from pathlib import Path
from typing import Dict, List, Optional, Set

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)-8s] %(name)-45s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

class TracebackMonitor:
    """Monitor log files for Python tracebacks and send to monitoring endpoint."""
    
    def __init__(self, watch_dirs: List[str], monitor_url: str = "https://monitor.featrix.com/traceback"):
        self.watch_dirs = [Path(d) for d in watch_dirs]
        self.monitor_url = monitor_url
        self.file_sizes: Dict[Path, int] = {}  # Track file sizes
        self.seen_traceback_ids: Set[str] = set()  # Track traceback IDs we've already sent
        self.hostname = socket.gethostname()
        self.running = False
        self.thread = None
        
        # Ensure watch directories exist
        for watch_dir in self.watch_dirs:
            watch_dir.mkdir(parents=True, exist_ok=True)
            logger.info(f"📁 Watching log directory: {watch_dir}")
    
    def _get_code_context(self, filename: str, lineno: int, radius: int = 3) -> List[Dict]:
        """Get code context around a line number."""
        context = []
        for offset in range(-radius, radius + 1):
            lnum = lineno + offset
            if lnum <= 0:
                continue
            line = linecache.getline(filename, lnum)
            if not line:
                continue
            context.append({
                "lineno": lnum,
                "is_error_line": (offset == 0),
                "code": line.rstrip("\n"),
            })
        return context
    
    def _enhance_traceback(self, traceback_data: Dict) -> Dict:
        """Enhance traceback with source code context for /sphere/app/* files."""
        enhanced_frames = []
        
        for frame in traceback_data.get("frames", []):
            filename = frame.get("filename", "")
            lineno = frame.get("lineno", 0)
            
            # Only enhance frames from /sphere/app/* files
            if "/sphere/app/" in filename and lineno > 0:
                # Check if code_context already exists
                if not frame.get("code_context"):
                    try:
                        code_context = self._get_code_context(filename, lineno)
                        frame["code_context"] = code_context
                    except Exception as e:
                        logger.debug(f"Could not get code context for {filename}:{lineno}: {e}")
            
            enhanced_frames.append(frame)
        
        traceback_data["frames"] = enhanced_frames
        return traceback_data
    
    def _send_traceback(self, traceback_data: Dict, log_file: Path, process_name: str = None):
        """Send traceback to monitoring endpoint."""
        traceback_id = traceback_data.get("__id__")
        
        # Skip if we've already sent this traceback
        if traceback_id and traceback_id in self.seen_traceback_ids:
            return
        
        # Enhance with source code context
        enhanced_tb = self._enhance_traceback(traceback_data.copy())
        
        # Add monitoring metadata
        enhanced_tb["__monitor_metadata__"] = {
            "log_file": str(log_file),
            "hostname": self.hostname,
            "process": process_name or self._extract_process_name(log_file),
            "detected_at": datetime.now(timezone.utc).isoformat(),
        }
        
        try:
            # Send to monitoring endpoint
            data = json.dumps(enhanced_tb).encode('utf-8')
            req = urllib.request.Request(
                self.monitor_url,
                data=data,
                headers={'Content-Type': 'application/json'},
                method='POST'
            )
            
            with urllib.request.urlopen(req, timeout=10) as response:
                if response.status == 200:
                    logger.info(f"✅ Sent traceback {traceback_id[:8]} to monitor")
                    if traceback_id:
                        self.seen_traceback_ids.add(traceback_id)
                else:
                    logger.warning(f"⚠️  Monitor returned status {response.status} for traceback {traceback_id[:8]}")
                    
        except urllib.error.HTTPError as e:
            logger.error(f"❌ HTTP error sending traceback {traceback_id[:8]}: {e.code} {e.reason}")
        except urllib.error.URLError as e:
            logger.error(f"❌ URL error sending traceback {traceback_id[:8]}: {e}")
        except Exception as e:
            logger.error(f"❌ Error sending traceback {traceback_id[:8]}: {e}")
    
    def _extract_process_name(self, log_file: Path) -> str:
        """Extract process name from log file path."""
        # Try to extract from path patterns like:
        # /var/log/featrix/app.log -> app
        # /sphere/app/featrix_output/job_*/logs/stdout.log -> job_*
        try:
            if "featrix_output" in str(log_file):
                # Extract job name from path
                parts = log_file.parts
                for i, part in enumerate(parts):
                    if part == "featrix_output" and i + 1 < len(parts):
                        return parts[i + 1]
            elif log_file.name.endswith(".log"):
                return log_file.stem
        except Exception:
            pass
        return "unknown"
    
    def _parse_traceback_from_text(self, text: str) -> List[Dict]:
        """Parse JSON tracebacks from log text. Returns list of traceback dicts."""
        tracebacks = []
        
        # Look for JSON objects that match the traceback format
        # Pattern: {"__comment__": "FEATRIX TRACEBACK", ...}
        
        # Strategy: Find all potential JSON traceback starts
        lines = text.split('\n')
        
        for i, line in enumerate(lines):
            # Check if this line contains the traceback marker
            if '"__comment__": "FEATRIX TRACEBACK"' in line or '{"__comment__"' in line:
                # Try to extract JSON starting from this line
                json_candidates = []
                
                # Try single line first (most common case)
                try:
                    traceback_data = json.loads(line.strip())
                    if traceback_data.get("__comment__") == "FEATRIX TRACEBACK":
                        tracebacks.append(traceback_data)
                        continue
                except json.JSONDecodeError:
                    pass
                
                # Try multi-line JSON - collect lines until we have balanced braces
                json_lines = [line]
                brace_count = line.count('{') - line.count('}')
                
                for j in range(i + 1, min(i + 200, len(lines))):  # Limit to 200 lines
                    next_line = lines[j]
                    json_lines.append(next_line)
                    brace_count += next_line.count('{') - next_line.count('}')
                    
                    if brace_count == 0:
                        # Balanced braces - try to parse
                        json_text = '\n'.join(json_lines)
                        try:
                            traceback_data = json.loads(json_text)
                            if traceback_data.get("__comment__") == "FEATRIX TRACEBACK":
                                tracebacks.append(traceback_data)
                                break
                        except json.JSONDecodeError:
                            pass
                
                # Also try to find JSON in "featrix_trace" extra field format
                # Pattern: "featrix_trace": {...}
                if '"featrix_trace"' in line or "'featrix_trace'" in line:
                    # Look for the JSON object after the key
                    match = re.search(r'featrix_trace["\']?\s*:\s*(\{.*\})', line, re.DOTALL)
                    if match:
                        try:
                            traceback_data = json.loads(match.group(1))
                            if traceback_data.get("__comment__") == "FEATRIX TRACEBACK":
                                tracebacks.append(traceback_data)
                        except json.JSONDecodeError:
                            pass
        
        return tracebacks
    
    def _process_log_file(self, log_file: Path):
        """Process a log file for new tracebacks."""
        try:
            current_size = log_file.stat().st_size
            
            # Check if file was rotated (size decreased)
            if log_file in self.file_sizes and current_size < self.file_sizes[log_file]:
                logger.info(f"🔄 Log file rotated: {log_file} (size: {self.file_sizes[log_file]} -> {current_size})")
                # Read last 128KB of rotated file
                read_size = min(128 * 1024, current_size)
                offset = max(0, current_size - read_size)
            elif log_file not in self.file_sizes:
                # New file - read last 128KB
                read_size = min(128 * 1024, current_size)
                offset = max(0, current_size - read_size)
            else:
                # File grew - read delta
                old_size = self.file_sizes[log_file]
                read_size = current_size - old_size
                offset = old_size
            
            if read_size <= 0:
                self.file_sizes[log_file] = current_size
                return
            
            # Read the new content
            with open(log_file, 'rb') as f:
                f.seek(offset)
                content = f.read(read_size).decode('utf-8', errors='ignore')
            
            # Parse tracebacks from content
            tracebacks = self._parse_traceback_from_text(content)
            
            for traceback_data in tracebacks:
                logger.info(f"🔍 Found traceback in {log_file}: {traceback_data.get('exception_type')} - {traceback_data.get('exception_message', '')[:50]}")
                self._send_traceback(traceback_data, log_file)
            
            # Update file size
            self.file_sizes[log_file] = current_size
            
        except FileNotFoundError:
            # File was deleted
            if log_file in self.file_sizes:
                del self.file_sizes[log_file]
        except Exception as e:
            logger.error(f"Error processing log file {log_file}: {e}")
    
    def _scan_log_files(self):
        """Scan watch directories for .log files."""
        log_files = []
        
        for watch_dir in self.watch_dirs:
            if not watch_dir.exists():
                continue
            
            # Find all .log files recursively
            for log_file in watch_dir.rglob("*.log"):
                if log_file.is_file():
                    log_files.append(log_file)
        
        return log_files
    
    def _monitor_loop(self):
        """Main monitoring loop."""
        logger.info("🔍 Starting traceback monitor...")
        
        while self.running:
            try:
                # Scan for log files
                log_files = self._scan_log_files()
                
                # Process each log file
                for log_file in log_files:
                    self._process_log_file(log_file)
                
                # Sleep before next scan
                time.sleep(5)  # Check every 5 seconds
                
            except Exception as e:
                logger.error(f"Error in traceback monitor loop: {e}")
                time.sleep(10)  # Wait longer on error
    
    def start(self):
        """Start the traceback monitor in a background thread."""
        if self.running:
            return
        
        self.running = True
        self.thread = threading.Thread(target=self._monitor_loop, daemon=True)
        self.thread.start()
        logger.info("✅ Traceback monitor started")
    
    def stop(self):
        """Stop the traceback monitor."""
        self.running = False
        if self.thread:
            self.thread.join(timeout=5)
        logger.info("🛑 Traceback monitor stopped")


class SystemMonitor:
    def __init__(self, db_path="/sphere/app/system_monitor.db", poll_interval=20, enable_traceback_monitor=True):
        self.db_path = Path(db_path)
        self.poll_interval = poll_interval  # seconds
        self.last_job_scan = 0  # timestamp of last job directory scan
        self.known_jobs = set()  # job IDs we've already recorded completion for
        
        # RAM monitoring state (for high memory alerts)
        self.ram_alert_count = 0  # Number of RAM alerts sent (max 3)
        self.last_ram_alert_time = 0  # Timestamp of last RAM alert
        self.ram_alert_interval = 3600  # 1 hour in seconds
        self.ram_alert_threshold_gb = 100  # Alert if RAM usage > 100 GB
        self.max_ram_alerts = 3  # Maximum number of alerts to send
        
        # Initialize database
        self.init_database()
        
        # Load existing jobs on startup
        self.scan_existing_jobs()
        
        # Initialize traceback monitor
        self.traceback_monitor = None
        if enable_traceback_monitor:
            try:
                watch_dirs = ["/var/log/featrix", "/sphere/app/featrix_output"]
                self.traceback_monitor = TracebackMonitor(watch_dirs=watch_dirs)
                self.traceback_monitor.start()
            except Exception as e:
                logger.warning(f"Could not start traceback monitor: {e}")
        
        logger.info(f"System monitor initialized with {self.poll_interval}s polling interval")
        logger.info(f"Database: {self.db_path}")

    def init_database(self):
        """Initialize SQLite database for monitoring data."""
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("PRAGMA journal_mode=WAL")
            
            # System metrics table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS system_metrics (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp TEXT NOT NULL,
                    cpu_percent REAL,
                    memory_percent REAL,
                    memory_used_gb REAL,
                    memory_total_gb REAL,
                    gpu_count INTEGER,
                    gpu_utilization TEXT,  -- JSON array of GPU utilization %
                    gpu_memory_used TEXT,  -- JSON array of GPU memory used MB
                    gpu_memory_total TEXT, -- JSON array of GPU memory total MB
                    gpu_memory_percent TEXT, -- JSON array of GPU memory %
                    gpu_temperature TEXT,  -- JSON array of GPU temperatures
                    load_average_1m REAL,
                    load_average_5m REAL,
                    load_average_15m REAL,
                    disk_usage_percent REAL
                )
            """)
            
            # Job completion events table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS job_events (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    job_id TEXT NOT NULL,
                    job_type TEXT,
                    session_id TEXT,
                    event_type TEXT NOT NULL,  -- 'started', 'completed', 'failed'
                    timestamp TEXT NOT NULL,
                    runtime_seconds REAL,
                    log_file_size INTEGER,
                    working_directory TEXT,
                    cpu_time_user REAL,
                    cpu_time_system REAL,
                    peak_memory_mb REAL,
                    exit_code INTEGER
                )
            """)
            
            # Job resource correlation table (links jobs to system metrics during their runtime)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS job_resource_usage (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    job_id TEXT NOT NULL,
                    started_at TEXT,
                    completed_at TEXT,
                    avg_cpu_percent REAL,
                    max_cpu_percent REAL,
                    avg_memory_percent REAL,
                    max_memory_percent REAL,
                    avg_gpu_utilization TEXT,  -- JSON average across all GPUs
                    max_gpu_utilization TEXT,  -- JSON max across all GPUs
                    avg_gpu_memory_percent TEXT,
                    max_gpu_memory_percent TEXT,
                    runtime_seconds REAL
                )
            """)
            
            # Create indexes
            conn.execute("CREATE INDEX IF NOT EXISTS idx_metrics_timestamp ON system_metrics(timestamp)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_job_events_job_id ON job_events(job_id)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_job_events_timestamp ON job_events(timestamp)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_job_resource_job_id ON job_resource_usage(job_id)")
            
            conn.commit()
        
        logger.info("Database initialized successfully")

    def get_cpu_stats(self) -> Dict:
        """Get CPU usage statistics."""
        try:
            import psutil
            
            # CPU percentage
            cpu_percent = psutil.cpu_percent(interval=1)
            
            # Load averages
            load_avg = os.getloadavg()
            
            # Memory usage
            memory = psutil.virtual_memory()
            
            return {
                'cpu_percent': cpu_percent,
                'memory_percent': memory.percent,
                'memory_used_gb': memory.used / (1024**3),
                'memory_total_gb': memory.total / (1024**3),
                'load_average_1m': load_avg[0],
                'load_average_5m': load_avg[1],
                'load_average_15m': load_avg[2]
            }
        except ImportError:
            logger.warning("psutil not available - CPU/memory stats disabled")
            return {}
        except Exception as e:
            logger.error(f"Error getting CPU stats: {e}")
            return {}

    def check_ram_usage(self, memory_used_gb: float):
        """Check RAM usage and send Slack alerts if over threshold.
        
        Alerts are sent:
        - Once per hour (self.ram_alert_interval)
        - Maximum 3 times total (self.max_ram_alerts)
        - Only if memory_used_gb > self.ram_alert_threshold_gb (100 GB)
        """
        if memory_used_gb is None:
            return
        
        current_time = time.time()
        
        # Check if we're over threshold
        if memory_used_gb <= self.ram_alert_threshold_gb:
            # Memory is below threshold - reset alert count if it was high before
            if self.ram_alert_count > 0:
                logger.debug(f"RAM usage dropped below threshold ({memory_used_gb:.1f} GB < {self.ram_alert_threshold_gb} GB) - resetting alert count")
                self.ram_alert_count = 0
            return
        
        # We're over threshold - check if we should send an alert
        if self.ram_alert_count >= self.max_ram_alerts:
            # Already sent max alerts - don't send more
            return
        
        # Check if enough time has passed since last alert
        time_since_last = current_time - self.last_ram_alert_time
        if time_since_last < self.ram_alert_interval:
            # Not enough time has passed - skip this check
            return
        
        # Send alert
        try:
            from slack import send_slack_message
            import socket
            hostname = socket.gethostname()
            
            alert_number = self.ram_alert_count + 1
            slack_msg = (
                f"⚠️  **High RAM Usage Alert** ({alert_number}/{self.max_ram_alerts})\n"
                f"Host: {hostname}\n"
                f"Memory Used: {memory_used_gb:.1f} GB\n"
                f"Threshold: {self.ram_alert_threshold_gb} GB\n"
                f"\n"
                f"System is using over {self.ram_alert_threshold_gb} GB of RAM.\n"
                f"This alert will be sent up to {self.max_ram_alerts} times (once per hour).\n"
                f"\n"
                f"**Possible causes:**\n"
                f"- String cache growing unbounded\n"
                f"- SQLite cache size too large\n"
                f"- Memory leak in training process\n"
                f"- Large dataset loaded in memory"
            )
            
            send_slack_message(slack_msg, throttle=False, skip_hostname_prefix=True)
            logger.warning(f"🚨 High RAM usage alert sent: {memory_used_gb:.1f} GB (alert {alert_number}/{self.max_ram_alerts})")
            
            # Update state
            self.ram_alert_count += 1
            self.last_ram_alert_time = current_time
            
        except Exception as e:
            logger.error(f"Failed to send RAM usage alert: {e}")

    def get_gpu_stats(self) -> Dict:
        """Get GPU usage statistics using nvidia-smi."""
        try:
            import subprocess
            import json
            
            # Query GPU stats using nvidia-smi
            cmd = [
                'nvidia-smi', 
                '--query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu',
                '--format=csv,noheader,nounits'
            ]
            
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
            
            if result.returncode != 0:
                logger.warning("nvidia-smi command failed - GPU stats disabled")
                return {}
            
            lines = result.stdout.strip().split('\n')
            gpu_count = len(lines)
            
            gpu_utilization = []
            gpu_memory_used = []
            gpu_memory_total = []
            gpu_memory_percent = []
            gpu_temperature = []
            
            for line in lines:
                if line.strip():
                    parts = [p.strip() for p in line.split(',')]
                    if len(parts) >= 4:
                        util = float(parts[0]) if parts[0] != '[Not Supported]' else 0.0
                        mem_used = float(parts[1])
                        mem_total = float(parts[2])
                        temp = float(parts[3]) if parts[3] != '[Not Supported]' else 0.0
                        
                        gpu_utilization.append(util)
                        gpu_memory_used.append(mem_used)
                        gpu_memory_total.append(mem_total)
                        gpu_memory_percent.append((mem_used / mem_total * 100) if mem_total > 0 else 0)
                        gpu_temperature.append(temp)
            
            return {
                'gpu_count': gpu_count,
                'gpu_utilization': json.dumps(gpu_utilization),
                'gpu_memory_used': json.dumps(gpu_memory_used),
                'gpu_memory_total': json.dumps(gpu_memory_total),
                'gpu_memory_percent': json.dumps(gpu_memory_percent),
                'gpu_temperature': json.dumps(gpu_temperature)
            }
            
        except subprocess.TimeoutExpired:
            logger.warning("nvidia-smi timeout - GPU stats disabled")
            return {}
        except FileNotFoundError:
            logger.warning("nvidia-smi not found - GPU stats disabled")
            return {}
        except Exception as e:
            logger.error(f"Error getting GPU stats: {e}")
            return {}

    def get_disk_stats(self) -> Dict:
        """Get disk usage statistics."""
        try:
            import shutil
            
            # Check disk usage for /sphere/app
            total, used, free = shutil.disk_usage('/sphere/app')
            
            usage_percent = (used / total) * 100
            
            return {
                'disk_usage_percent': usage_percent
            }
        except Exception as e:
            logger.error(f"Error getting disk stats: {e}")
            return {}

    def collect_system_metrics(self):
        """Collect all system metrics and store in database."""
        timestamp = datetime.now(timezone.utc).isoformat()
        
        # Collect all metrics
        metrics = {
            'timestamp': timestamp,
            'cpu_percent': None,
            'memory_percent': None,
            'memory_used_gb': None,
            'memory_total_gb': None,
            'gpu_count': 0,
            'gpu_utilization': None,
            'gpu_memory_used': None,
            'gpu_memory_total': None,
            'gpu_memory_percent': None,
            'gpu_temperature': None,
            'load_average_1m': None,
            'load_average_5m': None,
            'load_average_15m': None,
            'disk_usage_percent': None
        }
        
        # Update with actual stats
        metrics.update(self.get_cpu_stats())
        metrics.update(self.get_gpu_stats())
        metrics.update(self.get_disk_stats())
        
        # Store in database
        try:
            with sqlite3.connect(self.db_path) as conn:
                columns = ', '.join(metrics.keys())
                placeholders = ', '.join(['?' for _ in metrics])
                
                conn.execute(
                    f"INSERT INTO system_metrics ({columns}) VALUES ({placeholders})",
                    list(metrics.values())
                )
                conn.commit()
                
            # Log summary (only every 5th reading to avoid spam)
            if int(time.time()) % (self.poll_interval * 5) < self.poll_interval:
                cpu_info = f"CPU: {metrics['cpu_percent']:.1f}%" if metrics['cpu_percent'] else "CPU: N/A"
                mem_info = f"RAM: {metrics['memory_percent']:.1f}%" if metrics['memory_percent'] else "RAM: N/A"
                
                gpu_info = "GPU: N/A"
                if metrics['gpu_utilization']:
                    gpu_utils = json.loads(metrics['gpu_utilization'])
                    gpu_mems = json.loads(metrics['gpu_memory_percent'])
                    if gpu_utils:
                        gpu_info = f"GPU: {gpu_utils[0]:.1f}% util, {gpu_mems[0]:.1f}% mem"
                
                logger.info(f"📊 {cpu_info} | {mem_info} | {gpu_info}")
                
        except Exception as e:
            logger.error(f"Error storing system metrics: {e}")

    def scan_existing_jobs(self):
        """Scan existing job directories to build initial known_jobs set."""
        try:
            output_dir = Path("/sphere/app/featrix_output")
            if not output_dir.exists():
                return
            
            for job_dir in output_dir.iterdir():
                if job_dir.is_dir():
                    job_name = job_dir.name
                    self.known_jobs.add(job_name)
                    
            logger.info(f"Loaded {len(self.known_jobs)} existing jobs")
            
        except Exception as e:
            logger.error(f"Error scanning existing jobs: {e}")

    def scan_for_job_events(self):
        """Scan for new job completions and record events."""
        try:
            output_dir = Path("/sphere/app/featrix_output")
            if not output_dir.exists():
                return
            
            current_jobs = set()
            new_completions = []
            
            for job_dir in output_dir.iterdir():
                if not job_dir.is_dir():
                    continue
                    
                job_name = job_dir.name
                current_jobs.add(job_name)
                
                # Skip if we've already processed this job
                if job_name in self.known_jobs:
                    continue
                
                # Parse job info
                job_parts = job_name.split('_')
                if len(job_parts) >= 3:
                    job_type = '_'.join(job_parts[:-2])
                    timestamp_str = job_parts[-2]
                    job_id = job_parts[-1]
                else:
                    job_type = "unknown"
                    job_id = job_name
                
                # Check if job has completed (has stdout.log)
                logs_dir = job_dir / "logs"
                stdout_log = logs_dir / "stdout.log"
                
                if stdout_log.exists():
                    # Job has produced output - record completion
                    try:
                        stat_info = stdout_log.stat()
                        log_size = stat_info.st_size
                        completion_time = datetime.fromtimestamp(stat_info.st_mtime, timezone.utc).isoformat()
                        
                        # Try to determine session ID from job directory structure
                        session_id = self.extract_session_id(job_dir)
                        
                        # Calculate runtime (rough estimate from directory creation to log modification)
                        job_creation_time = job_dir.stat().st_ctime
                        log_modification_time = stat_info.st_mtime
                        runtime_seconds = log_modification_time - job_creation_time
                        
                        new_completions.append({
                            'job_id': job_name,
                            'job_type': job_type,
                            'session_id': session_id,
                            'event_type': 'completed',
                            'timestamp': completion_time,
                            'runtime_seconds': runtime_seconds,
                            'log_file_size': log_size,
                            'working_directory': str(job_dir),
                            'cpu_time_user': None,  # Could parse from logs if needed
                            'cpu_time_system': None,
                            'peak_memory_mb': None,
                            'exit_code': None
                        })
                        
                    except Exception as e:
                        logger.error(f"Error processing job {job_name}: {e}")
            
            # Record new completions
            if new_completions:
                with sqlite3.connect(self.db_path) as conn:
                    for completion in new_completions:
                        conn.execute("""
                            INSERT INTO job_events 
                            (job_id, job_type, session_id, event_type, timestamp, 
                             runtime_seconds, log_file_size, working_directory, 
                             cpu_time_user, cpu_time_system, peak_memory_mb, exit_code)
                            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                        """, (
                            completion['job_id'],
                            completion['job_type'], 
                            completion['session_id'],
                            completion['event_type'],
                            completion['timestamp'],
                            completion['runtime_seconds'],
                            completion['log_file_size'],
                            completion['working_directory'],
                            completion['cpu_time_user'],
                            completion['cpu_time_system'],
                            completion['peak_memory_mb'],
                            completion['exit_code']
                        ))
                    conn.commit()
                
                logger.info(f"📝 Recorded {len(new_completions)} new job completions")
                for completion in new_completions:
                    logger.info(f"   ✅ {completion['job_type']} {completion['job_id'][:8]} ({completion['runtime_seconds']:.1f}s)")
            
            # Update known jobs
            self.known_jobs = current_jobs
            
        except Exception as e:
            logger.error(f"Error scanning for job events: {e}")

    def extract_session_id(self, job_dir: Path) -> Optional[str]:
        """Try to extract session ID from job directory or associated files."""
        try:
            # Look for session information in job directory
            # This is a best-effort attempt - might need refinement
            
            job_name = job_dir.name
            job_parts = job_name.split('_')
            
            if len(job_parts) >= 3:
                timestamp_str = job_parts[-2]
                job_id = job_parts[-1]
                
                # Session ID often follows pattern: YYYYMMDD-HHMMSS_jobid
                potential_session_id = f"{timestamp_str}_{job_id}"
                return potential_session_id
            
            return None
            
        except Exception:
            return None

    def calculate_job_resource_usage(self):
        """Calculate resource usage statistics for completed jobs."""
        try:
            with sqlite3.connect(self.db_path) as conn:
                # Find jobs that have completed but don't have resource usage calculated
                cursor = conn.execute("""
                    SELECT DISTINCT je.job_id, je.timestamp as completed_at
                    FROM job_events je
                    LEFT JOIN job_resource_usage jru ON je.job_id = jru.job_id
                    WHERE je.event_type = 'completed' 
                    AND jru.job_id IS NULL
                    AND je.timestamp > datetime('now', '-24 hours')
                """)
                
                uncalculated_jobs = cursor.fetchall()
                
                for job_id, completed_at in uncalculated_jobs:
                    # Calculate resource usage for this job
                    # Look for metrics during job runtime (rough estimate)
                    completed_dt = datetime.fromisoformat(completed_at.replace('Z', '+00:00'))
                    
                    # Estimate start time (this is rough - could be improved)
                    start_time = completed_dt.replace(microsecond=0) - timedelta(hours=2)  # Assume max 2 hour jobs
                    
                    # Query metrics during this timeframe
                    metrics_cursor = conn.execute("""
                        SELECT cpu_percent, memory_percent, gpu_utilization, gpu_memory_percent
                        FROM system_metrics
                        WHERE timestamp BETWEEN ? AND ?
                        AND cpu_percent IS NOT NULL
                    """, (start_time.isoformat(), completed_at))
                    
                    metrics_data = metrics_cursor.fetchall()
                    
                    if metrics_data:
                        # Calculate averages and maximums
                        cpu_values = [m[0] for m in metrics_data if m[0] is not None]
                        memory_values = [m[1] for m in metrics_data if m[1] is not None]
                        
                        avg_cpu = sum(cpu_values) / len(cpu_values) if cpu_values else None
                        max_cpu = max(cpu_values) if cpu_values else None
                        avg_memory = sum(memory_values) / len(memory_values) if memory_values else None
                        max_memory = max(memory_values) if memory_values else None
                        
                        # GPU calculations (more complex due to JSON arrays)
                        gpu_utils = []
                        gpu_mems = []
                        
                        for m in metrics_data:
                            if m[2]:  # gpu_utilization
                                try:
                                    utils = json.loads(m[2])
                                    if utils:
                                        gpu_utils.extend(utils)
                                except:
                                    pass
                                    
                            if m[3]:  # gpu_memory_percent
                                try:
                                    mems = json.loads(m[3])
                                    if mems:
                                        gpu_mems.extend(mems)
                                except:
                                    pass
                        
                        avg_gpu_util = sum(gpu_utils) / len(gpu_utils) if gpu_utils else None
                        max_gpu_util = max(gpu_utils) if gpu_utils else None
                        avg_gpu_mem = sum(gpu_mems) / len(gpu_mems) if gpu_mems else None
                        max_gpu_mem = max(gpu_mems) if gpu_mems else None
                        
                        # Store resource usage summary
                        conn.execute("""
                            INSERT INTO job_resource_usage 
                            (job_id, started_at, completed_at, avg_cpu_percent, max_cpu_percent,
                             avg_memory_percent, max_memory_percent, avg_gpu_utilization, 
                             max_gpu_utilization, avg_gpu_memory_percent, max_gpu_memory_percent)
                            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                        """, (
                            job_id, start_time.isoformat(), completed_at,
                            avg_cpu, max_cpu, avg_memory, max_memory,
                            json.dumps([avg_gpu_util]) if avg_gpu_util else None,
                            json.dumps([max_gpu_util]) if max_gpu_util else None,
                            json.dumps([avg_gpu_mem]) if avg_gpu_mem else None,
                            json.dumps([max_gpu_mem]) if max_gpu_mem else None
                        ))
                
                conn.commit()
                
                if uncalculated_jobs:
                    logger.info(f"📈 Calculated resource usage for {len(uncalculated_jobs)} jobs")
                    
        except Exception as e:
            logger.error(f"Error calculating job resource usage: {e}")

    def cleanup_old_data(self):
        """Clean up old monitoring data to prevent database bloat."""
        try:
            cutoff_days = 7  # Keep 7 days of detailed metrics
            cutoff_date = datetime.now(timezone.utc) - timedelta(days=cutoff_days)
            
            with sqlite3.connect(self.db_path) as conn:
                # Keep job events and resource usage longer (30 days)
                job_cutoff = datetime.now(timezone.utc) - timedelta(days=30)
                
                # Clean old system metrics
                cursor = conn.execute(
                    "DELETE FROM system_metrics WHERE timestamp < ?",
                    (cutoff_date.isoformat(),)
                )
                metrics_deleted = cursor.rowcount
                
                # Clean old job events  
                cursor = conn.execute(
                    "DELETE FROM job_events WHERE timestamp < ?",
                    (job_cutoff.isoformat(),)
                )
                jobs_deleted = cursor.rowcount
                
                conn.commit()
                
                if metrics_deleted > 0 or jobs_deleted > 0:
                    logger.info(f"🧹 Cleaned up {metrics_deleted} old metrics, {jobs_deleted} old job events")
                    
        except Exception as e:
            logger.error(f"Error cleaning up old data: {e}")

    def run(self):
        """Main monitoring loop."""
        logger.info("🚀 Starting system monitor...")
        logger.info(f"📊 Polling every {self.poll_interval} seconds")
        
        cleanup_counter = 0
        
        try:
            while True:
                start_time = time.time()
                
                # Collect system metrics
                self.collect_system_metrics()
                
                # Scan for job events (every cycle)
                self.scan_for_job_events()
                
                # Calculate job resource usage (every 5 cycles)
                if cleanup_counter % 5 == 0:
                    self.calculate_job_resource_usage()
                
                # Cleanup old data (every 100 cycles - roughly every 30 minutes)
                if cleanup_counter % 100 == 0:
                    self.cleanup_old_data()
                
                cleanup_counter += 1
                
                # Sleep for remaining time in polling interval
                elapsed = time.time() - start_time
                sleep_time = max(0, self.poll_interval - elapsed)
                
                if sleep_time > 0:
                    time.sleep(sleep_time)
                else:
                    logger.warning(f"Monitoring cycle took {elapsed:.1f}s (longer than {self.poll_interval}s interval)")
                    
        except KeyboardInterrupt:
            logger.info("🛑 Monitor stopped by user")
        except Exception as e:
            logger.error(f"💥 Monitor crashed: {e}")
            raise
        finally:
            # Stop traceback monitor
            if self.traceback_monitor:
                self.traceback_monitor.stop()

def main():
    """Main entry point."""
    # Install Featrix exception hook for better error tracking
    try:
        from lib.featrix_debug import install_featrix_excepthook
        install_featrix_excepthook()
    except Exception:
        pass  # Don't fail if debug module not available
    
    import argparse
    
    parser = argparse.ArgumentParser(description="Featrix System Resource Monitor")
    parser.add_argument('--interval', type=int, default=20, 
                       help='Polling interval in seconds (default: 20)')
    parser.add_argument('--db-path', default='/sphere/app/system_monitor.db',
                       help='Database path (default: /sphere/app/system_monitor.db)')
    parser.add_argument('--query', action='store_true',
                       help='Query recent data instead of running monitor')
    parser.add_argument('--no-traceback-monitor', action='store_true',
                       help='Disable traceback monitoring')
    
    args = parser.parse_args()
    
    if args.query:
        # Quick query mode for testing
        monitor = SystemMonitor(db_path=args.db_path, poll_interval=args.interval)
        
        with sqlite3.connect(args.db_path) as conn:
            # Show recent metrics
            cursor = conn.execute("""
                SELECT timestamp, cpu_percent, memory_percent, gpu_utilization
                FROM system_metrics 
                ORDER BY timestamp DESC 
                LIMIT 5
            """)
            
            print("Recent system metrics:")
            for row in cursor.fetchall():
                timestamp, cpu, memory, gpu = row
                gpu_info = ""
                if gpu:
                    try:
                        gpu_data = json.loads(gpu)
                        if gpu_data:
                            gpu_info = f"GPU: {gpu_data[0]:.1f}%"
                    except:
                        gpu_info = "GPU: N/A"
                
                print(f"  {timestamp}: CPU {cpu:.1f}%, RAM {memory:.1f}%, {gpu_info}")
            
            # Show recent job completions
            cursor = conn.execute("""
                SELECT job_id, job_type, timestamp, runtime_seconds
                FROM job_events
                WHERE event_type = 'completed'
                ORDER BY timestamp DESC
                LIMIT 5
            """)
            
            print("\nRecent job completions:")
            for row in cursor.fetchall():
                job_id, job_type, timestamp, runtime = row
                print(f"  {timestamp}: {job_type} {job_id[:16]} ({runtime:.1f}s)")
    else:
        # Normal monitoring mode
        monitor = SystemMonitor(
            db_path=args.db_path, 
            poll_interval=args.interval,
            enable_traceback_monitor=not args.no_traceback_monitor
        )
        monitor.run()

if __name__ == "__main__":
    # Handle missing dependencies gracefully
    try:
        from datetime import timedelta
    except ImportError:
        print("Missing required dependencies. Install with: pip install psutil")
        exit(1)
        
    main() 