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

import logging
from datetime import datetime, timedelta
from typing import Annotated

import pyjson5
import typer
from msgspec import ValidationError
from msgspec.json import decode as json_decode
from msgspec.json import encode as json_encode
from pyjson5 import Json5DecoderException
from pytz import UTC
from slices_clientlib_aiis.schema.job import JobStatus
from slices_clientlib_aiis.schema.new_job import NewJob
from slices_clientlib_core.base.exceptions import (
    ApiException,
    ForbiddenException,
)
from slices_clientlib_core.experiment import Experiment

from . import console, err_console
from .ai import ai_app
from .ai_common import duration_string_to_seconds
from .ai_context import AIContext
from .ai_helpers import (
    handle_some_exceptions,
)
from .ai_ssh_pubkey import SshKeyParamType, get_ssh_keys
from .ai_subcommands_wait import wait_for_job_status
from .utils import get_api_error_detail

logger = logging.getLogger(__name__)


def _submit_new_job_request(ctx: AIContext, new_job: NewJob, wait_run: bool, wait_done: bool, exp: Experiment):
    """Submit NewResourcesRequest to AI service.

    Optionally wait for the job to be running or finished.

    Show job_id in output.
    """
    try:
        created_job = (
            ctx.obj.api.for_project(ctx.obj.require_project_id())
            .for_experiment(exp.id)
            .ai(ctx.obj.site_id)
            .jobs.create(new_job)
        )
    except ForbiddenException as err:
        err_console.print("[bold red]No permission to create job[/] :x:")
        err_console.print(f"[red]{get_api_error_detail(err)}")
        raise typer.Exit(code=1) from None
    except ApiException as err:
        err_console.print("[bold red]Could not create job[/] :x:")
        err_console.print(f"[red]{get_api_error_detail(err)}")
        raise typer.Exit(code=1) from None

    if wait_run or wait_done:
        wait_for_job_status(ctx, created_job.id, JobStatus.RUNNING if wait_run else JobStatus.FINISHED)

    console.print(f"[bold green]:sparkles: Created Job {created_job.id}")


def _get_or_create_experiment(
    ctx: AIContext, project_id: str, experiment_name: str, expires_at: timedelta
) -> Experiment:
    now = datetime.now(UTC)
    slices_client = ctx.obj.api._slices_client
    user_project_id_token = slices_client.get_id_token(
        slices_client._central_api_clients_config.central_hub_services.experiment_authority.audience, project_id
    )
    if experiment_name.startswith("ai.slices-ri.eu"):
        description = f"ai.slices-ri.eu job '{experiment_name[len('ai.slices-ri.eu-') :]}'"
    else:
        description = None
    return ctx.obj.api.experiment.get_or_create(
        experiment_name, project_id, now + expires_at + timedelta(minutes=2), user_project_id_token, description
    )


def get_job_max_end_time_s(job_max_duration: str | None, job_not_after: datetime | str | None) -> int:
    """Return the maximum end time in seconds from now for the job."""
    now = datetime.now(UTC)
    max_duration_str = job_max_duration or "todo default"
    max_duration_s = duration_string_to_seconds(max_duration_str)
    if max_duration_s is None:
        err_console.print(
            f"[bold red]request.scheduling.max_duration is not a valid duration: {max_duration_str}[/] :x:"
        )
        raise typer.Exit(code=1)

    not_after_val = job_not_after or "1 week"
    if isinstance(not_after_val, str):
        not_after_duration_s = duration_string_to_seconds(not_after_val)

        if not_after_duration_s is None:
            # assume datetime string: parse iso date
            try:
                # note: before python 3.11, fromisoformat did not support 'Z' for UTC
                not_after_date = datetime.fromisoformat(not_after_val.replace("Z", "+00:00"))
            except ValueError:
                err_console.print(
                    f"[bold red]request.scheduling.not_after is not a valid date or duration: {not_after_val}[/] :x:"
                )
                raise typer.Exit(code=1) from None
            not_after_s = int((not_after_date - now).total_seconds())
        else:
            not_after_s = not_after_duration_s
    elif isinstance(not_after_val, datetime):
        # This won't occur because msgspec decodes to str. See NewJobRequestScheduling.
        not_after_s = int((not_after_val - now).total_seconds())
    else:
        err_console.print(f"[bold red]request.scheduling.not_after is not valid: {not_after_val!r}[/] :x:")
        raise typer.Exit(code=1)
    return not_after_s + max_duration_s


