"""Embedding commands - standalone embedding service management.

This is the canonical CLI entry-point for sageLLM embedding.  It mirrors the
LLM surface (serve / embed / health) so embeddings are first-class citizens
alongside inference, instead of being hidden behind ``sagellm serve --with-embedding``.

Architecture
------------
* ``sagellm embedding serve``  → foreground FastAPI server (sagellm_core.embedding_server)
* ``sagellm embedding embed``  → local one-shot inference via LLMEngine (no server needed)
                                  or HTTP call via --endpoint <url>
* ``sagellm embedding health`` → HTTP health check of a running server

Both serve and embed use ``LLMEngine(config=LLMEngineConfig(task_type="embedding"))``,
the same unified inference path as text generation -- no extra dependencies required.

SAGE compatibility
------------------
``sagellm.EmbeddingService.from_endpoint()`` defaults to
``http://127.0.0.1:{_SagePorts.EMBEDDING_DEFAULT}``, which matches the default port of
``sagellm embedding serve``.  No changes required on the SAGE side.
"""

from __future__ import annotations

import json
import os
import sys

import click

from sagellm_protocol.ports import SagePorts as _SagePorts

from ..utils.config import init_default_config
from ..utils.console import get_console

DEFAULT_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
DEFAULT_EMBEDDING_PORT: int = _SagePorts.EMBEDDING_DEFAULT  # centralised in SagePorts


@click.group()
def embedding() -> None:
    """Manage sageLLM embedding service.

    \b
    Quick start:
      sagellm embedding serve                    # start embedding server
      sagellm embedding embed "hello world"      # one-shot embedding (local)
      sagellm embedding health                   # check server health
    """


# ---------------------------------------------------------------------------
# serve
# ---------------------------------------------------------------------------


@embedding.command("serve")
@click.option(
    "--model",
    "-m",
    type=str,
    default=None,
    help="Embedding model id or local path",
)
@click.option(
    "--port",
    type=int,
    default=DEFAULT_EMBEDDING_PORT,
    show_default=True,
    help="Server port",
)
@click.option(
    "--host",
    type=str,
    default="0.0.0.0",
    show_default=True,
    help="Server bind host",
)
def serve(model: str | None, port: int, host: str) -> None:
    """Start a standalone embedding server (OpenAI-compatible /v1/embeddings).

    This replaces ``sagellm serve --with-embedding`` for pure-embedding
    deployments.

    \b
    Examples:
        sagellm embedding serve
        sagellm embedding serve --model BAAI/bge-m3 --port 8891
        sagellm embedding serve --host 127.0.0.1
    """
    init_default_config()
    console = get_console()

    model_name = model or DEFAULT_EMBEDDING_MODEL

    console.print("\n🔧 [bold blue]Starting sageLLM Embedding Server...[/bold blue]\n")
    console.print(f"  [green]Model:[/green]    {model_name}")
    console.print(f"  [green]Endpoint:[/green] http://{host}:{port}/v1/embeddings")
    console.print(f"  [green]Health:[/green]   http://{host}:{port}/health\n")

    os.environ["SAGELLM_EMBEDDING_MODEL"] = model_name
    os.environ["SAGELLM_EMBEDDING_DEVICE"] = "cpu"
    os.environ["SAGELLM_EMBEDDING_HOST"] = host
    os.environ["SAGELLM_EMBEDDING_PORT"] = str(port)

    try:
        import uvicorn

        import sagellm_core.embedding_server as embedding_server

        console.print("[bold green]🚀 Server starting — press Ctrl+C to stop[/bold green]\n")
        uvicorn.run(embedding_server.app, host=host, port=port, log_level="info")

    except Exception as e:
        console.print(f"[red]❌ Failed to start embedding server: {e}[/red]")
        sys.exit(1)


# ---------------------------------------------------------------------------
# embed  (one-shot inference, mirrors `sagellm run` for LLM)
# ---------------------------------------------------------------------------


