"""AI Infrastructure Service CLI - action commands."""

from datetime import datetime, timedelta
from typing import Annotated

import typer
from pytz import UTC
from slices_clientlib_core.base.exceptions import ConflictException

from . import console, err_console
from .ai import ai_app
from .ai_common import duration_string_to_seconds, get_ai_job, validate_duration_option
from .ai_context import AIContext
from .ai_helpers import (
    handle_some_exceptions,
)
from .ai_subcommands_create import get_job_max_end_time_s


@ai_app.command("cancel", rich_help_panel="Lifecycle commands")
@handle_some_exceptions
def cancel_job(
    ctx: AIContext,
    job_id: Annotated[str, typer.Argument(help="Job ID.", show_default=False)],
):
    """Cancel a job."""

    project_id = ctx.obj.get_current_project_id()
    if not project_id:
        raise typer.Exit(0)

    job = get_ai_job(ctx, job_id)

    ctx.obj.api.for_project(ctx.obj.require_project_id()).ai(ctx.obj.site_id).jobs.cancel(job.id)

    console.print(f"[green] :recycle: Cancelled Job {job.id}")


@ai_app.command("rm", rich_help_panel="Lifecycle commands")
@handle_some_exceptions
def delete_job(
    ctx: AIContext,
    job_id: Annotated[str, typer.Argument(help="Job ID.", show_default=False)],
):
    """Delete a job."""

    project_id = ctx.obj.get_current_project_id()
    if not project_id:
        raise typer.Exit(0)

    job = get_ai_job(ctx, job_id)

    ctx.obj.api.for_project(ctx.obj.require_project_id()).ai(ctx.obj.site_id).jobs.delete(job.id)

    console.print(f"[green] :recycle: Deleted Job {job.id}")


@ai_app.command("halt", rich_help_panel="Lifecycle commands")
@handle_some_exceptions
def halt_job(
    ctx: AIContext,
    job_id: Annotated[str, typer.Argument(help="Job ID.", show_default=False)],
):
    """Halt a job. Halted jobs may be automatically re-QUEUED later."""

    project_id = ctx.obj.get_current_project_id()
    if not project_id:
        raise typer.Exit(0)

    job = get_ai_job(ctx, job_id)

    ctx.obj.api.for_project(ctx.obj.require_project_id()).ai(ctx.obj.site_id).jobs.halt(job.id)

    console.print(f"[red] :stop_sign: [green] Requested Job halt for {job.id}")


@ai_app.command("hold", rich_help_panel="Lifecycle commands")
@handle_some_exceptions
def hold_job(
    ctx: AIContext,
    job_id: Annotated[str, typer.Argument(help="Job ID.", show_default=False)],
):
    """Hold a QUEUED job, preventing it from running."""

    project_id = ctx.obj.get_current_project_id()
    if not project_id:
        raise typer.Exit(0)

    job = get_ai_job(ctx, job_id)

    ctx.obj.api.for_project(ctx.obj.require_project_id()).ai(ctx.obj.site_id).jobs.hold(job.id)

    console.print(
        f"[green] :pause_button: Holding Job {job.id} (It will not run until you release it to the queue again.)"
    )


@ai_app.command("requeue", rich_help_panel="Lifecycle commands")
@handle_some_exceptions
def requeue_job(
    ctx: AIContext,
    job_id: Annotated[str, typer.Argument(help="Job ID.", show_default=False)],
):
    """Requeue/Release a held job, QUEUEing it, so it waits to run."""

    project_id = ctx.obj.get_current_project_id()
    if not project_id:
        raise typer.Exit(0)

    job = get_ai_job(ctx, job_id)

    ctx.obj.api.for_project(ctx.obj.require_project_id()).ai(ctx.obj.site_id).jobs.requeue(job.id)

    console.print(
        f"[green] :play_button: Released Job {job.id} to the queue. (Where it will wait for its turn to run.)"
    )


