#!/usr/bin/env python3
"""
macOS Real-Time System Monitor Server v2.0
WebSocket + GPU Profiling + Temperature
"""

from flask import Flask, jsonify
from flask_socketio import SocketIO, emit
from flask_cors import CORS
import subprocess
import json
import psutil
import platform
import re
from datetime import datetime
from threading import Thread
import time

app = Flask(__name__)
app.config['SECRET_KEY'] = 'macos-monitor-secret'
CORS(app)
socketio = SocketIO(app, cors_allowed_origins="*", async_mode='threading')

# Historical data (last 60 samples = 2 minutes at 2s interval)
history = {
    'cpu': [],
    'gpu': [],
    'memory': [],
    'temperature': [],
    'timestamps': []
}
MAX_HISTORY = 60

def run_command(cmd):
    """Execute shell command and return output"""
    try:
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=5)
        return result.stdout.strip()
    except Exception as e:
        return ""

def get_cpu_info():
    """Get real CPU information"""
    try:
        brand = run_command("sysctl -n machdep.cpu.brand_string")
        cores = run_command("sysctl -n hw.ncpu")
        perf_cores = run_command("sysctl -n hw.perflevel0.logicalcpu 2>/dev/null") or "0"
        eff_cores = run_command("sysctl -n hw.perflevel1.logicalcpu 2>/dev/null") or "0"

        cpu_percent = psutil.cpu_percent(interval=0.1, percpu=True)

        return {
            'brand': brand,
            'total_cores': int(cores) if cores else psutil.cpu_count(),
            'performance_cores': int(perf_cores) if perf_cores != "0" else None,
            'efficiency_cores': int(eff_cores) if eff_cores != "0" else None,
            'usage_percent': psutil.cpu_percent(interval=0.1),
            'usage_per_core': cpu_percent,
            'frequency_mhz': psutil.cpu_freq().current if psutil.cpu_freq() else None
        }
    except Exception as e:
        print(f"[Error] get_cpu_info: {e}")
        return {'error': str(e)}

def get_gpu_info():
    """Get GPU information with Metal API data"""
    try:
        # First try system_profiler for GPU info
        gpu_raw = run_command("system_profiler SPDisplaysDataType -json 2>/dev/null")

        gpu_list = []

        if gpu_raw:
            gpu_data = json.loads(gpu_raw)
            displays = gpu_data.get('SPDisplaysDataType', [])

            for display in displays:
                chipset = display.get('sppci_model', 'Unknown')
                vram = display.get('spdisplays_vram', 'Unknown')

                gpu_list.append({
                    'name': chipset,
                    'vram': vram,
                    'type': 'Integrated' if 'Apple' in chipset else 'Discrete',
                    'cores': None,
                    'usage_percent': None
                })

        # Try to get GPU cores info from ioreg (Apple Silicon)
        ioreg_gpu = run_command("ioreg -r -d 1 -w 0 -c IOAccelerator 2>/dev/null")

        if 'AGXAccelerator' in ioreg_gpu or 'AppleM' in ioreg_gpu:
            # Apple Silicon detected
            brand = run_command("sysctl -n machdep.cpu.brand_string")
            m_chip = re.search(r'Apple (M\d+[^\s]*)', brand)

            if m_chip:
                chip_name = m_chip.group(1)

                # M4 has 10 GPU cores (standard config)
                gpu_cores = 10
                if 'M4 Pro' in chip_name:
                    gpu_cores = 16
                elif 'M4 Max' in chip_name:
                    gpu_cores = 32

                gpu_list = [{
                    'name': f'{chip_name} GPU',
                    'vram': 'Unified Memory',
                    'type': 'Integrated',
                    'cores': gpu_cores,
                    'usage_percent': None  # Will be populated by powermetrics if available
                }]

        return gpu_list if gpu_list else [{'name': 'Unknown GPU', 'vram': 'N/A'}]

    except Exception as e:
        print(f"[Error] get_gpu_info: {e}")
        return [{'error': str(e)}]

