"""Serve command - Start sageLLM server."""

from __future__ import annotations

import json
import os
import re
import sys
import threading
import time
from urllib import error as urllib_error
from urllib import request as urllib_request

import click

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

DEFAULT_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"


def _env_bool(name: str, default: bool = False) -> bool:
    """Parse boolean env var values."""
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


def _http_get_ok(url: str, timeout_s: float = 2.0) -> bool:
    """Return True if GET returns 200."""
    try:
        with urllib_request.urlopen(url, timeout=timeout_s) as resp:  # noqa: S310
            return resp.status == 200
    except (urllib_error.URLError, TimeoutError):
        return False


def _http_post_json(url: str, payload: dict, timeout_s: float = 3.0) -> bool:
    """POST JSON and return True on 2xx."""
    req = urllib_request.Request(  # noqa: S310
        url=url,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib_request.urlopen(req, timeout=timeout_s) as resp:
            return 200 <= resp.status < 300
    except (urllib_error.URLError, TimeoutError):
        return False


def _assess_canary_output(text: str) -> tuple[bool, str]:
    """Assess output quality for startup canary.

    Returns:
        (is_healthy, reason)
    """
    content = (text or "").strip()
    if len(content) < 8:
        return False, "output too short"

    tokens = [t for t in re.split(r"\s+", content) if t]
    if len(tokens) >= 8:
        freq = {}
        for token in tokens:
            freq[token] = freq.get(token, 0) + 1
        max_repeat = max(freq.values())
        if max_repeat / len(tokens) >= 0.5:
            return False, "high token repetition"

    # Excessive non-language symbols often indicates corrupted output.
    # Keep CJK, latin letters, digits and common punctuation as valid text chars.
    weird_chars = re.findall(
        r"[^\u4e00-\u9fffA-Za-z0-9\s，。！？、；：“”‘’（）()【】\-_,.:;'\"!?/]", content
    )
    if len(weird_chars) / max(len(content), 1) > 0.15:
        return False, "abnormal symbol ratio"

    return True, "ok"


def _run_startup_canary(engine_port: int, model_name: str, console) -> None:
    """Run startup canary against local engine endpoint and fail fast on bad output."""
    canary_url = f"http://localhost:{engine_port}/v1/chat/completions"
    payload = {
        "model": model_name,
        "messages": [{"role": "user", "content": "请用中文一句话介绍SAGE。"}],
        "max_tokens": 48,
        "temperature": 0.7,
        "top_p": 0.9,
        "stream": False,
    }

    req = urllib_request.Request(  # noqa: S310
        url=canary_url,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    try:
        with urllib_request.urlopen(req, timeout=30.0) as resp:
            if not (200 <= resp.status < 300):
                raise RuntimeError(f"canary http status {resp.status}")
            body = json.loads(resp.read().decode("utf-8"))
    except Exception as e:
        raise RuntimeError(f"startup canary request failed: {e}") from e

    output_text = ""
    choices = body.get("choices") if isinstance(body, dict) else None
    if isinstance(choices, list) and choices:
        message = choices[0].get("message", {}) if isinstance(choices[0], dict) else {}
        output_text = message.get("content", "") if isinstance(message, dict) else ""

    ok, reason = _assess_canary_output(output_text)
    if not ok:
        preview = output_text[:120].replace("\n", " ")
        raise RuntimeError(f"startup canary failed: {reason}; preview='{preview}'")

    console.print("[green]✅ Startup canary passed[/green]")


def _get_model_candidates(primary_model: str) -> list[str]:
    """Build ordered model candidates from primary + fallback env config."""
    raw = os.getenv("SAGELLM_FALLBACK_MODEL", "").strip()
    fallbacks = [m.strip() for m in raw.split(",") if m.strip()]
    models: list[str] = [primary_model]
    for model in fallbacks:
        if model not in models:
            models.append(model)
    return models


def _preflight_model_canary_local(model_name: str) -> tuple[bool, str]:
    """Run local preflight canary via transformers before booting services."""
    try:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None,
        )
        messages = [{"role": "user", "content": "请用中文一句话介绍SAGE。"}]
        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

        inputs = tokenizer(prompt, return_tensors="pt")
        if torch.cuda.is_available():
            inputs = {key: value.to(model.device) for key, value in inputs.items()}

        output = model.generate(
            **inputs,
            max_new_tokens=48,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            top_k=40,
        )
        text = tokenizer.decode(
            output[0][inputs["input_ids"].shape[1] :],
            skip_special_tokens=True,
        )
        ok, reason = _assess_canary_output(text)
        if not ok:
            preview = text[:120].replace("\n", " ")
            return False, f"{reason}; preview='{preview}'"
        return True, "ok"
    except Exception as e:
        return False, f"preflight exception: {e}"


def _start_periodic_canary_monitor(engine_port: int, model_name: str, console) -> None:
    """Start background periodic canary checks and fail fast on sustained failures."""
    interval_sec = int(os.getenv("SAGELLM_CANARY_INTERVAL_SEC", "300"))
    fail_threshold = int(os.getenv("SAGELLM_CANARY_FAIL_THRESHOLD", "3"))

    def _runner() -> None:
        consecutive_failures = 0
        while True:
            time.sleep(max(interval_sec, 30))
            try:
                _run_startup_canary(engine_port, model_name, console)
                consecutive_failures = 0
            except Exception as e:
                consecutive_failures += 1
                console.print(
                    "[yellow]⚠️ Periodic canary failed "
                    f"({consecutive_failures}/{fail_threshold}): {e}[/yellow]"
                )
                if consecutive_failures >= fail_threshold:
                    console.print(
                        "[red]❌ Canary circuit breaker tripped. "
                        "Shutting down service for safe recovery.[/red]"
                    )
                    os._exit(78)

    thread = threading.Thread(target=_runner, name="sagellm-canary-monitor", daemon=True)
    thread.start()


@click.command()
@click.option(
    "--backend",
    type=str,
    default=None,
    help="Backend provider kind: cpu/cuda/ascend",
)
@click.option("--model", type=str, default=None, help="Model path or name")
@click.option("--port", type=int, default=8000, help="API server port")
@click.option("--host", type=str, default="0.0.0.0", help="Server host")
@click.option("--workers", type=int, default=1, help="Number of uvicorn workers")
@click.option(
    "--with-embedding/--no-embedding",
    default=_env_bool("SAGELLM_WITH_EMBEDDING", False),
    help="Start embedding engine together in full-stack mode",
)
@click.option(
    "--embedding-model",
    type=str,
    default=os.getenv("SAGELLM_EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2"),
    help="Embedding model id used by --with-embedding",
)
@click.option(
    "--embedding-port",
    type=int,
    default=int(os.getenv("SAGELLM_EMBEDDING_PORT", "8890")),
    help="Embedding server port used by --with-embedding",
)
@click.option(
    "--engine-port",
    type=int,
    default=int(os.getenv("SAGELLM_ENGINE_PORT", "0")),
    help="LLM engine HTTP port in full-stack mode (default: gateway port + 1)",
)
def serve(
    backend: str | None,
    model: str | None,
    port: int,
    host: str,
    workers: int,
    with_embedding: bool,
    embedding_model: str,
    embedding_port: int,
    engine_port: int,
) -> None:
    """Start sageLLM server with OpenAI-compatible API (like vLLM).

    By default, starts the full stack: Gateway + Control Plane + Engine.
    This provides an OpenAI-compatible REST API out of the box.

    \b
    Examples:
        sage-llm serve --model Qwen2-7B              # Full stack (recommended)
        sage-llm serve --backend cuda --model xxx    # GPU inference
    """
    init_default_config()
    console = get_console()
    _serve_full_stack(
        backend,
        model,
        port,
        host,
        workers,
        console,
        with_embedding=with_embedding,
        embedding_model=embedding_model,
        embedding_port=embedding_port,
        engine_port=engine_port if engine_port else port + 1,
    )


# Standalone embedding serve is now owned by `sagellm embedding serve`
# (cli.commands.embedding).  The old private helpers _serve_engine_debug and
# _serve_embedding_engine have been removed; use `sagellm embedding serve` instead.


def _start_server_in_background(
    app_import_path: str,
    env_vars: dict[str, str],
    port: int,
    thread_name: str,
    console,
) -> None:
    """Generic helper: start any FastAPI/ASGI app in a background daemon thread.

    ``app_import_path`` is a dotted string accepted by uvicorn, e.g.
    ``"sagellm_core.engine_server:app"``.  All ``env_vars`` are written to
    ``os.environ`` **before** the thread is started so the forked environment
    inherits them.

    Both ``_start_llm_engine_in_background`` and
    ``_start_embedding_engine_in_background`` delegate to this helper so the
    background-launch pattern is defined in one place.
    """
    for key, value in env_vars.items():
        os.environ[key] = value

    def _runner() -> None:
        try:
            import uvicorn

            uvicorn.run(app_import_path, host="0.0.0.0", port=port, log_level="warning")
        except Exception as exc:
            console.print(f"[yellow]⚠️ Server thread '{thread_name}' exited: {exc}[/yellow]")

    thread = threading.Thread(target=_runner, name=thread_name, daemon=True)
    thread.start()


def _start_embedding_engine_in_background(
    embedding_model: str,
    embedding_port: int,
    console,
) -> None:
    """Start embedding server in a background daemon thread.

    Delegates to ``_start_server_in_background`` so the launch pattern is
    shared with ``_start_llm_engine_in_background``.
    """
    _start_server_in_background(
        app_import_path="sagellm_core.embedding_server:app",
        env_vars={
            "SAGELLM_EMBEDDING_MODEL": embedding_model,
            "SAGELLM_EMBEDDING_DEVICE": "cpu",
            "SAGELLM_EMBEDDING_ENGINE_ID": f"embedding-cpu-{embedding_port}",
            "SAGELLM_EMBEDDING_HOST": "localhost",
            "SAGELLM_EMBEDDING_PORT": str(embedding_port),
        },
        port=embedding_port,
        thread_name="sagellm-embedding-server",
        console=console,
    )


def _register_embedding_engine_after_gateway_ready(
    gateway_host: str,
    gateway_port: int,
    embedding_model: str,
    embedding_port: int,
    console,
) -> None:
    """Register embedding engine to gateway via management API."""
    connect_host = "localhost" if gateway_host in ("0.0.0.0", "::") else gateway_host
    register_url = f"http://{connect_host}:{gateway_port}/v1/management/engines/register"
    payload = {
        "engine_id": f"embedding-cpu-{embedding_port}",
        "model_id": embedding_model,
        "host": "localhost",
        "port": embedding_port,
        "engine_kind": "embedding",
    }

    deadline = time.time() + 30.0
    while time.time() < deadline:
        if _http_post_json(register_url, payload):
            console.print("[green]✅ Embedding engine registered to gateway[/green]")
            return
        time.sleep(1.0)

    console.print("[yellow]⚠️ Embedding registration timeout after 30s[/yellow]")


def _start_llm_engine_in_background(
    backend_kind: str,
    model_name: str,
    engine_port: int,
    dtype: str,
    console,
) -> None:
    """Start LLM engine HTTP server in a background daemon thread.

    Delegates to ``_start_server_in_background`` so the launch pattern is
    shared with ``_start_embedding_engine_in_background``.
    """
    _start_server_in_background(
        app_import_path="sagellm_core.engine_server:app",
        env_vars={
            "SAGELLM_MODEL_PATH": model_name,
            "SAGELLM_DEVICE": backend_kind,
            "SAGELLM_DTYPE": dtype,
            "SAGELLM_ENGINE_ID": f"engine-{backend_kind}-{engine_port}",
        },
        port=engine_port,
        thread_name="sagellm-llm-engine",
        console=console,
    )


def _register_llm_engine_after_gateway_ready(
    gateway_host: str,
    gateway_port: int,
    model_name: str,
    engine_port: int,
    backend_kind: str,
    console,
) -> None:
    """Register LLM engine to gateway via management API."""
    # Always connect to localhost for registration — the bind address
    # (e.g. 0.0.0.0) is not necessarily reachable as a connect target.
    connect_host = "localhost" if gateway_host in ("0.0.0.0", "::") else gateway_host
    register_url = f"http://{connect_host}:{gateway_port}/v1/management/engines/register"
    payload = {
        "engine_id": f"engine-{backend_kind}-{engine_port}",
        "model_id": model_name,
        "host": "localhost",
        "port": engine_port,
        "engine_kind": "llm",
        "metadata": {"pre_verified_healthy": True},
    }

    deadline = time.time() + 60.0
    while time.time() < deadline:
        if _http_post_json(register_url, payload):
            console.print("[green]✅ LLM engine registered to gateway[/green]")
            return
        time.sleep(1.0)

    console.print("[yellow]⚠️ LLM engine registration timeout after 60s[/yellow]")


def _serve_full_stack(
    backend: str | None,
    model: str | None,
    port: int,
    host: str,
    workers: int,
    console,
    with_embedding: bool = False,
    embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
    embedding_port: int = 8890,
    engine_port: int = 0,
) -> None:
    """Start full stack: Gateway + Control Plane + Engine (like vLLM)."""
    console.print("\n🚀 [bold blue]Starting sageLLM server (full stack)...[/bold blue]\n")

    backend_kind, engine_kind = resolve_backend_and_engine(backend=backend)
    model_name = model or DEFAULT_MODEL

    if _env_bool("SAGELLM_PREFLIGHT_CANARY", True):
        candidates = _get_model_candidates(model_name)
        selected_model: str | None = None
        for candidate in candidates:
            console.print(f"[cyan]🔎 Preflight canary for model: {candidate}[/cyan]")
            ok, reason = _preflight_model_canary_local(candidate)
            if ok:
                selected_model = candidate
                if candidate != model_name:
                    console.print(
                        f"[yellow]⚠️ Primary model failed preflight; switched to fallback: {candidate}[/yellow]"
                    )
                break
            console.print(f"[yellow]⚠️ Preflight failed for {candidate}: {reason}[/yellow]")

        if selected_model is None:
            console.print("[red]❌ All model candidates failed preflight canary.[/red]")
            sys.exit(1)
        model_name = selected_model
    else:
        console.print("[dim]Preflight canary disabled (SAGELLM_PREFLIGHT_CANARY=0)[/dim]")

    llm_engine_port = engine_port or port + 1

    console.print(f"  [green]Backend: {backend_kind}[/green]")
    console.print(f"  Engine: {engine_kind}")
    console.print(f"  Model: {model_name}")
    console.print(f"  LLM engine port: {llm_engine_port}")
    console.print(f"  Embedding enabled: {with_embedding}")
    console.print(f"  API endpoint: http://{host}:{port}/v1/chat/completions")
    console.print(f"  API docs: http://{host}:{port}/docs\n")

    try:
        # Check dependencies
        import uvicorn

        try:
            import sagellm_gateway  # noqa: F401
        except ImportError:
            console.print("[red]❌ Gateway not installed[/red]")
            console.print("\n[yellow]💡 Install full stack: pip install 'isagellm'[/yellow]")
            sys.exit(1)

        # --- Start LLM Engine in background ---
        _start_llm_engine_in_background(backend_kind, model_name, llm_engine_port, "auto", console)

        llm_health_url = f"http://localhost:{llm_engine_port}/health"
        llm_deadline = time.time() + 60.0
        llm_healthy = False
        while time.time() < llm_deadline:
            if _http_get_ok(llm_health_url):
                llm_healthy = True
                break
            time.sleep(1.0)

        if llm_healthy:
            console.print(
                f"[green]✅ LLM engine is healthy at"
                f" http://localhost:{llm_engine_port}/health[/green]"
            )

            if _env_bool("SAGELLM_STARTUP_CANARY", True):
                console.print("[cyan]🔎 Running startup canary check...[/cyan]")
                try:
                    _run_startup_canary(llm_engine_port, model_name, console)
                except Exception as canary_error:
                    console.print(f"[red]❌ Startup canary failed: {canary_error}[/red]")
                    console.print(
                        "[red]Service startup aborted to protect users from unhealthy model output.[/red]"
                    )
                    sys.exit(1)
            else:
                console.print("[dim]Startup canary disabled (SAGELLM_STARTUP_CANARY=0)[/dim]")

            if _env_bool("SAGELLM_PERIODIC_CANARY", True):
                console.print("[cyan]🔁 Starting periodic canary monitor...[/cyan]")
                _start_periodic_canary_monitor(llm_engine_port, model_name, console)
            else:
                console.print("[dim]Periodic canary disabled (SAGELLM_PERIODIC_CANARY=0)[/dim]")

            threading.Thread(
                target=_register_llm_engine_after_gateway_ready,
                args=(host, port, model_name, llm_engine_port, backend_kind, console),
                name="sagellm-llm-register",
                daemon=True,
            ).start()
        else:
            console.print("[yellow]⚠️ LLM engine health check timed out after 60s[/yellow]")

        # --- Start Embedding Engine in background (if requested) ---
        if with_embedding:
            _start_embedding_engine_in_background(embedding_model, embedding_port, console)

            health_url = f"http://localhost:{embedding_port}/health"
            deadline = time.time() + 30.0
            healthy = False
            while time.time() < deadline:
                if _http_get_ok(health_url):
                    healthy = True
                    break
                time.sleep(1.0)

            if healthy:
                console.print(
                    f"[green]✅ Embedding engine is healthy at http://localhost:{embedding_port}/health[/green]"
                )

                registration_thread = threading.Thread(
                    target=_register_embedding_engine_after_gateway_ready,
                    args=(host, port, embedding_model, embedding_port, console),
                    name="sagellm-embedding-register",
                    daemon=True,
                )
                registration_thread.start()
            else:
                console.print(
                    "[yellow]⚠️ Embedding engine health check timed out after 30s[/yellow]"
                )

        console.print("[bold green]✅ Starting Gateway...[/bold green]\n")

        # Start Gateway with uvicorn (this provides the REST API)
        uvicorn.run(
            "sagellm_gateway.server:app",
            host=host,
            port=port,
            workers=workers,
            log_level="info",
        )

    except ImportError as e:
        console.print(f"[red]❌ Import error: {e}[/red]")
        console.print("\n[yellow]💡 Install full stack: pip install 'isagellm'[/yellow]")
        sys.exit(1)
    except Exception as e:
        console.print(f"[red]❌ Failed to start server: {e}[/red]")
        sys.exit(1)
