import asyncio
import logging
import signal
import sys
import traceback
from textwrap import dedent

import aiofiles
import aiohttp
from solidlab_perftest_common.agent import (
    AgentCommandResultStatus,
    AgentCommand,
    AgentCommandType,
    AgentCommandResult,
)
from solidlab_perftest_common.util import call_external_command_async

from solidlab_perftest_agent.api_calls import Api
from solidlab_perftest_agent.config import Config

main_loop_active = True


class CommandExecutor:
    # This is a class instead of a method, because to handle timeout, we need to have access to the partial result.
    def __init__(self, command: AgentCommand, config: Config, api: Api):
        self.command = command
        self.config = config
        self.api = api
        self.stdout_s = None
        self.stderr_s = None
        self.err_s = None
        self.debug_s = None
        self.trace = None
        self.return_value = None
        self.called = False

        # It is handy to make the config available to executed command via the environment
        self.env = {
            "PERFTEST_AGENT_API_ENDPOINT": config.api_endpoint,
            "PERFTEST_AGENT_MACHINE_ID": config.machine_id,
            "PERFTEST_AGENT_TEST_ENV_ID": config.test_env_id,
            "PERFTEST_AGENT_AUTH_TOKEN": config.auth_token,
        }
        if self.config.css_config_file:
            self.env["PERFTEST_AGENT_CSS_CONFIG_FILE"] = self.config.css_config_file
        if self.config.css_host:
            self.env["PERFTEST_AGENT_CSS_HOST"] = self.config.css_host
        if self.config.css_port:
            self.env["PERFTEST_AGENT_CSS_PORT"] = str(self.config.css_port)

    def _stdout_handler(self, msg: str):
        self.stdout_s += msg + "\n"

    def _stderr_handler(self, msg: str):
        self.stderr_s += msg + "\n"

    def _error_log_handler(self, msg: str):
        self.err_s += msg + "\n"

    def _debug_log_handler(self, msg: str):
        self.debug_s += msg + "\n"

    def get_status(self, status: AgentCommandResultStatus) -> AgentCommandResult:
        return AgentCommandResult(
            command=self.command,
            status=status,
            return_value=self.return_value,
            error_msg=self.err_s,
            trace=None,
            debug_msg=self.debug_s,
            stdout=self.stdout_s,
            stderr=self.stderr_s,
        )

    async def execute(self) -> bool:
        assert not self.called
        self.called = True
        try:
            if self.command.type == AgentCommandType.SET_CSS_CONFIG:
                assert isinstance(self.command.data, str)
                async with aiofiles.open(self.config.css_config_file, mode="w") as f:
                    await f.write(self.command.data)

                await self.api.report_command_status(
                    self.command, AgentCommandResultStatus.SUCCESS
                )
                return True

            self.stdout_s = ""
            self.stderr_s = ""
            self.err_s = ""
            self.debug_s = ""

            if self.command.type == AgentCommandType.STOP_SERVICE:
                assert isinstance(self.command.data, str)
                self.return_value = await call_external_command_async(
                    command=["/usr/bin/systemctl", "stop", self.command.data],
                    timeout_s=self.command.timeout_s,
                    stdout_handler=self._stdout_handler,
                    stderr_handler=self._stderr_handler,
                    debug_log_handler=self._debug_log_handler,
                    error_log_handler=self._error_log_handler,
                )
            elif self.command.type == AgentCommandType.START_SERVICE:
                assert isinstance(self.command.data, str)
                self.return_value = await call_external_command_async(
                    command=["/usr/bin/systemctl", "restart", self.command.data],
                    timeout_s=self.command.timeout_s,
                    stdout_handler=self._stdout_handler,
                    stderr_handler=self._stderr_handler,
                    debug_log_handler=self._debug_log_handler,
                    error_log_handler=self._error_log_handler,
                )
            elif self.command.type == AgentCommandType.EXEC_BASH:
                assert isinstance(self.command.data, list)
                self.return_value = await call_external_command_async(
                    command=self.command.data,
                    timeout_s=self.command.timeout_s,
                    stdout_handler=self._stdout_handler,
                    stderr_handler=self._stderr_handler,
                    debug_log_handler=self._debug_log_handler,
                    error_log_handler=self._error_log_handler,
                    env=self.env,
                )
            else:
                raise ValueError(f"Unsupported AgentCommandType: {self.command.type}")

            await self.api.report_command_result(
                self.command,
                self.get_status(
                    AgentCommandResultStatus.SUCCESS
                    if self.return_value == 0
                    else AgentCommandResultStatus.FAILURE
                ),
            )
            return True
        except Exception as e:
            self.trace = "".join(traceback.format_exception(None, e, e.__traceback__))
            await self.api.report_command_result(
                self.command, self.get_status(AgentCommandResultStatus.FAILURE)
            )
            return False


