"""AI Infrastructure Service CLI - Create Job."""

import os
import time
from datetime import datetime, timedelta
from typing import Annotated, Literal, Union
from uuid import UUID

import typer
from pytz import UTC
from rich.progress import Progress
from slices_clientlib_aiis.schema.job import JobStatus

from . import console
from .ai import ai_app
from .ai_context import AIContext
from .ai_helpers import (
    handle_some_exceptions,
)
from .utils import td_format


def _is_active_status(status: JobStatus) -> bool:
    """Return True if the status is active (i.e. the job is still running or pending)."""
    return status in (
        JobStatus.QUEUED,
        JobStatus.ASSIGNED,
        JobStatus.STARTING,
        JobStatus.RUNNING,
        JobStatus.MUSTHALT,
        JobStatus.HALTING,
        JobStatus.CANCELLING,
    )


class SimpleSpinner:
    """A very simple status spinner.

    It creates a spinner by overwriting the same line, using carriage return.
    """

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

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

    def start(self):
        """Start the spinner."""
        self.print_spinner()

    def print_spinner(self):
        """Print the spinner."""
        ch = self.SIMPLE_PROGRESS_CHARS[self.char_index % len(self.SIMPLE_PROGRESS_CHARS)]
        # console.print(f"{ch: <{self.simple_progress_cols}}", end="\r")
        console.print(f"{f'{ch} {self.msg}': <{self.simple_progress_cols}}", end="\r")
        self.char_index += 1

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


def wait_for_job_status(
    ctx: AIContext, job_id: Union[str, UUID], requested_status: JobStatus, max_wait_in_s: int = 60 * 15
):
    """Wait for the job to reach the requested status."""

    start_wait = datetime.now(UTC)
    deadline = start_wait + timedelta(seconds=max_wait_in_s)

    def time_str(show_relative: bool = True):
        if show_relative:
            rel_time = datetime.now(UTC) - start_wait
            rel_time_str = td_format(rel_time)
            needed_spaces = 10 - len(rel_time_str)
            if needed_spaces > 0:
                rel_time_str += " " * needed_spaces
            rel_time_str = f" {rel_time_str} -"
        else:
            rel_time_str = " " * 13
        return f"{datetime.now().astimezone():%Y-%m-%d %H:%M:%S %z} -{rel_time_str}"

    if console.is_dumb_terminal or not console.is_terminal:
        progress = SimpleSpinner(f"[blue]Waiting for job {job_id} to reach status {requested_status.name}...")
        out = console
    else:
        progress = Progress(transient=True)
        progress.start()
        out = progress.console
        progress.add_task(f"[blue]Waiting for job {job_id} to reach status {requested_status.name}...", total=None)
    try:
        last_status = None
        while True:
            if datetime.now(UTC) > deadline:
                out.print(
                    f"[bold red]:x: Waited for job {job_id} to reach status {requested_status.name} "
                    f"for more than {max_wait_in_s} seconds, giving up."
                )
                raise typer.Exit(1)
            job = ctx.obj.api.for_project(ctx.obj.require_project_id()).ai(ctx.obj.site_id).jobs.get(job_id)
            cur_status = job.state.status

            if job.id != job_id:
                # if job_id was a partial ID, it is replaced with a full job UUID
                # if job_id was a str, it is replaced with a UUID
                job_id = job.id

            if cur_status == requested_status:
                out.print(f"[bold green]:sparkles:Job {job_id} reached status {requested_status.name}.")
                return

            if not _is_active_status(cur_status):
                out.print(
                    f"[bold red]:x: Job {job_id} reached final status {cur_status.name}, "
                    f"which is not the requested status {requested_status}."
                )
                raise typer.Exit(1)

            if last_status is None:
                out.print(f"{time_str(show_relative=False)} Job status: {cur_status.name}")
            elif last_status != cur_status:
                out.print(f"{time_str()} Job changed status to: {cur_status.name}")

            if isinstance(progress, SimpleSpinner):
                progress.print_spinner()
            time.sleep(2)
            last_status = cur_status
    finally:
        if progress:
            progress.stop()


@ai_app.command(rich_help_panel="Auxiliary commands")
@handle_some_exceptions
def wait(
    ctx: AIContext,
    job_id: Annotated[str, typer.Argument(help="job ID", show_default=False)],
    target_status: Annotated[
        Literal["RUNNING", "FINISHED"], typer.Argument(help="Target job status to wait for.", case_sensitive=False)
    ] = "FINISHED",
):
    """Wait until the requested status has been reached (or can never be reached)."""
    target_status_value = JobStatus(target_status)
    wait_for_job_status(ctx, job_id, target_status_value)
