#!/usr/bin/env python3

from pathlib import Path
import subprocess
import sys, os, datetime
##import utils.Logger as L
#import threading #, multiprocessing, os
import time

#from django.db import transaction
#from common.models import Command

from Agent import Agent, build_agent, log, extract_parameters
import socket
from common.models import AgentCmd

from src.core.pyros_django.obsconfig.obsconfig_class import OBSConfig

class AgentSST(Agent):
    computer = "MainComputer"
    _previous_dir = ""
    PROJECT_ROOT_PATH = ""
    VENV_PYTHON = ""
    subprocess_dict = {}

    AGENT_SPECIFIC_COMMANDS = [
        "KILL_AGENT",
        "RESTART_AGENT"
    ]

    def __init__(self, name:str=None,sim_computer=None):
        
        super().__init__()
        self.PROJECT_ROOT_PATH = os.environ["PROJECT_ROOT_PATH"]
        if name is None:
            name = self.__class__.__name__
        # if sim_computer:
        #     self.computer = sim_computer
        # else:
        #     self.computer = socket.gethostname()
        WITH_DOCKER = False
        if os.environ.get("WITH_DOCKER"):
            WITH_DOCKER = True
        if WITH_DOCKER:
            VENV_ROOT = ""
            VENV = ""
            VENV_BIN = ""
        else:
            VENV_ROOT = "venv"
            VENV = "venv_py3_pyros"
            VENV_BIN = (
                self.PROJECT_ROOT_PATH
                + os.sep + VENV_ROOT
                + os.sep + VENV
                + os.sep + "bin"
                + os.sep
            )
        self.VENV_PYTHON = VENV_BIN + "python3"
        log.info(f"PC hostname is {self.computer}")
        self.start_agents()

    def start_agents(self,agent_name=None):
        """
        Start all agents or one agent of obs_config

        Args:
            agent_name (_type_, optional): Specific agent name to start. Defaults to None.
        """
        obs_config = self.get_config()
        
        agents = obs_config.get_agents_per_computer(obs_config.unit_name).get(self.computer)
        if agents is None:
            print("Computer not found in obs config")
            exit(1)
        #self.change_dir(self.PROJECT_ROOT_PATH)
        if agent_name:
            agent = agent_name
            # Start a specific agent of obs_config (restart)
            agent_informations = obs_config.get_agent_information(obs_config.unit_name,agent)
            protocol = agent_informations.get("protocol")
            if protocol:
                protocol_folder_abs_path = os.path.join(self.PROJECT_ROOT_PATH, os.path.dirname(protocol))
                
                protocol_script_name = protocol.split("/")[-1]
                if os.path.exists(protocol_folder_abs_path + os.sep + protocol_script_name):
                    cmd = self.VENV_PYTHON +" "+ protocol_folder_abs_path + os.sep + protocol_script_name
                
                    process = subprocess.Popen(f"{cmd}", shell=True, stdout=subprocess.DEVNULL,stderr=subprocess.STDOUT)
                    self.subprocess_dict[agent] = process
                    log.info(f"Start agent {agent} with cmd {cmd}")

        else:
            # Start every agent of obs_config (initial start)
            for agent in agents:
                agent_informations = obs_config.get_agent_information(obs_config.unit_name,agent)
                protocol = agent_informations.get("protocol")
                if protocol:
                    protocol_folder_abs_path = os.path.join(self.PROJECT_ROOT_PATH, os.path.dirname(protocol))
                    
                    protocol_script_name = protocol.split("/")[-1]
                    if os.path.exists(protocol_folder_abs_path + os.sep + protocol_script_name):
                        cmd = self.VENV_PYTHON +" "+ protocol_folder_abs_path + os.sep + protocol_script_name
                        
                        process = subprocess.Popen(f"{cmd}", shell=True, stdout=subprocess.DEVNULL,stderr=subprocess.STDOUT)
                        self.subprocess_dict[agent] = process
                        log.info(f"Start agent {agent} with cmd {cmd}")

    def start_agent(self, agent_name:str):
        """
        Start a specific agent of obs_config (Restart)

        Args:
            agent_name (str): Name of agent to start
        """
        self.start_agents(agent_name)


    def cmd_kill_agent(self, cmd:AgentCmd, agent:str)->None:
        if agent in self.subprocess_dict.keys():
            self.send_cmd_to(agent,"do_exit")

    def cmd_restart_agent(self, cmd:AgentCmd ,agent:str)->None:
        self.start_agent(agent)
        if agent in self.subprocess_dict.keys():
            cmd.set_result(f"Agent {agent} restarted")
            cmd.set_as_processed()
        else:
            cmd.set_result(f"Agent {agent} failed to restart")
            log.debug(f"Agent {agent} failed to restart")

    def force_kill_agent(self, cmd:AgentCmd, agent:str)->None:
        if agent in self.subprocess_dict.get(agent) is not None:
            process = self.subprocess_dict.get(agent)
            process.terminate()
            if process.poll() is not None:
                cmd.set_as_processed()
        else:
            return None
            
    def routine_process_before_body(self):
        try:
            now_time = datetime.datetime.utcnow()
            last_running_commands = AgentCmd.get_commands_sent_by_agent("AgentSST").filter(state="CMD_RUNNING")

            for cmd in last_running_commands:
                last_running_cmd, agent = cmd.full_name.split(" ")[:1]
                if last_running_cmd == "KILL_AGENT" and cmd.is_expired():
                    self.force_killed_agent(agent)
            next_cmd = self._get_next_valid_and_not_running_command()
            if next_cmd is not None:
                cmd, agent = next_cmd.full_name.split(" ")[:1]
                if cmd == "KILL_AGENT":
                    self.cmd_kill_agent(next_cmd,agent)
                elif cmd == "RESTART_AGENT":
                    self.cmd_restart_agent(next_cmd,agent)
            log.info("Check status of process")
            for agent in self.subprocess_dict:
                proc = self.subprocess_dict.get(agent)
            log.info(f"{agent} poll result is {proc.poll()}")
        except KeyboardInterrupt:

            for process in self.subprocess_dict.values():
                print(process)
                process.kill()
            exit(0)


if __name__ == "__main__":

    agent = build_agent(AgentSST)
    agent.run()