"""Utility to follow progress of a BI task."""

import os
from time import sleep

from rich.progress import BarColumn, Progress, TaskID, TaskProgressColumn, TextColumn
from slices_clientlib_bi.bi_task import BiTaskClient, TaskStatus

from slices_cli_bi import console

MAX_PROGRESS_UPDATES_PER_REQUEST = 40  # stay within the limit of how long a GET-url can become

progress_columns = [
    TextColumn("[progress.description]{task.description}"),
    BarColumn(),
    TaskProgressColumn(),
]


class SimpleProgressPrint:
    """A very simple status printer.

    It keeps overwriting the same line, by using carriage return.
    """

    SIMPLE_PROGRESS_CHARS = ["/", "-", "\\", "|"]

    def __init__(self, single_task: bool):
        try:
            self.simple_progress_cols = int(os.getenv("COLUMNS", "30"))
        except ValueError:
            self.simple_progress_cols = 30
        self.char_index = 0
        self.single_task = single_task

    def print(self, status: TaskStatus | None):
        """Print a status line."""
        ch = self.SIMPLE_PROGRESS_CHARS[self.char_index % len(self.SIMPLE_PROGRESS_CHARS)]
        if self.single_task and status:
            perc = (status.completed or 0) * 100.0 / (status.total or status.completed or 100)
            console.print(f"{f'{perc}% {ch} {status.description}': <{self.simple_progress_cols}}", end="\r")
        else:
            console.print(f"{ch: <{self.simple_progress_cols}}", end="\r")
        self.char_index += 1

    def clear(self):
        """Clear the line."""
        console.print(" " * self.simple_progress_cols, end="\r")


def follow_progress(
    bi_tasks_api: BiTaskClient,
    task_ids: list[str],
    task_description_prefixes: list[str],
    *,
    transient=False,
) -> dict[str, TaskStatus]:
    """Periodically query BI tasks API to check on task while showing progress bar."""
    assert len(task_ids) == len(task_description_prefixes)
    unique_task_id_len = len(set(task_ids))

    # If the console is not a terminal, or a dumb terminal, progress won't do anything until it is done.
    if console.is_dumb_terminal or not console.is_terminal:
        simple_progress = SimpleProgressPrint(unique_task_id_len == 1)
    else:
        simple_progress = None

    finished_tasks: dict[str, TaskStatus] = {}
    a_task: TaskStatus | None = None
    with Progress(*progress_columns, transient=transient, console=console) as progress:
        bi_task_ids_to_progress_t_ids: dict[str, list[TaskID]] = {}
        bi_task_ids_to_description_prefixes: dict[str, list[str]] = {}

        for task_id, task_description_prefix in zip(task_ids, task_description_prefixes, strict=True):
            progress_t_id = progress.add_task(task_description_prefix, total=None)
            bi_task_ids_to_progress_t_ids.setdefault(task_id, []).append(progress_t_id)
            bi_task_ids_to_description_prefixes.setdefault(task_id, []).append(task_description_prefix)

        while not progress.finished:
            assert len(finished_tasks) < unique_task_id_len

            tasks_to_query = list({task_id for task_id in task_ids if task_id not in finished_tasks})
            bi_task_status: list[TaskStatus] = []
            while tasks_to_query:
                bi_task_status.extend(
                    bi_tasks_api.get_tasks(
                        task_ids=tasks_to_query[0:MAX_PROGRESS_UPDATES_PER_REQUEST],  # type: ignore
                    )
                )
                tasks_to_query = tasks_to_query[MAX_PROGRESS_UPDATES_PER_REQUEST:]

            for task_status in bi_task_status:
                a_task = task_status
                task_finished = task_status.status in ("finished", "failed")
                task_success = task_status.status == "finished"
                progress_t_ids = bi_task_ids_to_progress_t_ids[task_status.id]
                for progress_t_id in progress_t_ids:
                    completed = task_status.completed or 0
                    total = task_status.total or 100

                    # The API only use task.status for finished. task.completed and task.total are only informative.
                    # But rich.progress uses completed >= total to detect finished.
                    #
                    # So to use rich.progress we need to auto fix bad completed counts coming from the API:
                    # - If status=finished but `completed` is not `total`, set `completed` to `total`
                    #   This prevents finished from never being reached because backend code forgot to set 100%.
                    # - If not status=finished but `completed` equal to `total`, set `completed` to `total-1`
                    #   This one prevents the informative 100% from being considered as finished
                    #   before the task is really finished.
                    completed = min(completed, total - 1) if not task_finished else task_status.total

                    progress.update(progress_t_id, total=task_status.total, completed=completed)

                    task_description = task_status.description
                    if task_status.failure_reason:
                        # make the failure reason message show up here as well
                        if task_description:
                            task_description = task_description + " # " + task_status.failure_reason.message
                        else:
                            task_description = task_status.failure_reason.message
                    if task_description:
                        for description_prefix in bi_task_ids_to_description_prefixes[task_status.id]:
                            progress.update(
                                progress_t_id,
                                description=f"{description_prefix}: {task_description}",
                                visible=not task_success,
                            )

                if task_finished:
                    finished_tasks[task_status.id] = task_status

            if not progress.finished:
                if simple_progress:
                    simple_progress.print(a_task)
                sleep(0.5)
            elif simple_progress:
                # cleanup simple busy view
                simple_progress.clear()

    return finished_tasks


def follow_one_progress(
    bi_tasks_api: BiTaskClient,
    task_id,
    task_description_prefix,
    *,
    transient=False,
):
    """Periodically query BI tasks API to check on task while showing progress bar."""
    follow_progress(bi_tasks_api, [task_id], [task_description_prefix], transient=transient)
