endpoints.py

from __future__ import annotations

import json
import re
import sys
from functools import wraps
from typing import Any, Callable, Dict, List, Literal, Sequence, TypeVar, Union

import click
from tabulate import tabulate

from together import Together
from together.error import InvalidRequestError
from together.types import DedicatedEndpoint, ListEndpoint
from together.types.endpoints import HardwareWithStatus


def print_endpoint(
    endpoint: Union[DedicatedEndpoint, ListEndpoint],
) -> None:
    """Print endpoint details in a Docker-like format or JSON."""

    # Print header info
    click.echo(f"ID:\t\t{endpoint.id}")
    click.echo(f"Name:\t\t{endpoint.name}")

    # Print type-specific fields
    if isinstance(endpoint, DedicatedEndpoint):
        click.echo(f"Display Name:\t{endpoint.display_name}")
        click.echo(f"Hardware:\t{endpoint.hardware}")
        click.echo(
            f"Autoscaling:\tMin={endpoint.autoscaling.min_replicas}, "
            f"Max={endpoint.autoscaling.max_replicas}"
        )

    click.echo(f"Model:\t\t{endpoint.model}")
    click.echo(f"Type:\t\t{endpoint.type}")
    click.echo(f"Owner:\t\t{endpoint.owner}")
    click.echo(f"State:\t\t{endpoint.state}")
    click.echo(f"Created:\t{endpoint.created_at}")


F = TypeVar("F", bound=Callable[..., Any])


def print_api_error(
    e: InvalidRequestError,
    endpoint_id: str | None = None,
) -> None:
    error_details = e.api_response.message
    error_lower = error_details.lower() if error_details else ""

    if "credentials" in error_lower or "authentication" in error_lower:
        click.echo("Error: Invalid API key or authentication failed", err=True)
    elif "not found" in error_lower and "endpoint" in error_lower:
        endpoint_display = f"'{endpoint_id}'" if endpoint_id else ""
        click.echo(f"Error: Endpoint {endpoint_display} not found.", err=True)
        click.echo(
            "The endpoint may have been deleted or the ID may be incorrect.",
            err=True,
        )
        click.echo(
            "Use 'together endpoints list --mine true' to see your endpoints.",
            err=True,
        )
    elif "permission" in error_lower or "forbidden" in error_lower or "unauthorized" in error_lower:
        click.echo("Error: You don't have permission to access this endpoint.", err=True)
        click.echo(
            "This endpoint may belong to another user or organization.",
            err=True,
        )
    else:
        click.echo(f"Error: {error_details}", err=True)


def handle_api_errors(f: F) -> F:
    """Decorator to handle common API errors in CLI commands."""

    @wraps(f)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        try:
            return f(*args, **kwargs)
        except InvalidRequestError as e:
            # Try to extract endpoint_id from kwargs for better error messages
            endpoint_id = kwargs.get("endpoint_id")
            print_api_error(e, endpoint_id=endpoint_id)
            sys.exit(1)
        except Exception as e:
            click.echo(f"Error: An unexpected error occurred - {str(e)}", err=True)
            sys.exit(1)

    return wrapper  # type: ignore


@click.group()
@click.pass_context
def endpoints(ctx: click.Context) -> None:
    """Endpoints API commands"""
    pass


