from __future__ import annotations

import sys
import json
from typing import Any, Dict, Literal, TypeVar, Callable, cast
from functools import wraps

import click

from together import Together, omit
from together.types import DedicatedEndpoint
from together._exceptions import APIError
from together.lib.utils.serializer import datetime_serializer
from together.types.endpoint_list_response import Data as DedicatedEndpointListItem


def print_endpoint(endpoint: DedicatedEndpoint | DedicatedEndpointListItem) -> 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}")

    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}, 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: APIError) -> None:
    error_details = cast(Dict[str, Any], e.body)["error"]["message"]

    if error_details and ("credentials" in error_details.lower() or "authentication" in error_details.lower()):
        click.echo("Error: Invalid API key or authentication failed", 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 APIError as e:
            print_api_error(e)
            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


@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(["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 | None,
    no_speculative_decoding: bool | None,
    no_auto_start: bool,
    inactive_timeout: int | None,
    availability_zone: str | None,
    wait: bool,
) -> None:
    """Create a new dedicated inference endpoint."""
    # Map GPU types to their full hardware ID names
    gpu_map = {
        "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,
            autoscaling={
                "min_replicas": min_replicas,
                "max_replicas": max_replicas,
            },
            display_name=display_name or omit,
            disable_prompt_cache=no_prompt_cache or omit,
            disable_speculative_decoding=no_speculative_decoding or omit,
            state="STOPPED" if no_auto_start else "STARTED",
            inactive_timeout=inactive_timeout,
            extra_query={"availability_zone": availability_zone or omit},
        )
    except APIError as e:
        print_api_error(e)
        if "check the hardware api" in str(e.args[0]).lower() or "invalid hardware provided" in str(e.args[0]).lower():
            fetch_and_print_hardware_options(client=client, model=model, print_json=False, available=True)

        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.retrieve(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.retrieve(endpoint_id)
    if json:
        import json as json_lib

        click.echo(json_lib.dumps(endpoint.model_dump(), indent=2, default=datetime_serializer))
    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 fetch_and_print_hardware_options(client: Together, model: str | None, print_json: bool, available: bool) -> None:
    """Print hardware options for a model."""

    message = "Available hardware options:" if available else "All hardware options:"
    click.echo(message, err=True)
    hardware_options = client.hardware.list(model=model or omit)
    # hardware_options = client.endpoints.list_hardware(model)
    if available:
        hardware_options.data = [
            hardware
            for hardware in hardware_options.data
            if hardware.availability is not None and hardware.availability.status == "available"
        ]

    if print_json:
        json_output = [hardware.model_dump() for hardware in hardware_options.data]
        click.echo(json.dumps(json_output, default=datetime_serializer, indent=2))
    else:
        for hardware in hardware_options.data:
            click.echo(f"  {hardware.id}", err=True)


@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.retrieve(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.retrieve(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 = client.endpoints.list(
        type=type or omit,
        usage_type=usage_type or omit,
        mine=mine if mine is not None else omit,
    )

    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.data], default=datetime_serializer, indent=2
            )
        )
    else:
        for endpoint in endpoints.data:
            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)

    # 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 or max_replicas is not None:
        kwargs["autoscaling"] = {}
        if min_replicas is not None:
            kwargs["autoscaling"]["min_replicas"] = min_replicas
        if max_replicas is not None:
            kwargs["autoscaling"]["max_replicas"] = max_replicas

    if inactive_timeout is not None:
        kwargs["inactive_timeout"] = inactive_timeout

    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:
        click.echo(f"  Min replicas: {min_replicas}", err=True)
    if max_replicas:
        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.model_dump(), indent=2))
    else:
        click.echo("Available zones:", err=True)
        for availability_zone in sorted(avzones.avzones):
            click.echo(f"  {availability_zone}")