async def async_main(config: Config) -> int:
    api = Api(config)
    await api.agent_hello()
    while main_loop_active:
        commands = await api.request_commands()
        if commands:
            for command in commands:
                command_executor = CommandExecutor(command, config, api)
                try:
                    command_success = await asyncio.wait_for(
                        command_executor.execute(), timeout=command.timeout_s
                    )
                    if not command_success:
                        # give server time to digest command failure info
                        await asyncio.sleep(0.5)
                        # Forget the current list of commands
                        # The server will (send them again)/(not remove them) if they don't need to be cancelled.
                        break
                except asyncio.TimeoutError:
                    await api.report_command_result(
                        command,
                        command_executor.get_status(AgentCommandResultStatus.TIMEOUT),
                    )
                    # give server time to digest this info
                    await asyncio.sleep(0.5)
                    # Forget the current list of commands
                    # The server will (send them again)/(not remove them) if they don't need to be cancelled.
                    break

    await api.agent_goodbye()

    return 0


def main() -> int:
    # agent is not meant to be called manually, so it has very strict parameters
    show_help = True
    config = None
    if len(sys.argv) == 7:
        api_endpoint = sys.argv[1].strip()
        show_help = not api_endpoint.startswith("http")
        assert api_endpoint.startswith("http")
        assert "/perftest/" not in api_endpoint
        assert "/artifact/" not in api_endpoint
        if not api_endpoint.endswith("/"):
            api_endpoint += "/"

        machine_id = sys.argv[2].strip()
        show_help = show_help or not machine_id

        try:
            test_env_id = int(sys.argv[3].strip())
            show_help = show_help or not test_env_id
        except ValueError:
            show_help = True
            test_env_id = None

        auth_token = sys.argv[4].strip()
        show_help = show_help or not auth_token

        css_config_file = sys.argv[5].strip() if sys.argv[5].strip() else None
        css_host = sys.argv[6].strip() if sys.argv[6].strip() else None
        css_port = (
            int(sys.argv[7].strip())
            if sys.argv[7].strip() and sys.argv[7].strip().lower != "none"
            else None
        )
        show_help = css_config_file or not auth_token

        config = Config(
            api_endpoint=api_endpoint,
            machine_id=machine_id,
            test_env_id=test_env_id,
            auth_token=auth_token,
            css_config_file=css_config_file,
            css_host=css_host,
            css_port=css_port,
        )

    if show_help:
        print(
            f"""Usage: {sys.argv[0]} <api_endpoint> <machine_id> <test_env_id> <auth_token> <css_config_file> <css_host> <css_port>
   api_endpoint: the URL of the solidlab-perftest-server API endpoint (example: 'https://example.com/perftest_api/v1/')
   machine_id: the ID of the machine this agent runs on (str)
   test_env_id: the ID of the TestEnv (int)
   auth_token: the auth token needed to POST to api_endpoint (str)
   css_config_file: the full path of the css config file (or "none" if not on the server itself) (str)
   css_host: the hostname on which the css can be reached (or "none" on css server itself) (str)
   css_port: the port on which the css can be reached (or "none" on css server itself) (str)
        """,
            file=sys.stderr,
        )
        return 1

    # noinspection PyUnusedLocal
    def signal_handler(sig, frame):
        print(f"{signal.Signals(sig).name} received. Will finishing work.")
        # Stop agent gracefully
        global main_loop_active
        main_loop_active = False

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    retval = asyncio.run(async_main(config))

    return retval


if __name__ == "__main__":
    sys.exit(main())