def _print_hardware_error(
    client: Together,
    model: str,
    hardware_id: str,
    gpu: str,
    gpu_count: int,
    *,
    speculative_decoding_enabled: bool = False,
) -> None:
    """Print a detailed error message when hardware selection fails."""
    click.echo(
        f"Error: Cannot create endpoint with {gpu_count}x {gpu.upper()} for model '{model}'",
        err=True,
    )

    # Fetch hardware options for this model to provide specific guidance
    try:
        hardware_options = client.endpoints.list_hardware(model)
    except Exception:
        # If we can't fetch hardware options, just show a generic message
        click.echo(
            "\nUse 'together endpoints hardware --model <model>' to see available options.",
            err=True,
        )
        return

    # Check if the requested hardware exists for this model
    requested_hw = next((hw for hw in hardware_options if hw.id == hardware_id), None)

    if requested_hw is None:
        # Hardware configuration is not compatible with this model
        click.echo(
            f"\nThe hardware configuration '{hardware_id}' is not compatible with this model.",
            err=True,
        )
    elif requested_hw.availability:
        status = requested_hw.availability.status
        if status == "unavailable":
            click.echo(
                f"\nThe {gpu_count}x {gpu.upper()} configuration is currently unavailable. "
                "This hardware type has no available capacity at this time.",
                err=True,
            )
        elif status == "insufficient":
            click.echo(
                f"\nThe {gpu_count}x {gpu.upper()} configuration has insufficient capacity. "
                "Not enough GPUs available for the requested number of replicas.",
                err=True,
            )
        elif status == "available":
            # Hardware is available but request failed - suggest toggling speculative decoding
            if speculative_decoding_enabled:
                click.echo(
                    "\nHardware is available but this configuration is not supported. "
                    "Try adding --no-speculative-decoding.",
                    err=True,
                )
            else:
                click.echo(
                    "\nHardware is available but this configuration is not supported. "
                    "Try removing --no-speculative-decoding to enable speculative decoding.",
                    err=True,
                )
            return

    # Show available alternatives
    available_options = [
        hw
        for hw in hardware_options
        if hw.availability is not None and hw.availability.status == "available"
    ]

    if available_options:
        click.echo("\nAvailable hardware options for this model:", err=True)
        click.echo("", err=True)
        _format_hardware_options(available_options, show_availability=True)
    else:
        click.echo(
            "\nNo hardware is currently available for this model. Please try again later.",
            err=True,
        )
        click.echo("\nAll hardware options for this model:", err=True)
        click.echo("", err=True)
        _format_hardware_options(hardware_options, show_availability=True)