def validate_duration_or_date_option(ctx: typer.Context, _param: typer.CallbackParam, value: str) -> datetime | None:
    """Validate that the value is either a duration (e.g. 1h30m) or a date (e.g. 2024-12-31T23:59:59Z).

    This always returns a date.
    If the input is a duration, it returns the current time + duration.
    If the input is a date, it returns the date.
    """

    if value is None:
        return None

    if ctx.resilient_parsing:
        return None

    try:
        # if it's a valid duration, it is OK.
        validate_duration_option(ctx, _param, value)
        s: int = duration_string_to_seconds(value)  # pyright: ignore [reportAssignmentType] cannot be None
        return datetime.now(UTC) + timedelta(seconds=s)
    except ValueError:
        pass

    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError:
        raise ValueError(
            f"Invalid value '{value}'. "
            f"Must be either a duration (e.g. 1h, 30m, ...) or a date (e.g. 2024-12-31T23:59:59Z)."
        ) from None


def validate_duration_option_or_none(ctx: typer.Context, param: typer.CallbackParam, value: str | None) -> str | None:
    """Validate that the value is either None or a duration (e.g. 1h, 30m, ...)."""

    if value is None:
        return None

    return validate_duration_option(ctx, param, value)


@ai_app.command("modify", rich_help_panel="Auxiliary commands")
@handle_some_exceptions
def modify_job(
    ctx: AIContext,
    job_id: Annotated[str, typer.Argument(help="Job ID.", show_default=False)],
    min_duration: Annotated[
        str | None,
        typer.Option("--min-duration", help="New minDuration for the job.", callback=validate_duration_option_or_none),
    ] = None,
    max_duration: Annotated[
        str | None,
        typer.Option("--max-duration", help="New maxDuration for the job.", callback=validate_duration_option_or_none),
    ] = None,
    not_after: Annotated[
        datetime | None,
        typer.Option("--not-after", help="New notAfter for the job.", callback=validate_duration_or_date_option),
    ] = None,
):
    """Change maxDuration, minDuration or notAfter of a job."""
    project_id = ctx.obj.require_project_id()

    if min_duration is None and max_duration is None and not_after is None:
        err_console.print(
            "[red] No changes specified. Use --min-duration, --max-duration or --not-after to specify changes."
        )
        raise typer.Exit(1)

    job = get_ai_job(ctx, job_id)

    experiment_id = job.owner.experiment_id
    if experiment_id is None:
        err_console.print(
            f"[red] :x: Job {job_id} does not have an experiment_id, cannot modify scheduling parameters."
        )
        raise typer.Exit(1)

    modify_job_end = max_duration is not None or not_after is not None

    required_expires_at = datetime.now(UTC) + timedelta(
        seconds=get_job_max_end_time_s(
            max_duration or job.request.scheduling.max_duration, not_after or job.request.scheduling.not_after
        )
    )

    experiment = ctx.obj.api.experiment.get(experiment_id)
    if modify_job_end and experiment.expires_at < required_expires_at:
        # Extend expiration time if the existing experiment expires before required expire time
        ctx.obj.api.experiment.update(
            experiment_id,
            project_id,
            expires_at=required_expires_at,
            description=experiment.description,
        )
        console.print("[green] :clock10: Extended experiment.")

    try:
        if max_duration is not None:
            ctx.obj.api.for_project(project_id).for_experiment(experiment_id).ai(
                ctx.obj.site_id
            ).jobs.update_max_duration(job_id, max_duration)
            console.print("[green] :pencil: Modified maxDuration.")

        if min_duration is not None:
            ctx.obj.api.for_project(project_id).for_experiment(experiment_id).ai(
                ctx.obj.site_id
            ).jobs.update_min_duration(job_id, min_duration)
            console.print("[green] :pencil: Modified minDuration.")

        if not_after is not None:
            ctx.obj.api.for_project(project_id).for_experiment(experiment_id).ai(ctx.obj.site_id).jobs.update_not_after(
                job_id, not_after
            )
            console.print("[green] :pencil: Modified notAfter.")
    except ConflictException as e:
        err_console.print(f"[red] :x: Could not modify job: {e}")
        raise typer.Exit(1) from None
