"""Custom streaming utilities for Mithril CLI.

This module provides a polling-based approach to monitoring SkyPilot requests,
giving full control over the UI rather than parsing SkyPilot's log output.
"""

from __future__ import annotations

import time
from typing import TYPE_CHECKING, TypeVar

from rich.live import Live
from rich.spinner import Spinner
from rich.text import Text

from mithril.cli.utils.skypilot_passthrough import SKY_ALIAS_COMMANDS

T = TypeVar("T")

if TYPE_CHECKING:
    from collections.abc import Callable

    from rich.console import Console
    from sky.server.common import RequestId

    from mithril.sky import SkyClient

TERMINAL_STATUSES = {"SUCCEEDED", "FAILED", "CANCELLED"}


def stream_job_logs(
    cluster_name: str,
    job_id: int,
    console: Console,
    *,
    sky: SkyClient,
    follow: bool = True,
    filter_fn: Callable[[str], str | None] | None = None,
) -> int:
    """Stream job logs with optional filtering.

    Args:
        cluster_name: The cluster name
        job_id: The job ID to stream logs for
        console: Rich console for output
        sky: SkyClient instance
        follow: Whether to follow the logs (like tail -f)
        filter_fn: Optional function to filter/transform each line.
                   Return the line to print, or None to skip it.

    Returns:
        Exit code (0 for success, non-zero for failure)
    """
    # Use preload_content=False to get an iterator
    log_iter = sky.tail_logs(
        cluster_name=cluster_name,
        job_id=job_id,
        follow=follow,
        preload_content=False,
    )

    for line in log_iter:
        if line is None:
            # End of stream
            break

        if filter_fn:
            filtered = filter_fn(line)
            if filtered is not None:
                console.print(filtered, end="", markup=False, highlight=False)
        else:
            console.print(line, end="", markup=False, highlight=False)

    # Get the exit code by checking job status
    job_statuses = sky.get(sky.job_status(cluster_name, job_ids=[job_id]))
    status = job_statuses.get(job_id)

    if status is None:
        return 1  # Job not found

    # JobStatus enum - check if succeeded
    return 0 if str(status) == "SUCCEEDED" else 1


def poll_launch(  # noqa: UP047 - PEP 695 syntax requires 3.12+
    request_id: RequestId[T],
    console: Console,
    *,
    sky: SkyClient,
    poll_interval: float = 1.0,
    quiet: bool = False,
    initial_cluster_name: str | None = None,
) -> tuple[T, str | None]:
    """Poll a launch request until completion, showing custom progress.

    Args:
        request_id: The SkyPilot request ID from sky.launch()
        console: Rich console for output
        sky: SkyClient instance
        poll_interval: How often to poll for status updates (seconds)
        quiet: If True, suppress progress output
        initial_cluster_name: Optional cluster name to seed status messages

    Returns:
        The result from the launch request (job_id, handle)
    """
    cluster_name = initial_cluster_name
    initial_renderable, last_display_text = _render_status_message(
        None,
        cluster_name,
    )

    spinner = Spinner("dots", text="Launching...")
    live = Live(spinner, console=console, refresh_per_second=10, transient=True)

    try:
        if not quiet:
            spinner.update(text=initial_renderable)
            live.start()
        last_display_text, cluster_name = _run_poll_loop(
            request_id=request_id,
            poll_interval=poll_interval,
            spinner=spinner,
            live=live,
            last_display_text=last_display_text,
            cluster_name=cluster_name,
            quiet=quiet,
            sky=sky,
        )
    except KeyboardInterrupt:
        if not quiet and cluster_name:
            _print_provision_hint(console, cluster_name)
        raise
    finally:
        if not quiet:
            live.stop()

    # Get the final result (this will raise if there was an error)
    result = sky.get(request_id)

    return result, cluster_name


def format_status_message(status_msg: str, cluster_name: str | None) -> Text:
    """Format a status message for display.

    Appends the cluster name to status messages so users know which cluster
    is being worked on. For example: "Launching" becomes "Launching (my-cluster)".

    The cluster name is optional because it isn't available at the start of a
    launch—SkyPilot assigns it during provisioning. Early status messages won't
    have a cluster name yet.
    """
    formatted = status_msg
    if cluster_name:
        formatted = f"{formatted} ({cluster_name})"
        formatted = (
            f"{formatted} [dim]View logs: {_cli_command('logs')} --provision "
            f"{cluster_name}[/dim]"
        )
    return Text.from_markup(formatted)


def _default_launch_message(cluster_name: str | None) -> str:
    if cluster_name:
        return (
            "Launching [dim]View logs: "
            f"{_cli_command('logs')} --provision {cluster_name}[/dim]"
        )
    return "Launching..."


def _render_status_message(
    status_msg: str | None,
    cluster_name: str | None,
) -> tuple[Text, str]:
    if status_msg:
        renderable = format_status_message(status_msg, cluster_name)
        return renderable, renderable.plain
    message = _default_launch_message(cluster_name)
    renderable = Text.from_markup(message)
    return renderable, renderable.plain


def _run_poll_loop(
    *,
    request_id: str,
    poll_interval: float,
    spinner: Spinner,
    live: Live,
    last_display_text: str,
    cluster_name: str | None,
    quiet: bool,
    sky: SkyClient,
) -> tuple[str, str | None]:
    while True:
        statuses = sky.api_status(request_ids=[request_id])

        if not statuses:
            time.sleep(poll_interval)
            continue

        status = statuses[0]
        current_status = status.status
        status_msg = status.status_msg
        cluster_name = status.cluster_name or cluster_name

        if not quiet:
            renderable, display_text = _render_status_message(
                status_msg,
                cluster_name,
            )
            if display_text != last_display_text:
                last_display_text = display_text
                spinner.update(text=renderable)
                live.update(spinner)

        if current_status in TERMINAL_STATUSES:
            break

        time.sleep(poll_interval)

    return last_display_text, cluster_name


def _print_provision_hint(console: Console, cluster_name: str) -> None:
    console.print(
        f"[dim]View logs: {_cli_command('logs')} --provision {cluster_name}[/dim]"
    )


def _cli_command(command: str) -> str:
    if command in SKY_ALIAS_COMMANDS:
        return f"ml {command}"
    return f"ml sky {command}"
