from getpass import getuser
from pathlib import Path
from unittest.mock import patch

import pytest
from airflow.models.pool import Pool
from airflow_balancer.testing import pools, variables
from airflow_config import load_config
from airflow_pydantic import SSHOperatorArgs


def test_hydra_config():
    try:
        from airflow_supervisor import SupervisorSSHAirflowConfiguration
    except ImportError:
        pytest.skip("SupervisorSSHAirflowConfiguration not available")

    with (
        patch("supervisor_pydantic.config.supervisor.gettempdir") as p1,
    ):
        with pools():
            pth = Path(__file__).resolve().parent.parent.parent.parent / ".pytest_cache"
            p1.return_value = str(pth)
            cfg = load_config("config", "config")
            assert "balancer" in cfg.extensions
            assert "supervisor" in cfg.extensions

            supervisor_cfg: SupervisorSSHAirflowConfiguration = cfg.extensions["supervisor"]

            # Basic Supervisor checks
            assert supervisor_cfg.command_prefix == ""
            assert supervisor_cfg.working_dir == pth / f"supervisor-{getuser()}-sleep-echo"
            assert supervisor_cfg.inet_http_server.port == "*:9001"
            assert str(supervisor_cfg.supervisorctl.serverurl) == "http://localhost:9001/"

            assert "sleep" in supervisor_cfg.program
            assert "echo" in supervisor_cfg.program
            assert supervisor_cfg.program["sleep"].command == "sleep 1000"
            assert supervisor_cfg.program["echo"].command == 'echo "hello"'

            # SSH Operator checks
            ssh_args: SSHOperatorArgs = supervisor_cfg.ssh_operator_args
            assert ssh_args is not None
            assert ssh_args.cmd_timeout == 63


def test_hydra_config_render():
    with (
        patch("supervisor_pydantic.config.supervisor.gettempdir") as p1,
        patch("supervisor_pydantic.config.supervisor.getuser") as p2,
    ):
        with pools():
            pth = Path(__file__).resolve().parent.parent.parent.parent / ".pytest_cache"
            p1.return_value = str(pth)
            p2.return_value = "user"
            cfg = load_config("config", "config_tasks")

            dag = cfg.dags["example_dag"]
            assert (
                dag.render()
                == """# Generated by airflow-config
from datetime import datetime, time, timedelta
from pathlib import Path

from airflow.models import DAG

from airflow_supervisor.airflow.ssh import SupervisorSSH

with DAG(
    schedule="0 3 * * *",
    start_date=datetime.fromisoformat("2023-01-01T00:00:00"),
    dag_id="example_dag",
    default_args={},
) as dag:
    run = SupervisorSSH(
        cfg={
            "inet_http_server": {"port": "*:9001", "username": None, "password": None},
            "program": {
                "echo": {
                    "command": 'echo "hello"',
                    "autostart": False,
                    "startsecs": 1,
                    "startretries": None,
                    "autorestart": False,
                    "exitcodes": [0],
                    "stopsignal": "TERM",
                    "stopwaitsecs": 30,
                    "stopasgroup": True,
                    "killasgroup": True,
                    "stdout_logfile": Path("/data/echo/output.log"),
                    "stderr_logfile": Path("/data/echo/error.log"),
                    "directory": Path("/data/echo"),
                },
                "sleep": {
                    "command": "sleep 1000",
                    "autostart": False,
                    "startsecs": 1,
                    "startretries": None,
                    "autorestart": False,
                    "exitcodes": [0],
                    "stopsignal": "TERM",
                    "stopwaitsecs": 30,
                    "stopasgroup": True,
                    "killasgroup": True,
                    "stdout_logfile": Path("/data/sleep/output.log"),
                    "stderr_logfile": Path("/data/sleep/error.log"),
                    "directory": Path("/data/sleep"),
                },
            },
            "rpcinterface": {"supervisor": {"rpcinterface_factory": "supervisor.rpcinterface:make_main_rpcinterface"}},
            "config_path": Path("/data/supervisord.conf"),
            "working_dir": Path("/data"),
            "check_interval": timedelta(10.0),
            "check_timeout": timedelta(60.0),
            "runtime": timedelta(360.0),
            "endtime": time(23, 0, 0, 0),
            "maxretrigger": 5,
            "stop_on_exit": False,
            "cleanup": False,
            "restart_on_initial": True,
            "restart_on_retrigger": False,
            "ssh_operator_args": {"cmd_timeout": 63},
        },
        task_id="run",
        dag=dag,
    )
"""
            )

        exec(dag.render())