@ai_app.command(rich_help_panel="Lifecycle commands")
@handle_some_exceptions
def submit(
    ctx: AIContext,
    file: Annotated[typer.FileText, typer.Argument(..., help="File with job definition.", show_default=False)],
    wait_run: Annotated[
        bool, typer.Option("--wait-run", help="Wait till the job is running.", show_default=False)
    ] = False,
    wait_done: Annotated[
        bool, typer.Option("--wait-done", help="Wait till the job is done.", show_default=False)
    ] = False,
    ssh_key: Annotated[
        list[str] | None,
        typer.Option(
            help="(Extra) SSH public key to register for login to the created job.",
            show_default=False,
            envvar="SLICES_SSH_KEY",
            click_type=SshKeyParamType(),
        ),
    ] = None,
    ssh_key_file: Annotated[
        list[typer.FileText] | None,
        typer.Option(
            help="(Extra) SSH public key file for login to the created job.",
            show_default=False,
            envvar="SLICES_SSH_KEY_FILE",
        ),
    ] = None,
    experiment: Annotated[
        str | None, typer.Option("--experiment", help="Name of the experiment (default: job.name).", show_default=False)
    ] = None,
):
    """Submit a job described in a file."""
    request = _read_and_process_job_request_file(ctx, file)
    ssh_authorized_keys = get_ssh_keys(ctx, ssh_key, ssh_key_file)
    combined_ssh_authorized_keys = set(request.request.extra.ssh_pub_keys)
    combined_ssh_authorized_keys.update(ssh_authorized_keys)  # TODO use smart join of SSH keys (ignore comments)
    request.request.extra.ssh_pub_keys = list(combined_ssh_authorized_keys)

    project_id = ctx.obj.require_project_id()
    expires_at = timedelta(
        seconds=get_job_max_end_time_s(
            request.request.scheduling.max_duration, request.request.scheduling.not_after or "5 min"
        )
    )
    experiment_name = experiment if experiment else f"ai.slices-ri.eu {request.name[:30]}"
    exp = _get_or_create_experiment(ctx, project_id, experiment_name, expires_at)

    _submit_new_job_request(ctx, request, wait_run, wait_done, exp)


def _read_and_process_job_request_file(_ctx: AIContext, file: typer.FileText) -> NewJob:
    """Return NewJob request."""
    try:
        # We support JSON5 instead of plain JSON, which allows trailing comma's, comments, etc.
        # See https://json5.org/
        request_dict = pyjson5.load(file)  # pyright: ignore [reportArgumentType] FileText is compatible enough it seems
    except Json5DecoderException as e:
        err_console.print(f"[red]Errors parsing JSON5 in request file: {e}")
        raise typer.Exit(code=1) from e
    except OSError as e:
        err_console.print(f"[red]Error reading request file: {type(e)}: {e}")
        raise typer.Exit(code=1) from e

    request_json = json_encode(request_dict)

    logger.debug("Processed request json: %s", request_json)

    try:
        res = json_decode(request_json, type=NewJob)
        if res.request.scheduling.interactive:
            # interactive jobs may not have not_after set
            res.request.scheduling.not_after = None
        logger.debug("Processed NewJob: %s", res)
    except ValidationError as e:
        err_console.print("[red]Errors in request file:")
        err_console.print(f"[red]{str(e)}")
        raise typer.Exit(code=1) from e
    else:
        return res


# TODO: add
# run_interactive