def get_temperature():
    """Get system temperature from powermetrics or istats"""
    try:
        # Try powermetrics (requires sudo, but works without for some metrics)
        temp_raw = run_command("sudo powermetrics --samplers smc -i 500 -n 1 2>/dev/null | grep -i 'CPU die temperature' | head -1")

        if temp_raw:
            temp_match = re.search(r'(\d+\.\d+)', temp_raw)
            if temp_match:
                cpu_temp = float(temp_match.group(1))

                # Try GPU temp
                gpu_temp_raw = run_command("sudo powermetrics --samplers smc -i 500 -n 1 2>/dev/null | grep -i 'GPU die temperature' | head -1")
                gpu_temp = None
                if gpu_temp_raw:
                    gpu_match = re.search(r'(\d+\.\d+)', gpu_temp_raw)
                    if gpu_match:
                        gpu_temp = float(gpu_match.group(1))

                return {
                    'cpu_temp_c': cpu_temp,
                    'gpu_temp_c': gpu_temp,
                    'source': 'powermetrics'
                }

        # Fallback: try istats (if installed)
        istats = run_command("istats cpu temp 2>/dev/null")
        if istats:
            temp_match = re.search(r'(\d+\.\d+)', istats)
            if temp_match:
                return {
                    'cpu_temp_c': float(temp_match.group(1)),
                    'gpu_temp_c': None,
                    'source': 'istats'
                }

        # Fallback: estimate based on CPU usage (for demo)
        cpu_usage = psutil.cpu_percent(interval=0.1)
        estimated_temp = 35 + (cpu_usage * 0.5)  # Rough estimate

        return {
            'cpu_temp_c': round(estimated_temp, 1),
            'gpu_temp_c': None,
            'source': 'estimated',
            'note': 'Install istats or run with sudo for real temperature'
        }

    except Exception as e:
        print(f"[Error] get_temperature: {e}")
        return {'error': str(e)}

def get_memory_info():
    """Get real memory information"""
    try:
        mem = psutil.virtual_memory()
        swap = psutil.swap_memory()

        return {
            'total_gb': round(mem.total / (1024**3), 2),
            'used_gb': round(mem.used / (1024**3), 2),
            'available_gb': round(mem.available / (1024**3), 2),
            'percent': mem.percent,
            'swap_total_gb': round(swap.total / (1024**3), 2),
            'swap_used_gb': round(swap.used / (1024**3), 2),
            'swap_percent': swap.percent
        }
    except Exception as e:
        return {'error': str(e)}

def get_disk_info():
    """Get disk information"""
    try:
        disk = psutil.disk_usage('/')
        io = psutil.disk_io_counters()

        return {
            'total_gb': round(disk.total / (1024**3), 2),
            'used_gb': round(disk.used / (1024**3), 2),
            'free_gb': round(disk.free / (1024**3), 2),
            'percent': disk.percent,
            'read_mb': round(io.read_bytes / (1024**2), 2) if io else 0,
            'write_mb': round(io.write_bytes / (1024**2), 2) if io else 0
        }
    except Exception as e:
        return {'error': str(e)}

def get_battery_info():
    """Get battery information"""
    try:
        battery = psutil.sensors_battery()

        if battery:
            return {
                'percent': battery.percent,
                'power_plugged': battery.power_plugged,
                'seconds_left': battery.secsleft if battery.secsleft != psutil.POWER_TIME_UNLIMITED else None,
                'time_left_formatted': format_time(battery.secsleft) if battery.secsleft != psutil.POWER_TIME_UNLIMITED else 'Unlimited'
            }
        else:
            return {'error': 'No battery detected (desktop?)'}
    except Exception as e:
        return {'error': str(e)}

def get_network_info():
    """Get network information"""
    try:
        net_io = psutil.net_io_counters()
        connections = len(psutil.net_connections())

        return {
            'bytes_sent_mb': round(net_io.bytes_sent / (1024**2), 2),
            'bytes_recv_mb': round(net_io.bytes_recv / (1024**2), 2),
            'packets_sent': net_io.packets_sent,
            'packets_recv': net_io.packets_recv,
            'connections_count': connections
        }
    except Exception as e:
        return {'error': str(e)}

def get_top_processes(limit=10):
    """Get top processes by CPU usage"""
    try:
        processes = []
        for proc in psutil.process_iter(['pid', 'name', 'cpu_percent', 'memory_percent']):
            try:
                pinfo = proc.info
                processes.append({
                    'pid': pinfo['pid'],
                    'name': pinfo['name'],
                    'cpu_percent': pinfo['cpu_percent'],
                    'memory_percent': round(pinfo['memory_percent'], 2)
                })
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass

        processes.sort(key=lambda x: x['cpu_percent'] if x['cpu_percent'] else 0, reverse=True)
        return processes[:limit]

    except Exception as e:
        return []

def format_time(seconds):
    """Format seconds to readable time"""
    if seconds is None or seconds < 0:
        return "N/A"
    hours = seconds // 3600
    minutes = (seconds % 3600) // 60
    return f"{int(hours)}h {int(minutes)}m"

