"""AI Infrastructure Service CLI - Query jobs."""

import re
from datetime import datetime, timedelta
from enum import Enum
from typing import TYPE_CHECKING, Annotated, Any

import typer
from msgspec.json import encode as json_encode
from msgspec.yaml import encode as yaml_encode
from pytz import UTC
from rich.table import Column, Table
from slices_clientlib_aiis.schema.job import (
    Job,
    JobPortMapping,
    JobStatus,  # noqa: TCH002 required outside type checking! ruff is wrong.
)
from slices_clientlib_aiis.schema.job_debug_logs import JobDebugLogType

from . import console
from .ai import ai_app
from .ai_common import get_ai_job, get_requested_ai_jobs, project_id_to_name, status_to_color, user_id_to_name
from .ai_context import AIContext
from .ai_helpers import (
    handle_some_exceptions,
)
from .utils import friendly_datetime, td_format

if TYPE_CHECKING:
    from collections.abc import Collection


class ListResourcesFormat(Enum):
    """List Resources Formats."""

    table = "table"
    json = "json"
    yaml = "yaml"


class JobOwnerSelection(Enum):
    """Show jobs with which owner."""

    user = "user"
    project = "project"
    all = "all"


def _get_cluster_id(job: Job) -> int | None:
    """Get cluster ID from job, or return None if not available."""
    if job.state.resources is not None:
        return job.state.resources.cluster_id
    if job.state.scheduling is not None and job.state.scheduling.assigned_cluster_id is not None:
        return job.state.scheduling.assigned_cluster_id
    if job.request.resources.cluster_id is not None:
        return job.request.resources.cluster_id
    return None


@ai_app.command("list", rich_help_panel="Information commands")
@handle_some_exceptions
def list_jobs(
    ctx: AIContext,
    job_ids: Annotated[list[str] | None, typer.Argument(help="Zero, one or more job IDs.", show_default=False)] = None,
    include_pending: Annotated[
        bool, typer.Option("--show-pending/--hide-pending", help="Show/hide QUEUED and ONHOLD jobs.")
    ] = True,
    include_running: Annotated[
        bool, typer.Option("--show-running/--hide-running", help="Show/hide jobs in any running state.")
    ] = True,
    include_finished: Annotated[
        bool, typer.Option("--show-finished/--hide-finished", help="Show/hide jobs that have stopped.")
    ] = True,
    include_deleted: Annotated[
        bool, typer.Option("--show-deleted/--hide-deleted", help="Show/hide deleted jobs.")
    ] = False,
    format: Annotated[
        ListResourcesFormat,
        typer.Option("--format", "-f", help="Format in which the resources should be listed."),
    ] = "table",  # type: ignore
    filter_owner: Annotated[
        JobOwnerSelection,
        typer.Option("--owner-filter", "-o", help="Show only your own jobs, those in your project, or all jobs."),
    ] = JobOwnerSelection.project,
    limit: Annotated[
        int,
        typer.Option("--limit", help="Show only a limited number of jobs.", min=1, max=500),
    ] = 20,
):
    """List your jobs."""

    jobs_to_show: Collection[Job] = get_requested_ai_jobs(
        ctx,
        job_ids,
        filter_user=filter_owner == JobOwnerSelection.user,
        filter_project=filter_owner == JobOwnerSelection.project,
        include_pending=include_pending,
        include_running=include_running,
        include_finished=include_finished,
        include_deleted=include_deleted,
    )

    total_job_count = len(jobs_to_show)
    if total_job_count > limit:
        limited_jobs_to_show = []
        for j in jobs_to_show:
            limited_jobs_to_show.append(j)
            if len(limited_jobs_to_show) == limit:
                break
        jobs_to_show = limited_jobs_to_show

    if format == ListResourcesFormat.table:
        cols = [
            "ID",
            "Name",
            "Status",
            "Created At",
            "Cluster ID",
            "User",
            "Project",
        ]
        title = "Jobs"
        if filter_owner == JobOwnerSelection.user:
            title = "Your Jobs"
        elif filter_owner == JobOwnerSelection.project:
            title = f"Jobs in Project {ctx.obj.require_project_friendly_name()}"
        title += f" ({len(jobs_to_show)}/{total_job_count})"
        table = Table(*cols, title=title)
        for j in jobs_to_show:
            cluster_id = _get_cluster_id(j)
            rows = [
                str(j.id),
                j.name,
                f"{status_to_color(j.state.status)}{j.state.status.name}",
                friendly_datetime(j.state.event_times.created.astimezone()),
                str(cluster_id) if cluster_id else "-",
                user_id_to_name(ctx, j.owner.user_id),
                project_id_to_name(ctx, j.owner.project_id),
            ]
            table.add_row(*rows)
        console.print(table)

    elif format == ListResourcesFormat.json:
        console.print_json(json_encode(list(jobs_to_show)).decode())
    elif format == ListResourcesFormat.yaml:
        console.print(yaml_encode(list(jobs_to_show)).decode())


