from typing import Callable, List

import requests
from airflow.datasets import Dataset
from airflow.decorators import task, task_group
from airflow.exceptions import AirflowFailException
from airflow.models.taskinstance import TaskInstance
from airflow.sensors.base import PokeReturnValue


def jenkins_job(
    group_id: str,
    jenkins_url: str,
    username: str,
    password: str,
    job_name: str,
    input_fn: Callable[[TaskInstance], list],
    outlets: List[Dataset] = [],
):
    @task
    def input_params(ti: TaskInstance):
        return list(input_fn(ti))

    @task.virtualenv(
        requirements=[
            "jenkinsapi==0.3.13",
        ],
        retries=3,
    )
    def call(jenkins_url, username, password, job_name, params_list):
        import time

        from jenkinsapi.jenkins import Jenkins

        jenkins = Jenkins(jenkins_url, username, password)
        builds = []

        for params in params_list:
            qi = jenkins.get_job(job_name).invoke(build_params=params)
            # we need to get build as soon as possible because qi has a short life time
            while "executable" not in qi.poll():
                print("waiting to build...")
                time.sleep(3)
            build = qi.get_build()
            builds.append(build.baseurl)

        return builds

    @task.sensor(
        poke_interval=30, timeout=6 * 60 * 60, mode="reschedule", outlets=outlets
    )
    def wait(builds: List[str], ti: TaskInstance):
        status = {}
        for baseurl in builds:
            r = requests.get(
                baseurl + "/api/json",
                auth=(username, password),
            )
            if r.status_code != 200:
                raise AirflowFailException(r.text)
            build = r.json()
            print(build)
            if build["building"]:
                return PokeReturnValue(is_done=False)
            status[baseurl] = build["result"]
        ti.xcom_push("builds", builds)
        ti.xcom_push("status", status)
        return PokeReturnValue(is_done=True)

    @task_group(group_id=group_id)
    def group():
        wait(call(jenkins_url, username, password, job_name, input_params()))

    return group()