import json
import os
import signal
import subprocess
import sys
import time
from threading import Lock, Thread

from loguru import logger

from pm5.argparsers.pm5 import get_app_args

LOCK_FILE = "process_lock.json"  # File to store process IDs to manage process locking

lock = Lock()  # Lock to handle thread synchronization

shutdown = False  # Global flag to indicate if the system is shutting down


# Global list to keep track of running processes and a dictionary to map process IDs to service names
processes = []
process_service_map = {}


# Function to read the ecosystem configuration from a JSON file
def read_config(file_path):
    with open(file_path, "r") as file:
        return json.load(file)


# Function to start a service instance
def start_service(service, instance_id):
    # Construct the command to start the service
    command = (
        [service["interpreter"]]  # Interpreter to run the script (e.g., python)
        + service.get("interpreter_args", [])  # Optional interpreter arguments
        + [service.get("script", "")]  # Script to execute
        + service.get("args", [])  # Optional script arguments
    )

    env = os.environ.copy()  # Copy current environment variables

    env.update(
        {k: str(v) for k, v in service.get("env", {}).items()}
    )  # Update with service-specific environment variables

    cwd = service.get(
        "cwd", os.getcwd()
    )  # Use specified cwd or current working directory if not set

    process = subprocess.Popen(
        command, env=env, cwd=cwd, preexec_fn=os.setsid
    )  # Start the process

    with lock:  # Ensure thread-safe access
        processes.append(process)  # Add process to the list

        process_service_map[process.pid] = service[
            "name"
        ]  # Map process ID to service name

        update_lock_file()  # Update the lock file with running process IDs

    logger.info(
        f"Starting instance {instance_id} of service '{service['name']}' with command: {' '.join(command)} (PID: {process.pid}) in directory: {cwd}"
    )

    return process


# Function to monitor a service and restart it if necessary
def monitor_service(service, process, instance_id):
    global shutdown

    max_restarts = service.get(
        "max_restarts", 0
    )  # Get maximum number of restarts allowed

    restarts = 0  # Initialize restart count

    while restarts <= max_restarts and not shutdown:
        process.wait()  # Wait for the process to finish

        with lock:
            if process in processes:
                processes.remove(process)  # Remove process from the list

                update_lock_file()  # Update the lock file

        if shutdown:
            break

        if process.returncode != 0:  # If the process exited with an error
            logger.error(
                f"Instance {instance_id} of service '{service['name']}' exited with error code {process.returncode}"
            )

        if restarts < max_restarts and service.get("autorestart", False):
            logger.info(
                f"Restarting instance {instance_id} of service '{service['name']}' (Restart {restarts + 1})"
            )

            process = start_service(service, instance_id)  # Restart the service

            restarts += 1  # Increment restart count

        elif restarts == max_restarts:
            if process.returncode != 0:
                logger.warning(
                    f"Instance {instance_id} of service '{service['name']}' has exceeded the maximum number of restarts ({max_restarts}). Stopping all services."
                )

                handle_exit(signal.SIGTERM, None)  # Exit if max restarts exceeded

        else:
            if process.returncode != 0:
                logger.error(
                    f"Instance {instance_id} of service '{service['name']}' has exited with an error and will not be restarted"
                )

            break