class ShowResourceFormat(Enum):
    """Show Resource Formats."""

    table = "table"
    json = "json"
    yaml = "yaml"


DOCKER_IMAGE_PATTERN = (
    r"^((?P<username>[\w\-+_.]+):(?P<password>[^@ ]+)@)?"  # passwords with space or @ are not supported
    r"(?P<repository>[\w.\-_]+(?::\d+|)\/)?"
    r"(?P<image>[a-z0-9.\-_]+(?:\/[a-z0-9.\-_]+|)+)(:(?P<tag>[\w.\-_]{1,127}))?$"
)

DOCKER_IMAGE_REGEX = re.compile(DOCKER_IMAGE_PATTERN)


@ai_app.command("show", rich_help_panel="Job interaction commands")
@handle_some_exceptions
def show_job(  # noqa: PLR0915, C901
    ctx: AIContext,
    job_id: Annotated[str, typer.Argument(help="job ID", show_default=False)],
    format: Annotated[
        ShowResourceFormat,
        typer.Option("--format", "-f", help="Format in which the job should be shown."),
    ] = "table",  # type: ignore
):
    """Show Job Details."""

    job: Job = get_ai_job(ctx, job_id)

    if format == ShowResourceFormat.table:

        def handle_none(i: Any | None) -> Any:
            if i is None:
                return "-"
            return i

        def image_nopass(image: str | None) -> str | None:
            """Docker image with the password removed."""
            if not image:
                return None
            match = DOCKER_IMAGE_REGEX.match(image)

            if match and match.group("password"):
                return (
                    f"{match.group('username')}:XXXX@"
                    f"{match.group('repository')}{match.group('image')}:{match.group('tag')}"
                )

            return image

        def command_as_str(command: str | list | None) -> str | None:
            """Command as a string."""
            if not command:
                return None
            if isinstance(command, str):
                return command
            return " ".join(command)

        def _ssh_port_command(port: int | None) -> str:
            if port and port != 22:  # noqa: PLR2004
                return f" -oPort={port}"
            return ""

        def _port_mapping_format(port_mappings: list[JobPortMapping]) -> str:
            try:
                res = [f"{port_mapping.container_port} -> {port_mapping.host_port}" for port_mapping in port_mappings]
                return ", ".join(res)
            except:  # noqa: E722
                return f"Error processing port mappings: {port_mappings}"

        table = Table.grid(Column(header="Key", justify="right", style="bold"), "Value", padding=(0, 1, 0, 1))

        table.add_row("Job ID:", str(job.uuid))
        table.add_row("Name:", handle_none(job.name))
        table.add_row("Description:", handle_none(job.description))
        table.add_row("Project:", handle_none(job.owner.project_id) if job.owner else None)
        table.add_row("User ID:", handle_none(job.owner.user_id) if job.owner else None)
        table.add_row("Docker image:", handle_none(image_nopass(job.request.docker.image)))
        table.add_row("Command:", handle_none(command_as_str(job.request.docker.command)))

        table.add_row(
            "Status:",
            f"{status_to_color(job.state.status)}{job.state.status.name}" if job.state and job.state.status else "-",
        )

        if job.state == JobStatus.FAILED:
            table.add_row("", f"To print more info about FAILED: slices ai debug {job.uuid}")

        table.add_row(
            "Cluster ID:",
            str(handle_none(job.state.resources.cluster_id)) if job.state and job.state.resources else "-",
        )
        # No longer useful info:
        # table.add_row(
        #     "Worker ID:",
        #     str(handle_none(job.state.resources.worker_id)) if job.state and job.state.resources else "-"
        # )
        table.add_row(
            "Machine ID:", handle_none(job.state.resources.machine_id) if job.state and job.state.resources else "-"
        )
        if job.state.status == JobStatus.QUEUED:
            table.add_row(
                "Queued because:",
                handle_none(" \n".join(job.state.scheduling.queued_explanations))
                if job.state and job.state.scheduling
                else "-",
            )

        if job.state and job.state.resources:
            if job.state.resources.slave_name:
                table.add_row(
                    "Worker Name:",
                    handle_none(job.state.resources.slave_name if job.state and job.state.resources else None),
                )
            if job.state.resources.port_mappings:
                table.add_row(
                    "Port Mappings:",
                    handle_none(
                        _port_mapping_format(job.state.resources.port_mappings)
                        if job.state and job.state.resources
                        else None
                    ),
                )
                if job.state.resources.slave_host:
                    table.add_row(
                        "Worker Host:",
                        handle_none(job.state.resources.slave_host if job.state and job.state.resources else None),
                    )
            if job.state.resources.ssh_username and job.state.resources.ssh_host and job.state.resources.ssh_port:
                ssh_login_str = (
                    f"ssh{_ssh_port_command(job.state.resources.ssh_port)} "
                    f"{job.state.resources.ssh_username}@{job.state.resources.ssh_host}"
                )
                if job.state.resources.ssh_proxy_host:
                    ssh_proxy_login_str = (
                        f' -oProxyCommand="ssh{_ssh_port_command(job.state.resources.ssh_proxy_port)} '
                        f'{job.state.resources.ssh_proxy_username}@{job.state.resources.ssh_proxy_host} -W %h:%p"'
                    )
                    ssh_login_str += ssh_proxy_login_str
                table.add_row("SSH Login:", ssh_login_str)

        timing_table = Table.grid(Column(header="Key", justify="right", style="bold"), "Value", padding=(0, 1, 0, 1))
        if job.state:

            def _show_rel(shown_dt: datetime | None, earlier_dt: datetime | None, after_subject) -> str:
                if not shown_dt or not earlier_dt:
                    return ""
                rel_str = td_format(shown_dt - earlier_dt)
                return f" ({rel_str} after {after_subject})"

            def _show_ago(shown_dt: datetime | None) -> str:
                if not shown_dt:
                    return ""
                rel_str = td_format(datetime.now(UTC) - shown_dt)
                return f" ({rel_str} ago)"

            def _date_to_str(dt: datetime | None) -> str | None:
                if not dt:
                    return None
                used_dt = dt - timedelta(microseconds=dt.microsecond)
                # change timezone to local
                used_dt = used_dt.astimezone(tz=None)
                return used_dt.isoformat()

            if job.state.event_times:
                timing_table.add_row("", "[bold]Timing:")
                timing_table.add_row(
                    "Created:",
                    handle_none(_date_to_str(job.state.event_times.created)) + _show_ago(job.state.event_times.created),
                )
                timing_table.add_row(
                    "Queued:",
                    handle_none(_date_to_str(job.state.event_times.QUEUED))
                    + _show_rel(job.state.event_times.QUEUED, job.state.event_times.created, "job creation"),
                )
                timing_table.add_row(
                    "Assigned:",
                    handle_none(_date_to_str(job.state.event_times.ASSIGNED))
                    + _show_rel(job.state.event_times.ASSIGNED, job.state.event_times.QUEUED, "QUEUED"),
                )
                timing_table.add_row(
                    "Starting:",
                    handle_none(_date_to_str(job.state.event_times.STARTING))
                    + _show_rel(job.state.event_times.STARTING, job.state.event_times.ASSIGNED, "ASSIGNED"),
                )
                timing_table.add_row(
                    "Running:",
                    handle_none(_date_to_str(job.state.event_times.RUNNING))
                    + _show_rel(job.state.event_times.RUNNING, job.state.event_times.STARTING, "STARTING"),
                )

                if job.state.event_times.MUSTHALT and job.state.event_times.HALTING:
                    timing_table.add_row(
                        "HALT request:",
                        handle_none(_date_to_str(job.state.event_times.MUSTHALT))
                        + _show_rel(job.state.event_times.MUSTHALT, job.state.event_times.RUNNING, "RUNNING"),
                    )
                    timing_table.add_row(
                        "HALT procedure start:",
                        handle_none(_date_to_str(job.state.event_times.HALTING))
                        + _show_rel(job.state.event_times.HALTING, job.state.event_times.MUSTHALT, "HALT request"),
                    )
                    if job.state.event_times.HALTED:
                        timing_table.add_row(
                            "HALT success:",
                            handle_none(_date_to_str(job.state.event_times.HALTED))
                            + _show_rel(
                                job.state.event_times.HALTED, job.state.event_times.HALTING, "HALT procedure start"
                            ),
                        )
                    if job.state.event_times.FAILED:
                        timing_table.add_row(
                            "HALT failure:",
                            handle_none(_date_to_str(job.state.event_times.FAILED))
                            + _show_rel(
                                job.state.event_times.FAILED, job.state.event_times.HALTING, "HALT procedure start"
                            ),
                        )
                timing_table.add_row(
                    "Ended:",
                    handle_none(_date_to_str(job.state.event_times.end_date))
                    + _show_ago(job.state.event_times.end_date),
                )
                duration = job.state.event_times.duration_or_none
                timing_table.add_row("Duration:", td_format(duration) if duration else "-")
                timing_table.add_row(
                    "State Updated:",
                    handle_none(_date_to_str(job.state.event_times.status_updated))
                    + _show_ago(job.state.event_times.status_updated),
                )

            console.print(table)
            if timing_table.columns:
                console.print("")
                # console.print("\n[bold]Timing:")
                console.print(timing_table)
    elif format == ShowResourceFormat.json:
        console.print_json(json_encode(job).decode())
    elif format == ShowResourceFormat.yaml:
        console.print(yaml_encode(job).decode(), highlight=False)