def get_system_info():
    """Get general system information"""
    try:
        boot_time = datetime.fromtimestamp(psutil.boot_time())
        uptime_seconds = (datetime.now() - boot_time).total_seconds()

        return {
            'platform': platform.system(),
            'platform_release': platform.release(),
            'platform_version': platform.version(),
            'architecture': platform.machine(),
            'hostname': platform.node(),
            'boot_time': boot_time.isoformat(),
            'uptime_seconds': int(uptime_seconds),
            'uptime_formatted': format_time(uptime_seconds)
        }
    except Exception as e:
        return {'error': str(e)}

def collect_all_data():
    """Collect all system data"""
    data = {
        'timestamp': datetime.now().isoformat(),
        'system': get_system_info(),
        'cpu': get_cpu_info(),
        'gpu': get_gpu_info(),
        'memory': get_memory_info(),
        'disk': get_disk_info(),
        'temperature': get_temperature(),
        'battery': get_battery_info(),
        'network': get_network_info(),
        'top_processes': get_top_processes(10)
    }

    # Add to history
    history['timestamps'].append(datetime.now().isoformat())
    history['cpu'].append(data['cpu'].get('usage_percent', 0))
    history['memory'].append(data['memory'].get('percent', 0))

    temp_data = data['temperature']
    if isinstance(temp_data, dict) and 'cpu_temp_c' in temp_data:
        history['temperature'].append(temp_data['cpu_temp_c'])
    else:
        history['temperature'].append(None)

    # Keep only last MAX_HISTORY samples
    for key in history:
        if len(history[key]) > MAX_HISTORY:
            history[key] = history[key][-MAX_HISTORY:]

    data['history'] = history

    return data

# WebSocket event handlers
@socketio.on('connect')
def handle_connect():
    print(f"[WebSocket] Client connected: {request.sid if 'request' in dir() else 'unknown'}")
    # Send initial data immediately
    data = collect_all_data()
    emit('system_update', data)

@socketio.on('disconnect')
def handle_disconnect():
    print(f"[WebSocket] Client disconnected")

@socketio.on('request_update')
def handle_request_update():
    """Handle manual update request"""
    data = collect_all_data()
    emit('system_update', data)

# Background thread to push updates
def background_updates():
    """Push system updates every 2 seconds"""
    while True:
        time.sleep(2)
        data = collect_all_data()
        socketio.emit('system_update', data, namespace='/')

# REST API endpoints (for compatibility)
@app.route('/api/system/all', methods=['GET'])
def get_all_data():
    """Get all system data in one request"""
    return jsonify(collect_all_data())

@app.route('/api/health', methods=['GET'])
def health_check():
    return jsonify({'status': 'ok', 'message': 'System Monitor Server v2.0 is running', 'websocket': True})

@app.route('/', methods=['GET'])
def index():
    return jsonify({
        'name': 'macOS System Monitor API v2.0',
        'version': '2.0.0',
        'features': ['WebSocket Real-Time', 'GPU Profiling', 'Temperature Monitoring', 'Historical Data'],
        'websocket': {
            'url': 'ws://localhost:5555',
            'events': {
                'connect': 'Auto-sends initial data',
                'system_update': 'Pushed every 2 seconds',
                'request_update': 'Manual update request'
            }
        },
        'endpoints': {
            '/api/system/all': 'Get all system data (REST fallback)',
            '/api/health': 'Health check'
        }
    })

if __name__ == '__main__':
    print("=" * 70)
    print("🖥️  macOS System Monitor Server v2.0")
    print("=" * 70)
    print("✨ Features: WebSocket + GPU Profiling + Temperature + History")
    print("\nPort: 5555")
    print("CORS: Enabled (all origins)")
    print("\n📡 WebSocket:")
    print("  → ws://localhost:5555")
    print("  → Event: 'system_update' (pushed every 2s)")
    print("\n🌐 REST API:")
    print("  → http://localhost:5555/api/system/all")
    print("  → http://localhost:5555/api/health")
    print("\nAccess via Tailscale:")
    print("  → ws://100.75.88.8:5555")
    print("  → http://100.75.88.8:5555/api/system/all")
    print("=" * 70)

    # Start background update thread
    update_thread = Thread(target=background_updates, daemon=True)
    update_thread.start()
    print("\n[Background] Update thread started (2s interval)")

    print("\n[Server] Starting on port 5555...")
    print("[WebSocket] Ready for connections\n")

    socketio.run(app, host='0.0.0.0', port=5555, debug=True, allow_unsafe_werkzeug=True)