# Function to clean up all running processes
def cleanup_processes():
    global shutdown

    if shutdown is True:
        logger.warning("Shutting down all services. Please hold...")

        return

    shutdown = True

    with lock:
        for process in processes:
            service_name = process_service_map.get(process.pid, "Unknown service")

            try:
                logger.info(
                    f"Sending SIGTERM to process group {os.getpgid(process.pid)} of service '{service_name}'"
                )

                os.killpg(
                    os.getpgid(process.pid), signal.SIGTERM
                )  # Send SIGTERM to process group

            except Exception as e:
                logger.info(
                    f"Error sending SIGTERM to process group {os.getpgid(process.pid)} of service '{service_name}': {e}"
                )

        logger.info("Cleaning up services...")

        time.sleep(1)  # Give some time for processes to terminate gracefully

        for process in processes:
            service_name = process_service_map.get(process.pid, "Unknown service")

            if process.poll() is None:  # Check if process is still running
                try:
                    logger.info(
                        f"Sending SIGKILL to process group {os.getpgid(process.pid)} of service '{service_name}'"
                    )

                    os.killpg(
                        os.getpgid(process.pid), signal.SIGKILL
                    )  # Send SIGKILL to process group

                except Exception as e:
                    logger.info(
                        f"Error sending SIGKILL to process group {os.getpgid(process.pid)} of service '{service_name}': {e}"
                    )

        for process in processes:
            try:
                process.wait(timeout=5)  # Wait for process to terminate

            except subprocess.TimeoutExpired:
                service_name = process_service_map.get(process.pid, "Unknown service")

                logger.warning(
                    f"Process {process.pid} of service '{service_name}' did not terminate in time"
                )

        clear_lock_file()  # Clear the lock file

    logger.info("Service cleanup complete.")


# Function to handle script exit
def handle_exit(signum, frame):
    global shutdown

    if shutdown is True:
        return

    logger.info("Terminating services...")

    cleanup_processes()

    with lock:
        if len(processes) > 0:
            for process in processes:
                if process.poll() is None:
                    logger.warning(
                        f"Process {process.pid} is still running, forcing exit..."
                    )

                    os.killpg(
                        os.getpgid(process.pid), signal.SIGKILL
                    )  # Force kill remaining processes

    os._exit(1)  # Force exit with error


# Function to read the lock file containing process IDs
def read_lock_file():
    if os.path.exists(LOCK_FILE):
        with open(LOCK_FILE, "r") as file:
            try:
                return json.load(file)

            except json.JSONDecodeError:
                return []

    return []


# Function to update the lock file with current process IDs
def update_lock_file():
    pids = [process.pid for process in processes]

    with open(LOCK_FILE, "w") as file:
        json.dump(pids, file)


# Function to clear the lock file
def clear_lock_file():
    if os.path.exists(LOCK_FILE):
        os.remove(LOCK_FILE)


# Function to terminate existing processes from the lock file
def terminate_existing_processes():
    pids = read_lock_file()

    for pid in pids:
        try:
            logger.info(f"Terminating existing process group with pid: {pid}")

            os.killpg(pid, signal.SIGTERM)  # Send SIGTERM to process group

        except ProcessLookupError:
            logger.warning(f"No process group found with pid: {pid}")

    clear_lock_file()  # Clear the lock file


# Main function to start the process manager
def main(**kwargs):
    # Print the PID of the current Python process
    logger.debug(f"The process manager process id is: {os.getpid()}")

    # Register signal handlers for graceful exit
    signal.signal(signal.SIGTERM, handle_exit)
    signal.signal(signal.SIGINT, handle_exit)

    # Terminate any existing processes from previous runs
    terminate_existing_processes()

    # Read the ecosystem configuration file
    config = read_config(kwargs["config_file"])
    services = config["services"]

    # Start all service instances
    total_cpus = os.cpu_count()

    services_started = False

    for service in services:
        if service.get("disabled", False):
            logger.info(f"Service '{service['name']}' is disabled. Skipping...")

            continue

        instances = service.get("instances", 1)

        if instances < 0:
            instances = max(
                1, total_cpus - abs(instances)
            )  # Ensure at least 1 instance

        for i in range(instances):
            process = start_service(service, i)

            services_started = True

            if service.get("wait_ready", False):
                thread = Thread(target=monitor_service, args=(service, process, i))

                thread.start()

    if not services_started:
        logger.error("Error: No services to start. Exiting...")

        sys.exit(1)

    # Wait for script termination
    try:
        while True:
            time.sleep(1)

    except KeyboardInterrupt:
        logger.info("Terminating services via keyboard exit...")
        cleanup_processes()
        sys.exit(0)


def app():
    main(**get_app_args())