@ai_app.command("output", rich_help_panel="Job interaction commands")
@handle_some_exceptions
def get_job_output(
    ctx: AIContext,
    job_id: Annotated[str, typer.Argument(help="Job ID.")],
):
    """Show Job output."""

    job = get_ai_job(ctx, job_id)
    logs = ctx.obj.api.for_project(ctx.obj.require_project_id()).ai(ctx.obj.site_id).jobs.output.get(job.id)
    # output without markup
    typer.echo(logs)


class DebugFormat(Enum):
    """Job Debug Info Formats."""

    table = "table"
    json = "json"
    yaml = "yaml"


@ai_app.command("debug", rich_help_panel="Auxiliary commands")
@handle_some_exceptions
def get_job_debug(
    ctx: AIContext,
    job_id: Annotated[str, typer.Argument(help="Job ID.")],
):
    """Show internal event log for a job."""

    def _4charsto5(v: str) -> str:
        return " " + v if len(v) == 4 else v  # noqa: PLR2004

    job = get_ai_job(ctx, job_id)
    debug_output = ctx.obj.api.for_project(ctx.obj.require_project_id()).ai(ctx.obj.site_id).jobs.logs.get(job.id)
    for out in debug_output:
        msg = (
            f"[blue]{_4charsto5(out.type.name)}[/blue]: {out.msg}"
            if out.type != JobDebugLogType.STATUS_CHANGE
            else f"[green bold]Status to {out.new_state.name if out.new_state else '?'}"
        )
        console.print(f"[bright_black]{out.time}[/bright_black]: {msg}", highlight=False)


# TODO: add
# get_slave_info
# get_info
# wait