@endpoints.command()
@click.option(
    "--model",
    required=True,
    help="The model to deploy (e.g. meta-llama/Llama-4-Scout-17B-16E-Instruct)",
)
@click.option(
    "--min-replicas",
    type=int,
    default=1,
    help="Minimum number of replicas to deploy",
)
@click.option(
    "--max-replicas",
    type=int,
    default=1,
    help="Maximum number of replicas to deploy",
)
@click.option(
    "--gpu",
    type=click.Choice(["b200", "h200", "h100", "a100", "l40", "l40s", "rtx-6000"]),
    required=True,
    help="GPU type to use for inference",
)
@click.option(
    "--gpu-count",
    type=int,
    default=1,
    help="Number of GPUs to use per replica",
)
@click.option(
    "--display-name",
    help="A human-readable name for the endpoint",
)
@click.option(
    "--no-prompt-cache",
    is_flag=True,
    help="Disable the prompt cache for this endpoint",
)
@click.option(
    "--no-speculative-decoding",
    is_flag=True,
    help="Disable speculative decoding for this endpoint",
)
@click.option(
    "--no-auto-start",
    is_flag=True,
    help="Create the endpoint in STOPPED state instead of auto-starting it",
)
@click.option(
    "--inactive-timeout",
    type=int,
    help="Number of minutes of inactivity after which the endpoint will be automatically stopped. Set to 0 to disable.",
)
@click.option(
    "--availability-zone",
    help="Start endpoint in specified availability zone (e.g., us-central-4b)",
)
@click.option(
    "--wait/--no-wait",
    default=True,
    help="Wait for the endpoint to be ready after creation",
)
@click.pass_obj
@handle_api_errors
def create(
    client: Together,
    model: str,
    min_replicas: int,
    max_replicas: int,
    gpu: str,
    gpu_count: int,
    display_name: str | None,
    no_prompt_cache: bool,
    no_speculative_decoding: bool,
    no_auto_start: bool,
    inactive_timeout: int | None,
    availability_zone: str | None,
    wait: bool,
) -> None:
    """Create a new dedicated inference endpoint."""
    # Client-side validation for replicas
    if min_replicas < 0:
        click.echo(
            f"Error: --min-replicas must be non-negative, got {min_replicas}", err=True
        )
        sys.exit(1)
    if max_replicas < 0:
        click.echo(
            f"Error: --max-replicas must be non-negative, got {max_replicas}", err=True
        )
        sys.exit(1)
    if min_replicas > max_replicas:
        click.echo(
            f"Error: --min-replicas ({min_replicas}) cannot be greater than "
            f"--max-replicas ({max_replicas})",
            err=True,
        )
        sys.exit(1)

    # Validate GPU count
    valid_gpu_counts = [1, 2, 4, 8]
    if gpu_count not in valid_gpu_counts:
        click.echo(
            f"Error: --gpu-count must be one of {valid_gpu_counts}, got {gpu_count}",
            err=True,
        )
        sys.exit(1)

    # Validate availability zone if specified
    if availability_zone:
        try:
            valid_zones = client.endpoints.list_avzones()
            if availability_zone not in valid_zones:
                click.echo(
                    f"Error: Invalid availability zone '{availability_zone}'", err=True
                )
                if valid_zones:
                    click.echo("Available zones:", err=True)
                    for zone in sorted(valid_zones):
                        click.echo(f"  {zone}", err=True)
                sys.exit(1)
        except Exception:
            # If we can't fetch zones, let the API validate it
            pass

    # Map GPU types to their full hardware ID names
    gpu_map = {
        "b200": "nvidia_b200_180gb_sxm",
        "h200": "nvidia_h200_140gb_sxm",
        "h100": "nvidia_h100_80gb_sxm",
        "a100": "nvidia_a100_80gb_pcie" if gpu_count == 1 else "nvidia_a100_80gb_sxm",
        "l40": "nvidia_l40",
        "l40s": "nvidia_l40s",
        "rtx-6000": "nvidia_rtx_6000_ada",
    }

    hardware_id = f"{gpu_count}x_{gpu_map[gpu]}"

    try:
        response = client.endpoints.create(
            model=model,
            hardware=hardware_id,
            min_replicas=min_replicas,
            max_replicas=max_replicas,
            display_name=display_name,
            disable_prompt_cache=no_prompt_cache,
            disable_speculative_decoding=no_speculative_decoding,
            state="STOPPED" if no_auto_start else "STARTED",
            inactive_timeout=inactive_timeout,
            availability_zone=availability_zone,
        )
    except InvalidRequestError as e:
        error_msg = str(e.args[0]).lower() if e.args else ""
        if (
            "check the hardware api" in error_msg
            or "invalid hardware provided" in error_msg
            or "the selected configuration" in error_msg
        ):
            # speculative decoding is enabled when --no-speculative-decoding is NOT passed
            speculative_decoding_enabled = not no_speculative_decoding
            _print_hardware_error(
                client,
                model,
                hardware_id,
                gpu,
                gpu_count,
                speculative_decoding_enabled=speculative_decoding_enabled,
            )
        elif "model" in error_msg and (
            "not found" in error_msg
            or "invalid" in error_msg
            or "does not exist" in error_msg
            or "not supported" in error_msg
        ):
            click.echo(
                f"Error: Model '{model}' was not found or is not available for "
                "dedicated endpoints.",
                err=True,
            )
            click.echo(
                "Please check that the model name is correct and that it supports "
                "dedicated endpoint deployment.",
                err=True,
            )
            click.echo(
                "You can browse available models at: https://api.together.ai/models",
                err=True,
            )
        elif "availability" in error_msg and "zone" in error_msg:
            click.echo(
                f"Error: Availability zone '{availability_zone}' is not valid.",
                err=True,
            )
            try:
                valid_zones = client.endpoints.list_avzones()
                if valid_zones:
                    click.echo("\nAvailable zones:", err=True)
                    for zone in sorted(valid_zones):
                        click.echo(f"  {zone}", err=True)
            except Exception:
                pass
        else:
            print_api_error(e)
        sys.exit(1)

    # Print detailed information to stderr
    click.echo("Created dedicated endpoint with:", err=True)
    click.echo(f"  Model: {model}", err=True)
    click.echo(f"  Min replicas: {min_replicas}", err=True)
    click.echo(f"  Max replicas: {max_replicas}", err=True)
    click.echo(f"  Hardware: {hardware_id}", err=True)
    if display_name:
        click.echo(f"  Display name: {display_name}", err=True)
    if no_prompt_cache:
        click.echo("  Prompt cache: disabled", err=True)
    if no_speculative_decoding:
        click.echo("  Speculative decoding: disabled", err=True)
    if no_auto_start:
        click.echo("  Auto-start: disabled", err=True)
    if inactive_timeout is not None:
        click.echo(f"  Inactive timeout: {inactive_timeout} minutes", err=True)
    if availability_zone:
        click.echo(f"  Availability zone: {availability_zone}", err=True)

    click.echo(f"Endpoint created successfully, id: {response.id}", err=True)

    if wait:
        import time

        click.echo("Waiting for endpoint to be ready...", err=True)
        while client.endpoints.get(response.id).state != "STARTED":
            time.sleep(1)
        click.echo("Endpoint ready", err=True)

    # Print only the endpoint ID to stdout
    click.echo(response.id)


@endpoints.command()
@click.argument("endpoint-id", required=True)
@click.option("--json", is_flag=True, help="Print output in JSON format")
@click.pass_obj
@handle_api_errors
def get(client: Together, endpoint_id: str, json: bool) -> None:
    """Get a dedicated inference endpoint."""
    endpoint = client.endpoints.get(endpoint_id)
    if json:
        import json as json_lib

        click.echo(json_lib.dumps(endpoint.model_dump(), indent=2))
    else:
        print_endpoint(endpoint)


