#!/usr/bin/env python3
"""
 ____ ____  _   _   ____                       __  __           _
/ ___|  _ \| | | | |  _ \ _____      _____ _ _|  \/  | ___   __| | ___
| |   | |_) | | | | | |_) / _ \ \ /\ / / _ \ '__| |\/| |/ _ \ / _` |/ _ \
| |___|  __/| |_| | |  __/ (_) \ V  V /  __/ |  | |  | | (_) | (_| |  __/
\____|_|    \___/  |_|   \___/ \_/\_/ \___|_|  |_|  |_|\___/ \__,_|\___|

CPU Power Mode Management Script
Switches between Normal and Power Saving CPU modes.

Modes:
- normal: CPU governor ondemand (600MHz-1500MHz dynamic)
- powersave: CPU governor powersave (600MHz fixed)

Usage:
  /usr/bin/python3 /var/www/nebuleair_pro_4g/power/set_cpu_mode.py <mode>
  /usr/bin/python3 /var/www/nebuleair_pro_4g/power/set_cpu_mode.py normal
  /usr/bin/python3 /var/www/nebuleair_pro_4g/power/set_cpu_mode.py powersave

Or get current mode:
  /usr/bin/python3 /var/www/nebuleair_pro_4g/power/set_cpu_mode.py get
"""

import sqlite3
import subprocess
import sys
import datetime
import json

# Paths
DB_PATH = "/var/www/nebuleair_pro_4g/sqlite/sensors.db"
LOG_PATH = "/var/www/nebuleair_pro_4g/logs/app.log"

# Available modes
MODES = {
    "normal": {
        "governor": "ondemand",
        "description": "Normal mode - CPU 600MHz-1500MHz dynamic",
        "min_freq": 600000,  # 600 MHz in kHz
        "max_freq": 1500000  # 1500 MHz in kHz
    },
    "powersave": {
        "governor": "powersave",
        "description": "Power saving mode - CPU 600MHz fixed",
        "min_freq": 600000,
        "max_freq": 600000
    }
}

def log_message(message):
    """Write message to log file with timestamp"""
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_entry = f"[{timestamp}] [CPU Power Mode] {message}\n"
    try:
        with open(LOG_PATH, "a") as log_file:
            log_file.write(log_entry)
    except Exception as e:
        print(f"Failed to write to log: {e}", file=sys.stderr)
    print(log_entry.strip())

def get_cpu_count():
    """Get number of CPU cores"""
    try:
        result = subprocess.run(
            ["nproc"],
            capture_output=True,
            text=True,
            timeout=5
        )
        return int(result.stdout.strip())
    except Exception as e:
        log_message(f"Failed to get CPU count: {e}")
        return 4  # Default to 4 cores for CM4

def set_cpu_governor(governor, cpu_id):
    """Set CPU governor for specific CPU core"""
    try:
        with open(f"/sys/devices/system/cpu/cpu{cpu_id}/cpufreq/scaling_governor", "w") as f:
            f.write(governor)
        return True
    except Exception as e:
        log_message(f"Failed to set governor for CPU{cpu_id}: {e}")
        return False

def set_cpu_freq(min_freq, max_freq, cpu_id):
    """Set CPU frequency limits for specific CPU core"""
    try:
        # Set min frequency
        with open(f"/sys/devices/system/cpu/cpu{cpu_id}/cpufreq/scaling_min_freq", "w") as f:
            f.write(str(min_freq))

        # Set max frequency
        with open(f"/sys/devices/system/cpu/cpu{cpu_id}/cpufreq/scaling_max_freq", "w") as f:
            f.write(str(max_freq))

        return True
    except Exception as e:
        log_message(f"Failed to set frequency for CPU{cpu_id}: {e}")
        return False

