import logging
from typing import Callable, List

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

task_logger = logging.getLogger("airflow.task")


def github_action(
    group_id: str,
    repo: str,
    workflow: str,
    input_fn: Callable[[TaskInstance], list],
    ref: str = "main",
    outlets: List[Dataset] = [],
    skip_on_empty: bool = False,
):
    @task.short_circuit(ignore_downstream_trigger_rules=False)
    def call(ti: TaskInstance):
        ti.xcom_push(key="github_action_begin", value=pendulum.now().isoformat())
        token = Variable.get("git_action_bot_token")
        inputs_list = list(input_fn(ti))
        if len(inputs_list) == 0 and skip_on_empty:
            task_logger.info("empty inputs list, skip")
            return False
        task_logger.info("inputs list: %(inputs_list)s", inputs_list)
        for inputs in inputs_list:
            r = requests.post(
                f"https://api.github.com/repos/{repo}/actions/workflows/{workflow}/dispatches",
                json={
                    "ref": ref,
                    "inputs": inputs,
                },
                headers={
                    "Authorization": f"Bearer {token}",
                },
            )

            if r.status_code != 204:
                raise AirflowFailException(r.text)
        return True

    @task.sensor(
        poke_interval=30,
        timeout=6 * 60 * 60,
        mode="reschedule",
        outlets=outlets,
    )
    def wait(ti: TaskInstance):
        batch_size = len(list(input_fn(ti)))
        if batch_size == 0:
            return PokeReturnValue(is_done=True)

        begin_time = ti.xcom_pull(
            key="github_action_begin", task_ids=f"{group_id}.call"
        )
        token = Variable.get("git_action_bot_token")
        r = requests.get(
            f"https://api.github.com/repos/{repo}/actions/workflows/{workflow}/runs",
            params={
                "created": f">{begin_time}",
                "branch": ref,
            },
            headers={
                "Authorization": f"Bearer {token}",
            },
        )
        if r.status_code != 200:
            raise AirflowFailException(r.text)
        data = r.json()
        runs = data["workflow_runs"]
        task_logger.info("workflow_runs: %(runs)s", runs)
        if len(runs) < batch_size:
            return PokeReturnValue(is_done=False)
        elif len(runs) > batch_size:
            raise AirflowFailException("github runs more actions than expected")
        for run in runs:
            if run["status"] == "failure" or run["conclusion"] in [
                "failure",
                "cancelled",
            ]:
                raise AirflowFailException(f"workflow run {run['id']} failed")
            in_process_status = [
                "action_required",
                "in_progress",
                "queued",
                "requested",
                "waiting",
                "pending",
            ]
            if run["status"] in in_process_status:
                return PokeReturnValue(is_done=False)
        return PokeReturnValue(is_done=True)

    @task_group(group_id=group_id)
    def group():
        call() >> wait()

    return group()