@endpoints.command()
@click.option("--model", help="Filter hardware options by model")
@click.option("--json", is_flag=True, help="Print output in JSON format")
@click.option(
    "--available",
    is_flag=True,
    help="Print only available hardware options (can only be used if model is passed in)",
)
@click.pass_obj
@handle_api_errors
def hardware(client: Together, model: str | None, json: bool, available: bool) -> None:
    """List all available hardware options, optionally filtered by model."""
    fetch_and_print_hardware_options(client, model, json, available)


def _format_hardware_options(
    hardware_options: Sequence[HardwareWithStatus],
    show_availability: bool = True,
) -> None:
    """Print hardware options in a formatted table using tabulate."""
    if not hardware_options:
        click.echo("  No hardware options found.", err=True)
        return

    display_list: List[Dict[str, Any]] = []

    for hw in hardware_options:
        data: Dict[str, Any] = {
            "Hardware ID": hw.id,
            "GPU": (
                re.sub(r"\-\d+[a-zA-Z][a-zA-Z]$", "", hw.specs.gpu_type)
                if hw.specs and hw.specs.gpu_type
                else "N/A"
            ),
            "Memory": f"{int(hw.specs.gpu_memory)}GB" if hw.specs else "N/A",
            "Count": hw.specs.gpu_count if hw.specs else "N/A",
            "Price (per minute)": (
                f"${hw.pricing.cents_per_minute / 100:.2f}" if hw.pricing else "N/A"
            ),
        }

        if show_availability:
            status_display = "—"
            if hw.availability:
                status = hw.availability.status
                # Add visual indicators for status
                if status == "available":
                    status_display = click.style("✓ available", fg="green")
                elif status == "unavailable":
                    status_display = click.style("✗ unavailable", fg="red")
                else:  # insufficient
                    status_display = click.style("⚠ insufficient", fg="yellow")
            data["Availability"] = status_display

        display_list.append(data)

    click.echo(tabulate(display_list, headers="keys", numalign="left"))


def fetch_and_print_hardware_options(
    client: Together, model: str | None, print_json: bool, available: bool
) -> None:
    """Print hardware options for a model."""
    hardware_options = client.endpoints.list_hardware(model)

    if available:
        hardware_options = [
            hardware
            for hardware in hardware_options
            if hardware.availability is not None
            and hardware.availability.status == "available"
        ]
        message = (
            f"Available hardware options for model '{model}':"
            if model
            else "Available hardware options:"
        )
    else:
        message = (
            f"Hardware options for model '{model}':"
            if model
            else "All hardware options:"
        )

    click.echo(message, err=True)
    click.echo("", err=True)

    if print_json:
        json_output = [hardware.model_dump() for hardware in hardware_options]
        click.echo(json.dumps(json_output, indent=2))
    else:
        # Show availability column only when model is specified (availability info is only returned with model filter)
        show_availability = model is not None
        _format_hardware_options(hardware_options, show_availability=show_availability)


@endpoints.command()
@click.argument("endpoint-id", required=True)
@click.option(
    "--wait/--no-wait",
    default=True,
    help="Wait for the endpoint to stop",
)
@click.pass_obj
@handle_api_errors
def stop(client: Together, endpoint_id: str, wait: bool) -> None:
    """Stop a dedicated inference endpoint."""
    client.endpoints.update(endpoint_id, state="STOPPED")
    click.echo("Successfully marked endpoint as stopping", err=True)

    if wait:
        import time

        click.echo("Waiting for endpoint to stop...", err=True)
        while client.endpoints.get(endpoint_id).state != "STOPPED":
            time.sleep(1)
        click.echo("Endpoint stopped", err=True)

    click.echo(endpoint_id)


@endpoints.command()
@click.argument("endpoint-id", required=True)
@click.option(
    "--wait/--no-wait",
    default=True,
    help="Wait for the endpoint to start",
)
@click.pass_obj
@handle_api_errors
def start(client: Together, endpoint_id: str, wait: bool) -> None:
    """Start a dedicated inference endpoint."""
    client.endpoints.update(endpoint_id, state="STARTED")
    click.echo("Successfully marked endpoint as starting", err=True)

    if wait:
        import time

        click.echo("Waiting for endpoint to start...", err=True)
        while client.endpoints.get(endpoint_id).state != "STARTED":
            time.sleep(1)
        click.echo("Endpoint started", err=True)

    click.echo(endpoint_id)