def get_current_cpu_state():
    """Get current CPU governor and frequencies"""
    try:
        with open("/sys/devices/system/cpu/cpu0/cpufreq/scaling_governor", "r") as f:
            governor = f.read().strip()

        with open("/sys/devices/system/cpu/cpu0/cpufreq/scaling_min_freq", "r") as f:
            min_freq = int(f.read().strip())

        with open("/sys/devices/system/cpu/cpu0/cpufreq/scaling_max_freq", "r") as f:
            max_freq = int(f.read().strip())

        with open("/sys/devices/system/cpu/cpu0/cpufreq/scaling_cur_freq", "r") as f:
            cur_freq = int(f.read().strip())

        return {
            "governor": governor,
            "min_freq_khz": min_freq,
            "max_freq_khz": max_freq,
            "current_freq_khz": cur_freq,
            "min_freq_mhz": min_freq / 1000,
            "max_freq_mhz": max_freq / 1000,
            "current_freq_mhz": cur_freq / 1000
        }
    except Exception as e:
        log_message(f"Failed to get current CPU state: {e}")
        return None

def update_config_db(mode):
    """Update cpu_power_mode in database"""
    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        cursor.execute(
            "UPDATE config_table SET value = ? WHERE key = 'cpu_power_mode'",
            (mode,)
        )
        conn.commit()
        conn.close()
        return True
    except sqlite3.Error as e:
        log_message(f"Database error: {e}")
        return False

def get_config_from_db():
    """Read cpu_power_mode from database"""
    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        cursor.execute("SELECT value FROM config_table WHERE key = 'cpu_power_mode'")
        result = cursor.fetchone()
        conn.close()
        return result[0] if result else None
    except sqlite3.Error as e:
        log_message(f"Database error: {e}")
        return None

def apply_mode(mode):
    """Apply CPU power mode"""
    if mode not in MODES:
        log_message(f"Invalid mode: {mode}. Valid modes: {', '.join(MODES.keys())}")
        return False

    mode_config = MODES[mode]
    log_message(f"Applying mode: {mode} - {mode_config['description']}")

    cpu_count = get_cpu_count()
    log_message(f"Detected {cpu_count} CPU cores")

    success = True

    # Apply settings to all CPU cores
    for cpu_id in range(cpu_count):
        # Set frequency limits first
        if not set_cpu_freq(mode_config['min_freq'], mode_config['max_freq'], cpu_id):
            success = False

        # Then set governor
        if not set_cpu_governor(mode_config['governor'], cpu_id):
            success = False

    if success:
        log_message(f"Successfully applied {mode} mode to all {cpu_count} cores")

        # Update database
        if update_config_db(mode):
            log_message(f"Updated database: cpu_power_mode = {mode}")
        else:
            log_message("Warning: Failed to update database")

        # Log current state
        state = get_current_cpu_state()
        if state:
            log_message(f"Current CPU state: {state['governor']} governor, "
                       f"{state['min_freq_mhz']:.0f}-{state['max_freq_mhz']:.0f} MHz "
                       f"(current: {state['current_freq_mhz']:.0f} MHz)")
    else:
        log_message(f"Failed to apply {mode} mode completely")

    return success

def main():
    """Main function"""
    if len(sys.argv) < 2:
        print("Usage: set_cpu_mode.py <mode>")
        print("Modes: normal, powersave")
        print("Or: set_cpu_mode.py get  (to get current state)")
        return 1

    command = sys.argv[1].lower()

    # Handle "get" command to return current state
    if command == "get":
        state = get_current_cpu_state()
        db_mode = get_config_from_db()

        if state:
            output = {
                "success": True,
                "cpu_state": state,
                "config_mode": db_mode
            }
            print(json.dumps(output))
            return 0
        else:
            output = {
                "success": False,
                "error": "Failed to read CPU state"
            }
            print(json.dumps(output))
            return 1

    # Handle mode setting
    log_message("=== CPU Power Mode Script Started ===")

    if apply_mode(command):
        log_message("=== CPU Power Mode Script Completed Successfully ===")
        output = {
            "success": True,
            "mode": command,
            "description": MODES[command]['description']
        }
        print(json.dumps(output))
        return 0
    else:
        log_message("=== CPU Power Mode Script Failed ===")
        output = {
            "success": False,
            "error": f"Failed to apply mode: {command}"
        }
        print(json.dumps(output))
        return 1

if __name__ == "__main__":
    sys.exit(main())