"""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 EmbeddingEngine (no server needed)
                                  or HTTP call via --endpoint <url>
* ``sagellm embedding health`` → HTTP health check of a running server

SAGE compatibility
------------------
``sagellm.embedding.EmbeddingService.from_endpoint()`` and
``sagellm.embedding.compat.get_embedding_model("sagellm", base_url=...)``
both default to ``http://127.0.0.1:8890``, 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 ..utils.config import init_default_config
from ..utils.console import get_console

DEFAULT_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
DEFAULT_EMBEDDING_PORT = 8890


@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 — no LLM stack required.

    \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")

    # Env vars are read by sagellm_core.embedding_server at startup via
    # its FastAPI lifespan handler (see sagellm_core/embedding_server.py).
    os.environ["SAGELLM_EMBEDDING_MODEL"] = model_name
    os.environ["SAGELLM_EMBEDDING_DEVICE"] = "cpu"
    os.environ["SAGELLM_EMBEDDING_ENGINE_ID"] = f"embedding-cpu-{port}"
    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 ImportError as e:
        console.print(f"[red]❌ Import error: {e}[/red]")
        console.print(
            "\n[yellow]💡 Install embedding deps: pip install 'isagellm[embedding]'[/yellow]"
        )
        sys.exit(1)
    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 EmbeddingEngine (no server needed), mirroring
    how ``sagellm run`` uses LLMEngine directly.  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:8890
        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.embedding.EmbeddingClient — the same component that
        # neuromem / sage-common use programmatically.
        # ------------------------------------------------------------------ #
        from sagellm.embedding 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 EmbeddingEngine directly — mirrors how `sagellm run`
        # uses LLMEngine without going through an HTTP hop.
        #
        # LLMEngine lifecycle:   await engine.start() → generate() → stop()
        # EmbeddingEngine lifecycle: engine.setup() → encode() → teardown()
        #
        # Both engines are standalone classes in sagellm-core; they share the
        # same lifecycle pattern but are not the same class (sentence-
        # transformers vs. causal-LM) so they can't share implementation,
        # but the CLI command structure is intentionally symmetric.
        # ------------------------------------------------------------------ #
        try:
            from sagellm_core.engines.embedding import EmbeddingEngine, EmbeddingEngineConfig
        except ImportError as exc:
            console.print(f"[red]❌ sagellm-core not available: {exc}[/red]")
            console.print(
                "[yellow]💡 pip install 'isagellm[embedding]' or use --endpoint <url>[/yellow]"
            )
            sys.exit(1)

        console.print(f"[dim]Loading model {model_name}...[/dim]")
        config = EmbeddingEngineConfig(
            model_id=model_name,
            device="auto",
            normalize=True,
        )
        engine = EmbeddingEngine(config)
        engine.setup()
        try:
            raw = engine.encode(text_list)
        finally:
            engine.teardown()

        # encode(list[str]) always returns list[list[float]], but guard anyway
        vectors: list[list[float]] = [raw] if isinstance(raw[0], float) else raw  # type: ignore[arg-type]

    # ---------------------------------------------------------------------- #
    # 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 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)
