#!/usr/bin/env python3
"""
통합 시스템 제어 스크립트 (Unified System Control Script)
기존의 개별 제어 스크립트들을 하나로 통합하였습니다.

지원 기능:
- power_on      (06-PowerON.py)
- power_off     (07-PowerOFF.py)
- log_clear     (05-clrsel.py)
- job_delete    (08-job_delete_all.py)
- tsr_collect   (03-tsr_log.py)
- tsr_save      (04-tsr_save.py)
"""

import os
import sys
import argparse
import subprocess
import time
import logging
import re
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Optional
from dotenv import load_dotenv

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

# Load .env relative to script location or current directory
load_dotenv() 

IDRAC_USER = os.getenv("IDRAC_USER", "root")
IDRAC_PASS = os.getenv("IDRAC_PASS", "calvin")
TSR_SHARE_URL = os.getenv("TSR_SHARE_URL", "//10.10.3.15/share/")
OME_USER = os.getenv("OME_USER", "")
OME_PASS = os.getenv("OME_PASS", "")
TSR_JOB_TIMEOUT = int(os.getenv("TSR_JOB_TIMEOUT", "900"))
TSR_JOB_POLL_INTERVAL = int(os.getenv("TSR_JOB_POLL_INTERVAL", "10"))
TSR_SAVE_WORKERS = int(os.getenv("TSR_SAVE_WORKERS", "30"))

# -----------------------------------------------------------------------------
# Command Definitions
# -----------------------------------------------------------------------------
# Each action maps to a function that returns the racadm command list
def cmd_power_on(ip):
    return ["racadm", "-r", ip, "-u", IDRAC_USER, "-p", IDRAC_PASS, "serveraction", "powerup"]

def cmd_power_off(ip):
    return ["racadm", "-r", ip, "-u", IDRAC_USER, "-p", IDRAC_PASS, "serveraction", "powerdown"]

def cmd_log_clear(ip):
    return ["racadm", "-r", ip, "-u", IDRAC_USER, "-p", IDRAC_PASS, "clrsel"]

def cmd_job_delete(ip):
    return ["racadm", "-r", ip, "-u", IDRAC_USER, "-p", IDRAC_PASS, "jobqueue", "delete", "-i", "ALL"]

def cmd_tsr_collect(ip):
    return ["racadm", "-r", ip, "-u", IDRAC_USER, "-p", IDRAC_PASS, "techsupreport", "collect"]

def cmd_tsr_save(ip):
    return [
        "racadm", "-r", ip, "-u", IDRAC_USER, "-p", IDRAC_PASS,
        "techsupreport", "export", "-l", TSR_SHARE_URL, "-u", OME_USER, "-p", OME_PASS
    ]

ACTION_MAP = {
    "power_on": cmd_power_on,
    "power_off": cmd_power_off,
    "log_clear": cmd_log_clear,
    "job_delete": cmd_job_delete,
    "tsr_collect": cmd_tsr_collect,
    "tsr_save": cmd_tsr_save,
}

# -----------------------------------------------------------------------------
# Execution Logic
# -----------------------------------------------------------------------------
def redact_cmd(cmd):
    redacted = []
    hide_next = False
    for part in cmd:
        if hide_next:
            redacted.append("***")
            hide_next = False
            continue
        redacted.append(part)
        if part == "-p":
            hide_next = True
    return redacted


def extract_job_id(output: str) -> Optional[str]:
    match = re.search(r"Job ID\s*=\s*(JID_\d+)", output or "")
    return match.group(1) if match else None


def wait_for_job(ip: str, job_id: str) -> bool:
    deadline = time.time() + TSR_JOB_TIMEOUT
    view_cmd = [
        "racadm", "-r", ip, "-u", IDRAC_USER, "-p", IDRAC_PASS,
        "jobqueue", "view", "-i", job_id
    ]

    while time.time() < deadline:
        result = subprocess.run(
            view_cmd,
            capture_output=True,
            text=True,
            timeout=120
        )
        output = "\n".join(part for part in [result.stdout.strip(), result.stderr.strip()] if part)
        logger.info(f"[{ip}] Job {job_id} status:\n{output}")

        if "Status=Completed" in output:
            return True
        if "Status=Failed" in output:
            return False

        time.sleep(TSR_JOB_POLL_INTERVAL)

    logger.error(f"[{ip}] Job {job_id} status check timed out after {TSR_JOB_TIMEOUT} seconds.")
    return False


def execute_action(ip: str, action_name: str) -> bool:
    """Run the racadm command for a single IP."""
    try:
        cmd_func = ACTION_MAP.get(action_name)
        if not cmd_func:
            logger.error(f"[{ip}] Unknown action: {action_name}")
            return False

        cmd = cmd_func(ip)
        logger.info(f"[{ip}] Executing {action_name}...")
        if action_name == "tsr_save":
            logger.info(f"[{ip}] Command: {redact_cmd(cmd)}")
        
        # Run command
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=600  # Generous timeout for TSR operations
        )
        output = "\n".join(part for part in [result.stdout.strip(), result.stderr.strip()] if part)

        if result.returncode == 0:
            logger.info(f"[{ip}] Success: {output}")
            if action_name == "tsr_save":
                job_id = extract_job_id(output)
                if not job_id:
                    logger.error(f"[{ip}] TSR Save did not return a Job ID.")
                    return False
                return wait_for_job(ip, job_id)
            return True
        else:
            logger.error(f"[{ip}] Failed: {output}")
            return False

    except Exception as e:
        logger.error(f"[{ip}] Exception: {e}")
        return False


def main():
    parser = argparse.ArgumentParser(description="Unified System Control Script")
    parser.add_argument("action", choices=ACTION_MAP.keys(), help="Action to perform")
    parser.add_argument("ip_file", type=Path, help="Path to the IP list file")
    
    # Optional: Allow passing IP directly via args instead of file (future extension)
    
    args = parser.parse_args()

    ip_file = args.ip_file
    action = args.action

    if not ip_file.is_file():
        logger.error(f"IP file not found: {ip_file}")
        sys.exit(1)

    # Read IPs
    try:
        with ip_file.open("r", encoding="utf-8") as f:
            ips = [line.strip() for line in f if line.strip()]
    except Exception as e:
        logger.error(f"Failed to read IP file: {e}")
        sys.exit(1)

    if not ips:
        logger.warning("No IPs found in file.")
        sys.exit(0)

    logger.info(f"Starting '{action}' for {len(ips)} servers...")
    start_time = time.time()

    success_count = 0
    failed_count = 0

    if action == "tsr_save":
        max_workers = max(1, min(TSR_SAVE_WORKERS, len(ips)))
        logger.info(f"TSR Save parallel workers={max_workers}")
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = [executor.submit(execute_action, ip, action) for ip in ips]
            for future in futures:
                if future.result():
                    success_count += 1
                else:
                    failed_count += 1
    else:
        # Parallel Execution
        # Adjust max_workers as needed. 20 is a safe defaults for I/O bound network tasks.
        with ThreadPoolExecutor(max_workers=20) as executor:
            futures = [executor.submit(execute_action, ip, action) for ip in ips]
            for future in futures:
                if future.result():
                    success_count += 1
                else:
                    failed_count += 1

    elapsed = time.time() - start_time
    logger.info(
        f"Completed '{action}' in {elapsed:.2f} seconds. "
        f"success={success_count}, failed={failed_count}"
    )

    if failed_count:
        sys.exit(1)


if __name__ == "__main__":
    main()