@embedding.command("embed")
@click.argument("texts", nargs=-1, required=True)
@click.option(
    "--model",
    "-m",
    type=str,
    default=None,
    help="Embedding model id or local path (local mode only)",
)
@click.option(
    "--endpoint",
    "-e",
    type=str,
    default=None,
    help=(
        "Call a running embedding server instead of loading the model locally. "
        "Example: http://localhost:8890"
    ),
)
@click.option(
    "--json-output",
    "json_output",
    is_flag=True,
    help="Print raw JSON array of vectors (pipe-friendly)",
)
def embed(
    texts: tuple[str, ...], model: str | None, endpoint: str | None, json_output: bool
) -> None:
    """Generate embeddings for one or more text inputs.

    By default runs locally via LLMEngine (no server needed), using the same
    unified inference path as text generation.  Pass ``--endpoint`` to use
    a running server instead.

    \b
    Examples:
        sagellm embedding embed "hello world"
        sagellm embedding embed "text1" "text2" --json-output
        sagellm embedding embed "hello" --endpoint http://localhost:{_SagePorts.EMBEDDING_DEFAULT}
        sagellm embedding embed "hello" -m BAAI/bge-m3
    """
    init_default_config()
    console = get_console()

    model_name = model or DEFAULT_EMBEDDING_MODEL
    text_list = list(texts)

    if endpoint:
        # ------------------------------------------------------------------ #
        # HTTP path: delegate to a running embedding server.
        # Uses sagellm.EmbeddingClient — stdlib-only, no torch required.
        # ------------------------------------------------------------------ #
        from sagellm.embedding_client import EmbeddingClient

        client = EmbeddingClient(base_url=endpoint, default_model=model_name)
        console.print(f"[dim]Calling embedding server at {endpoint}...[/dim]")
        try:
            vectors = client.embed(text_list, model=model_name)
        except RuntimeError as exc:
            console.print(f"[red]❌ {exc}[/red]")
            sys.exit(1)

    else:
        # ------------------------------------------------------------------ #
        # Local path: use LLMEngine directly — same unified inference path as
        # text generation, no extra deps required.
        # ------------------------------------------------------------------ #
        import asyncio

        try:
            from sagellm_core.llm_engine import LLMEngine, LLMEngineConfig
        except ImportError as exc:
            console.print(f"[red]❌ sagellm-core not available: {exc}[/red]")
            sys.exit(1)

        console.print(f"[dim]Loading model {model_name}...[/dim]")
        config = LLMEngineConfig(
            model_path=model_name,
            backend_type="auto",
            task_type="embedding",
            pooling_strategy="mean",
        )
        engine = LLMEngine(config=config)

        async def _local_embed() -> list[list[float]]:
            await engine.start()
            try:
                return await engine.embed(text_list, normalize=True)
            finally:
                await engine.stop()

        try:
            vectors: list[list[float]] = asyncio.run(_local_embed())
        except Exception as exc:
            console.print(f"[red]❌ Embedding failed: {exc}[/red]")
            sys.exit(1)

    # ---------------------------------------------------------------------- #
    # Output
    # ---------------------------------------------------------------------- #
    if json_output:
        click.echo(json.dumps(vectors))
    else:
        for i, (text, vec) in enumerate(zip(text_list, vectors, strict=False)):
            preview = text[:60] + ("..." if len(text) > 60 else "")
            console.print(f"[bold][{i}][/bold] [cyan]{preview}[/cyan]")
            console.print(f"     dim={len(vec)},  first5={[round(v, 4) for v in vec[:5]]}")


# ---------------------------------------------------------------------------
# health
# ---------------------------------------------------------------------------


@embedding.command("health")
@click.option(
    "--host",
    type=str,
    default="localhost",
    show_default=True,
    help="Embedding server host",
)
@click.option(
    "--port",
    type=int,
    default=DEFAULT_EMBEDDING_PORT,
    show_default=True,
    help="Embedding server port",
)
def health(host: str, port: int) -> None:
    """Check whether the embedding server is running and healthy.

    \b
    Examples:
        sagellm embedding health
        sagellm embedding health --port 8891
    """
    from sagellm.embedding_client import EmbeddingClient

    init_default_config()
    console = get_console()

    base_url = f"http://{host}:{port}"
    client = EmbeddingClient(base_url=base_url)
    ok = client.health()

    if ok:
        console.print(f"[green]✅ Embedding server at {base_url} is healthy[/green]")
    else:
        console.print(f"[red]❌ Embedding server at {base_url} is not reachable[/red]")
        sys.exit(1)