def test_hydra_config_render_hosts():
    with (
        patch("supervisor_pydantic.config.supervisor.gettempdir") as p1,
        patch("supervisor_pydantic.config.supervisor.getuser") as p2,
    ):
        with pools():
            pth = Path(__file__).resolve().parent.parent.parent.parent / ".pytest_cache"
            p1.return_value = str(pth)
            p2.return_value = "user"
            cfg = load_config("config", "config_tasks_host")

            dag = cfg.dags["example_dag"]
            assert (
                dag.render()
                == """# Generated by airflow-config
from datetime import datetime
from pathlib import Path

from airflow.models import DAG
from airflow_pydantic import Host, Port

from airflow_supervisor.airflow.ssh import SupervisorSSH

with DAG(
    schedule="0 3 * * *",
    start_date=datetime.fromisoformat("2023-01-01T00:00:00"),
    dag_id="example_dag",
    default_args={},
) as dag:
    run = SupervisorSSH(
        cfg={
            "inet_http_server": {"port": "*:9001", "username": None, "password": None},
            "program": {
                "echo": {
                    "command": 'echo "hello"',
                    "autostart": False,
                    "startsecs": 1,
                    "startretries": None,
                    "autorestart": False,
                    "exitcodes": [0],
                    "stopsignal": "TERM",
                    "stopwaitsecs": 30,
                    "stopasgroup": True,
                    "killasgroup": True,
                    "stdout_logfile": Path("/data/echo/output.log"),
                    "stderr_logfile": Path("/data/echo/error.log"),
                    "directory": Path("/data/echo"),
                }
            },
            "rpcinterface": {"supervisor": {"rpcinterface_factory": "supervisor.rpcinterface:make_main_rpcinterface"}},
            "config_path": Path("/data/supervisord.conf"),
            "working_dir": Path("/data"),
        },
        host=Host(name="my-remote-host", password_variable="blerg", password_variable_key="password", os="rhel"),
        port=Port(
            name="my-remote-port",
            host=Host(name="my-remote-host", password_variable="blerg", password_variable_key="password", os="rhel"),
            port=8080,
        ),
        task_id="run",
        dag=dag,
    )
"""
            )

        exec(dag.render())


def test_hydra_config_render_hosts_query():
    with (
        patch("supervisor_pydantic.config.supervisor.gettempdir") as p1,
        patch("supervisor_pydantic.config.supervisor.getuser") as p2,
    ):
        with pools():
            pth = Path(__file__).resolve().parent.parent.parent.parent / ".pytest_cache"
            p1.return_value = str(pth)
            p2.return_value = "user"
            cfg = load_config("config", "config_tasks_host_query")

            dag = cfg.dags["example_dag"]
            assert (
                dag.render()
                == """# Generated by airflow-config
from datetime import datetime
from pathlib import Path

from airflow.models import DAG
from airflow.models.pool import Pool
from airflow_pydantic import Host, Port

from airflow_supervisor.airflow.ssh import SupervisorSSH

with DAG(
    schedule="0 3 * * *",
    start_date=datetime.fromisoformat("2023-01-01T00:00:00"),
    dag_id="example_dag",
    default_args={},
) as dag:
    run = SupervisorSSH(
        cfg={
            "inet_http_server": {"port": "*:9001", "username": None, "password": None},
            "program": {
                "echo": {
                    "command": 'echo "hello"',
                    "autostart": False,
                    "startsecs": 1,
                    "startretries": None,
                    "autorestart": False,
                    "exitcodes": [0],
                    "stopsignal": "TERM",
                    "stopwaitsecs": 30,
                    "stopasgroup": True,
                    "killasgroup": True,
                    "stdout_logfile": Path("/data/echo/output.log"),
                    "stderr_logfile": Path("/data/echo/error.log"),
                    "directory": Path("/data/echo"),
                }
            },
            "rpcinterface": {"supervisor": {"rpcinterface_factory": "supervisor.rpcinterface:make_main_rpcinterface"}},
            "config_path": Path("/data/supervisord.conf"),
            "working_dir": Path("/data"),
        },
        host=Host(
            name="server2",
            username="user1",
            password_variable="myvar",
            password_variable_key="password",
            pool=Pool.create_or_update_pool(
                name="server2", slots=8, description="Balancer pool for host(server2)", include_deferred=False
            ).pool,
            size=8,
            tags=["tag2"],
        ),
        port=Port(
            name="my-app",
            host=Host(
                name="server2",
                username="user1",
                password_variable="myvar",
                password_variable_key="password",
                pool=Pool.create_or_update_pool(
                    name="server2", slots=8, description="Balancer pool for host(server2)", include_deferred=False
                ).pool,
                size=8,
                tags=["tag2"],
            ),
            host_name="server2",
            port=8080,
        ),
        task_id="run",
        dag=dag,
    )
"""
            )
        with pools(Pool()), variables({"username": "user1", "password": "myvar"}):
            exec(dag.render())