@endpoints.command()
@click.argument("endpoint-id", required=True)
@click.pass_obj
@handle_api_errors
def delete(client: Together, endpoint_id: str) -> None:
    """Delete a dedicated inference endpoint."""
    client.endpoints.delete(endpoint_id)
    click.echo("Successfully deleted endpoint", err=True)
    click.echo(endpoint_id)


@endpoints.command()
@click.option("--json", is_flag=True, help="Print output in JSON format")
@click.option(
    "--type",
    type=click.Choice(["dedicated", "serverless"]),
    help="Filter by endpoint type",
)
@click.option(
    "--mine",
    type=click.BOOL,
    default=None,
    help="true (only mine), default=all",
)
@click.option(
    "--usage-type",
    type=click.Choice(["on-demand", "reserved"]),
    help="Filter by endpoint usage type",
)
@click.pass_obj
@handle_api_errors
def list(
    client: Together,
    json: bool,
    type: Literal["dedicated", "serverless"] | None,
    usage_type: Literal["on-demand", "reserved"] | None,
    mine: bool | None,
) -> None:
    """List all inference endpoints (includes both dedicated and serverless endpoints)."""
    endpoints: List[ListEndpoint] = client.endpoints.list(
        type=type, usage_type=usage_type, mine=mine
    )

    if not endpoints:
        click.echo("No dedicated endpoints found", err=True)
        return

    click.echo("Endpoints:", err=True)
    if json:
        import json as json_lib

        click.echo(
            json_lib.dumps([endpoint.model_dump() for endpoint in endpoints], indent=2)
        )
    else:
        for endpoint in endpoints:
            print_endpoint(
                endpoint,
            )
            click.echo()


@endpoints.command()
@click.argument("endpoint-id", required=True)
@click.option(
    "--display-name",
    help="A new human-readable name for the endpoint",
)
@click.option(
    "--min-replicas",
    type=int,
    help="New minimum number of replicas to maintain",
)
@click.option(
    "--max-replicas",
    type=int,
    help="New maximum number of replicas to scale up to",
)
@click.option(
    "--inactive-timeout",
    type=int,
    help="Number of minutes of inactivity after which the endpoint will be automatically stopped. Set to 0 to disable.",
)
@click.pass_obj
@handle_api_errors
def update(
    client: Together,
    endpoint_id: str,
    display_name: str | None,
    min_replicas: int | None,
    max_replicas: int | None,
    inactive_timeout: int | None,
) -> None:
    """Update a dedicated inference endpoint's configuration."""
    if not any([display_name, min_replicas, max_replicas, inactive_timeout]):
        click.echo("Error: At least one update option must be specified", err=True)
        sys.exit(1)

    # If only one of min/max replicas is specified, we need both for the update
    if (min_replicas is None) != (max_replicas is None):
        click.echo(
            "Error: Both --min-replicas and --max-replicas must be specified together",
            err=True,
        )
        sys.exit(1)

    # Build kwargs for the update
    kwargs: Dict[str, Any] = {}
    if display_name is not None:
        kwargs["display_name"] = display_name
    if min_replicas is not None and max_replicas is not None:
        kwargs["min_replicas"] = min_replicas
        kwargs["max_replicas"] = max_replicas
    if inactive_timeout is not None:
        kwargs["inactive_timeout"] = inactive_timeout

    _response = client.endpoints.update(endpoint_id, **kwargs)

    # Print what was updated
    click.echo("Updated endpoint configuration:", err=True)
    if display_name:
        click.echo(f"  Display name: {display_name}", err=True)
    if min_replicas is not None and max_replicas is not None:
        click.echo(f"  Min replicas: {min_replicas}", err=True)
        click.echo(f"  Max replicas: {max_replicas}", err=True)
    if inactive_timeout is not None:
        click.echo(f"  Inactive timeout: {inactive_timeout} minutes", err=True)

    click.echo("Successfully updated endpoint", err=True)
    click.echo(endpoint_id)


@endpoints.command()
@click.option("--json", is_flag=True, help="Print output in JSON format")
@click.pass_obj
@handle_api_errors
def availability_zones(client: Together, json: bool) -> None:
    """List all availability zones."""
    avzones = client.endpoints.list_avzones()

    if not avzones:
        click.echo("No availability zones found", err=True)
        return

    if json:
        import json as json_lib

        click.echo(json_lib.dumps({"avzones": avzones}, indent=2))
    else:
        click.echo("Available zones:", err=True)
        for availability_zone in sorted(avzones):
            click.echo(f"  {availability_zone}")