"""Job logs command."""

from contextlib import suppress
from typing import Optional

import rich_click as click

from lightning_sdk.api.logs_api import SEVERITIES
from lightning_sdk.cli.utils.logging import LightningCommand
from lightning_sdk.cli.utils.logs import (
    LogSelection,
    make_logs_group,
    read_logs,
    resolve_time,
    run_download,
)
from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_job_machine, resolve_teamspace


@click.command("logs", cls=LightningCommand)
@click.argument("name", required=False, help="The job name. Required.")
@click.option(
    "--teamspace",
    default=None,
    help="Teamspace owner/name. Uses the configured default teamspace when omitted.",
)
@click.option("--follow", "-f", is_flag=True, default=False, help="Stream new log lines as they are produced.")
@click.option("--tail", type=int, default=None, help="Only show the last N lines.")
@click.option("--rank", type=int, default=None, help="Machine rank to read from in a multi-machine job.")
@click.option("--timestamps", is_flag=True, default=False, help="Prepend each line with its ISO-8601 timestamp.")
@click.option("--since", default=None, help='Only include lines at or after this time (e.g. "2h", RFC3339).')
@click.option("--until", default=None, help='Only include lines at or before this time (e.g. "30m", RFC3339).')
@click.option("--query", default=None, help="Only include lines containing every whitespace-separated term.")
@click.option(
    "--severity",
    type=click.Choice(SEVERITIES),
    default=None,
    help="Only include lines at or above this severity.",
)
@click.option("--json", "as_json", is_flag=True, default=False, help="Output entries as a JSON array.")
@click.option("--interactive", "-i", "tui", is_flag=True, default=False, help="Launch the interactive TUI log viewer.")
def _logs_job_cmd(
    name: Optional[str] = None,
    teamspace: Optional[str] = None,
    follow: bool = False,
    tail: Optional[int] = None,
    rank: Optional[int] = None,
    timestamps: bool = False,
    since: Optional[str] = None,
    until: Optional[str] = None,
    query: Optional[str] = None,
    severity: Optional[str] = None,
    as_json: bool = False,
    tui: bool = False,
) -> None:
    """Print the logs for a job.

    Prints a snapshot of the logs available so far. Pass --follow to stream new
    lines from a running job until it finishes or you press Ctrl-C. --query and
    --severity are applied by the server, to both the snapshot and the stream.
    Multi-machine logs are merged unless --rank selects one machine; the machine
    is resolved client-side and read through the logs API.
    """
    resolved_teamspace = resolve_teamspace(teamspace)
    resource = resolve_job(name, resolved_teamspace)
    selected_rank = resource.is_multi_machine is True and rank is not None

    if tui:
        from lightning_sdk.cli.logs_tui import run_tui

        job = resource
        if selected_rank:
            assert rank is not None
            job = resolve_job_machine(resource, rank)
        assert job.resource_id is not None, "job must have a resource_id"
        run_tui(
            LogSelection(teamspace_id=resolved_teamspace.id, job_ids=[job.resource_id]),
            follow=(follow or (since is None and until is None)),
            tail=tail,
            show_timestamps=True,
            since=since,
            until=until,
            query=query,
            title=f"{resolved_teamspace.owner.name}/{resolved_teamspace.name}/{job.name} logs",
        )
        return

    if as_json:
        if rank is not None and not selected_rank:
            raise click.ClickException("--rank is not supported with --json.")
        if selected_rank:
            assert rank is not None
            machine = resolve_job_machine(resource, rank)
            if machine.resource_id is None:
                raise click.ClickException("The selected machine does not have a resource ID.")
            selection = LogSelection(teamspace_id=resolved_teamspace.id, job_ids=[machine.resource_id])
        elif resource.is_multi_machine is True:
            labels: dict = {}
            with suppress(Exception):
                labels = {
                    machine.resource_id: machine.name
                    for machine in resource.machines
                    if machine.resource_id is not None
                }
            selection = LogSelection(
                teamspace_id=resolved_teamspace.id,
                mmt_id=resource.resource_id,
                labels=labels,
            )
        else:
            if resource.resource_id is None:
                raise click.ClickException("The selected job does not have a resource ID.")
            selection = LogSelection(teamspace_id=resolved_teamspace.id, job_ids=[resource.resource_id])
        read_logs(
            selection,
            query=query,
            severity=severity,
            since=resolve_time(since, "--since"),
            until=resolve_time(until, "--until"),
            tail=tail,
            follow=follow,
            as_json=True,
        )
        return

    try:
        # `rank` (when set) selects one machine of a multi-machine job; the SDK resolves it
        # client-side to that machine's own job and reads it through the logs API.
        logs = resource.logs(
            follow=follow,
            tail=tail,
            rank=rank,
            timestamps=timestamps,
            since=resolve_time(since, "--since"),
            until=resolve_time(until, "--until"),
            query=query,
            severity=severity,
        )
        if follow:
            for line in logs:
                click.echo(line)
        elif logs:
            click.echo(logs)
    except KeyboardInterrupt:
        pass
    except (RuntimeError, ValueError) as ex:
        raise click.ClickException(str(ex)) from ex


@click.command("download", cls=LightningCommand)
@click.argument("name", required=False, help="The job name. Required.")
@click.option(
    "--teamspace",
    default=None,
    help="Teamspace owner/name. Uses the configured default teamspace when omitted.",
)
@click.option("--timestamps", is_flag=True, default=False, help="Prepend each line with its ISO-8601 timestamp.")
def download_job(
    name: Optional[str] = None,
    teamspace: Optional[str] = None,
    timestamps: bool = False,
) -> None:
    """Download the complete logs for a job as a file."""
    resolved_teamspace = resolve_teamspace(teamspace)
    resource = resolve_job(name, resolved_teamspace)
    run_download(resource, timestamps=timestamps)


logs_job = make_logs_group(
    default_cmd=_logs_job_cmd,
    download_cmd=download_job,
    help_text="View the logs for a job.",
)