#TOTEST Cleared
import os
import subprocess
import time
from typing import Union
from variables import PATH_TO_RUN_SSH_FUNCTION,FUNCTION_REQUIRED_AGENT_TYPES,PATH_TO_VM_LOCAL_DECODER,PATH_TO_VM_LOCAL_RULES,PATH_TO_BASH_SCRIPTS,PATH_TO_RUN_SSH_SCRIPT
from agents import Agents 

class Remote(Agents):

    """
    Utility class for managing remote operations.
    """
    def __init__(self, remote_name: str,VM_conf_path:str = "/var/ossec/etc/ossec.conf",**kwargs):
        """
        Initializes Remote with a vagrant instance, as Remote cannot exist without it.
        
        Args:
            VM_conf_path (str): Path to the VM configuration file.
            remote_name (str): Name of the vagrant virtual machine
        """
        super().__init__(**kwargs)
        
        self.remote_name = remote_name  # Store the Agents instance as a dependency

        # HERE VM PATHS
        self.VM_conf_path = VM_conf_path
        self.VM_local_rules_path = PATH_TO_VM_LOCAL_RULES
        self.VM_local_decoder_path = PATH_TO_VM_LOCAL_DECODER

        assert isinstance(self.remote_name, str) and self.remote_name.strip(), "Remote host must be a valid non-empty string."

    
    def synchronize_with_VM(self, conf_path: str, vm_conf_path: str = None, from_vm: bool = False):
        """
        Synchronize a file with a remote VM.

        Args:
            conf_path (str): Local relative path (from /src) file to synchronize.
            vm_conf_path (str): Absolute path on the VM to the file targeted (default is ossec file).
            from_vm (bool): If True, synchronize files from VM to host; otherwise, from host to VM.

        Raises:
            AssertionError: If inputs are invalid.
            RuntimeError: If synchronization fails.
        """

        vm_conf_path = vm_conf_path or self.VM_conf_path
        self.logger.debug(f"Starting VM synchronization for file: {conf_path} to {self.remote_name}:{vm_conf_path}.")

        # Assertions for input validation
        assert isinstance(conf_path, str) and os.path.isfile(conf_path), f"Local configuration file {conf_path} does not exist."
        assert isinstance(self.remote_name, str) and self.remote_name.strip(), "remote_host must be a valid non-empty string."
        assert isinstance(vm_conf_path, str) and vm_conf_path.strip(), "VM_conf_path must be a valid non-empty string."

        conf_path_remote = os.path.abspath(conf_path)

        # Find the part of the path starting with '/vagrant'
        if "/vagrant" in conf_path_remote:
            conf_path_remote = conf_path_remote[conf_path_remote.index("/vagrant"):]

        # Determine the direction of synchronization
         # Determine the direction of synchronization
        if from_vm:
            # Synchronize from VM to host using cp command
            remote_command = ['cp', vm_conf_path, f'/vagrant/src/{conf_path}']
            result = self.run_remote_command(remote_command)
        else:
            # Synchronize from host to VM
            command = [PATH_TO_RUN_SSH_FUNCTION, self.remote_name, 'synchronize_file', conf_path_remote, vm_conf_path]

            result = self.run_command(command)

        if result.returncode != 0:
            self.logger.error(f"Synchronization failed: {result.stderr}")
            raise RuntimeError(f"Synchronization failed with error: {result.stderr}")
        self.logger.debug(f"Synchronization successful: {result.stdout}")


    def run_remote_command(self, command: list[str]):
        """
        Execute a command on the remote host.

        Args:
            command (list[str]): The command to execute as a list of arguments.

        Returns:
            subprocess.CompletedProcess: The result of the command execution.
        """
        self.logger.debug(f"Executing remote command on {self.remote_name}: {' '.join(command)}")
        # Build the SSH command dynamically
        ssh_command = ['ssh', f'root@{self.remote_name}'] + command
        result = self.run_command(ssh_command)
        self.logger.debug(f"Command executed successfully: {result.stdout}")
        return result

    def check_remote_file(self, remote_file_path: str):
        """
        Check if a file exists on the remote VM.

        Args:
            remote_file_path (str): Path to the file on the remote VM.

        Returns:
            bool: True if the file exists, False otherwise.
        """
        self.logger.debug(f"Checking existence of file on remote host: {remote_file_path}")

        result = self.run_remote_command(["test", "-f", remote_file_path])
        return result.returncode == 0

    def remove_remote_file(self, remote_file_path: str):
        """
        Remove a file from the remote VM if it exists.

        Args:
            remote_file_path (str): Path to the file on the remote VM.
        """
        if self.check_remote_file(remote_file_path):
            self.logger.debug(f"Removing file on remote host: {remote_file_path}")

            self.run_remote_command(["rm", "-f", remote_file_path])
            self.logger.debug(f"File removed successfully from remote host: {remote_file_path}")

    def run_function_on_remote_host(self, function: str, *args: list):
        """
        Run a Python or bash function on a remote host via SSH.

        This function constructs a command to execute a Python function on a remote host.
        It uses a Bash script to handle the SSH connection and execution of the Python function.

        Args:
            function (str): The remote_name of the Python or bash function to be executed (#NOTE : by default should be in utils.py or utils.sh!).
            *args: Additional arguments to be passed to the Python function , can be a string or a list or string.

        Returns:
        None
        """
        # Check if function can be run for the provide agent_type
        if function in FUNCTION_REQUIRED_AGENT_TYPES.keys():
            assert FUNCTION_REQUIRED_AGENT_TYPES[function]['agent_type'] == self.agent_type, self.logger.error(f"Agent of type {self.agent_type} can't use this function")
        
        # Construct the command to be executed
        # The command includes the path to the Bash script, the remote host, the Python function,
        # the source library path, and the arguments string
        command = [
            PATH_TO_RUN_SSH_FUNCTION,  # Path to the Bash script
            self.remote_name,                                # Remote host
            function,                            # Python function to run
        ]

        if args :
            command.extend(args)

        # Run the command using the run_command function
        # Assuming run_command is a function that takes a list of command arguments and executes them
        self.run_command(command)

        return


    def run_bash_script_on_remote_host(self, bash_script_name: str, *args: list):
        """
        Run a Bash script on a remote host via SSH.

        This function constructs a command to execute a Bash script on a remote host.
        It uses SSH to handle the connection and execution of the Bash script.

        Args:
            bash_script_name (str): The name of the Bash script to be executed.
            *args: Additional arguments to be passed to the Bash script.

        Returns:
            None
        """
            
    
        # Ensure the correct directory
        target_dir = 'src'
        if os.path.basename(os.getcwd()) != target_dir:
            os.chdir(target_dir)

        # Check if file is present : 
        
        path_to_bash_script = os.path.join(PATH_TO_BASH_SCRIPTS,bash_script_name)
        
        if os.path.isfile(path_to_bash_script):
            self.logger.debug(f"Script found at: {path_to_bash_script}")
        else:
            self.logger.error(f"Error: Script '{bash_script_name}' not found in '{PATH_TO_BASH_SCRIPTS}'")

        # Construct the command to be executed
        # The command includes the path to the Bash script, the remote host, the Bash script,
        # and the arguments string

        command = [
            PATH_TO_RUN_SSH_SCRIPT,                  # Path to bash script runner
            self.remote_name,                             # Remote host
            path_to_bash_script,                     # Absolut path to bash script
        ]

        # Add the Bash script content and arguments
        command.extend(args)

        # Run the command using subprocess
        self.run_command(command)

        return

    def restart(self, elements: Union[list[str], str]):
        """
        Restart one or multiple services to apply configuration changes on the remote machine

        Args:
            elements (str or list of str): remote_name(s) of the service(s) to restart.
            remote_name (str): remote machine targeted to restart the services.

        INFO: Adjust time.sleep regarding the element to restart, it can interfere with other commands !
        """
        if isinstance(elements, str):
            elements = [elements]  # Convert single element to a list

        for element in elements:
            RESTART_CMD = ["systemctl", "restart", element]
            try:
                self.run_remote_command(RESTART_CMD)
                time.sleep(2)
                self.logger.debug(f"{element} restarted successfully.")
            except subprocess.CalledProcessError as e:
                self.logger.error(f"Failed to restart {element}: {e}")
