#!/usr/bin/env python3
"""
entrypoint.py — VSS Docker Container Entrypoint (Python Version)

Boot sequence:
    Stage 0  Install dependencies from bundled requirements.txt (uv → pip fallback)
    Stage 1  Start Ollama server + pull LLM model
             (or start vLLM for LLM if VSS_LLM__USE_VLLM=true)
    Stage 2  Start vLLM VLM server (only when VSS_VLM_ENABLED=true)
    Stage 3  Verify ClickHouse reachability (non-fatal — app reconnects)
    Stage 4  exec → python3 main.py (replaces this process)

Usage:
    python -m matrice_vss.entrypoint
    # or
    python entrypoint.py
"""

from __future__ import annotations

import base64
import os
import signal
import subprocess
import sys
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional
from urllib.parse import urlparse
import urllib.request
import urllib.error

# ─── Constants ────────────────────────────────────────────────────────────────

SCRIPT_DIR = Path(__file__).parent.resolve()
LOG_DIR = Path(os.environ.get("VSS_LOG_DIR", "/app/logs"))
ENTRYPOINT_LOG = LOG_DIR / "entrypoint.log"
STARTUP_TIMEOUT_BANNER = (
    "Increase ENTRYPOINT_*_RETRIES or ENTRYPOINT_POLL_INTERVAL if the model needs more load time."
)

# Process registry for cleanup
CHILD_PIDS: list[int] = []


# ─── Logging ──────────────────────────────────────────────────────────────────


def _log(level: str, message: str) -> None:
    """Log to stderr and optionally to file."""
    ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    msg = f"[{ts}] [{level:5}] [entrypoint] {message}"
    print(msg, file=sys.stderr, flush=True)
    try:
        LOG_DIR.mkdir(parents=True, exist_ok=True)
        with open(ENTRYPOINT_LOG, "a") as f:
            f.write(msg + "\n")
    except Exception:
        pass


def log_info(msg: str) -> None:
    _log("INFO", msg)


def log_warn(msg: str) -> None:
    _log("WARN", msg)


def log_error(msg: str) -> None:
    _log("ERROR", msg)


def log_debug(msg: str) -> None:
    if os.environ.get("VSS_LOG_LEVEL", "INFO").upper() == "DEBUG":
        _log("DEBUG", msg)


def log_banner(msg: str) -> None:
    log_info("═" * 68)
    log_info(f"  {msg}")
    log_info("═" * 68)


def log_section(msg: str) -> None:
    log_info(f"──── {msg} " + "─" * (50 - len(msg)))


# ─── Process Management ───────────────────────────────────────────────────────


def register_pid(pid: int) -> None:
    """Register a background process PID for cleanup."""
    CHILD_PIDS.append(pid)
    log_debug(f"Registered background PID {pid}")


def cleanup() -> None:
    """Clean up background processes on shutdown."""
    log_info("Shutdown signal received — cleaning up background processes...")

    for pid in CHILD_PIDS:
        try:
            os.kill(pid, 0)  # Check if process exists
            log_info(f"  SIGTERM → PID {pid}")
            os.kill(pid, signal.SIGTERM)
        except ProcessLookupError:
            pass
        except Exception as e:
            log_warn(f"  Failed to signal PID {pid}: {e}")

    # Wait up to 10s for graceful exit
    for _ in range(10):
        all_gone = True
        for pid in CHILD_PIDS:
            try:
                os.kill(pid, 0)
                all_gone = False
                break
            except ProcessLookupError:
                pass
        if all_gone:
            break
        time.sleep(1)

    # Force kill remaining
    for pid in CHILD_PIDS:
        try:
            os.kill(pid, 0)
            log_warn(f"  Force-killing PID {pid} (SIGKILL)")
            os.kill(pid, signal.SIGKILL)
        except ProcessLookupError:
            pass

    log_info("Cleanup complete.")


def signal_handler(signum: int, frame) -> None:
    """Handle SIGTERM/SIGINT."""
    cleanup()
    sys.exit(143 if signum == signal.SIGTERM else 130)


# ─── URL Helpers ──────────────────────────────────────────────────────────────


def url_host(url: str) -> str:
    """Extract hostname from URL."""
    return urlparse(url).hostname or ""


def url_port(url: str) -> Optional[int]:
    """Extract port from URL."""
    return urlparse(url).port


def url_base(url: str) -> str:
    """Get scheme://host:port from URL."""
    u = urlparse(url)
    port_str = f":{u.port}" if u.port else ""
    return f"{u.scheme}://{u.hostname}{port_str}"


def is_local(host: str) -> bool:
    """Check if host is localhost."""
    return host in ("localhost", "127.0.0.1", "::1")


# ─── HTTP Health Check ────────────────────────────────────────────────────────


def wait_for_http(
    label: str,
    url: str,
    max_retries: int = 60,
    interval: int = 5,
) -> bool:
    """Wait for HTTP endpoint to become ready."""
    total = max_retries * interval
    log_info(f"Waiting for {label} → {url}")
    log_info(f"  (timeout ~{total}s | poll every {interval}s | max {max_retries} retries)")

    for attempt in range(1, max_retries + 1):
        try:
            req = urllib.request.Request(url, method="GET")
            with urllib.request.urlopen(req, timeout=5) as resp:
                if resp.status < 400:
                    log_info(f"  ✓ {label} is ready.")
                    return True
        except Exception:
            pass

        log_info(f"  [{attempt}/{max_retries}] {label} not yet ready — retrying in {interval}s...")
        time.sleep(interval)

    log_error(f"{label} did not become ready within ~{total}s.")
    log_error(STARTUP_TIMEOUT_BANNER)
    return False


def http_get_with_auth(
    url: str,
    username: str = "",
    password: str = "",
    timeout: int = 5,
) -> bool:
    """
    Perform HTTP GET with optional Basic Auth.
    Returns True if status < 400.
    """
    try:
        req = urllib.request.Request(url, method="GET")
        
        # Add Basic Auth header if credentials provided
        if username:
            credentials = f"{username}:{password}"
            encoded = base64.b64encode(credentials.encode("utf-8")).decode("ascii")
            req.add_header("Authorization", f"Basic {encoded}")
        
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status < 400
    except Exception:
        return False


# ─── Configuration ────────────────────────────────────────────────────────────


@dataclass
class Config:
    """Entrypoint configuration loaded from environment."""

    # LLM (Ollama or vLLM backend)
    llm_endpoint: str = "http://localhost:11434"
    llm_model: str = "llama3.1:8b"
    llm_gpu: str = "0"
    llm_use_vllm: bool = False
    llm_vllm_endpoint: str = "http://localhost:8001/v1"

    # VLM (Cosmos-Reason2-8B via vLLM)
    vlm_enabled: bool = False
    vlm_endpoint: str = "http://localhost:8000/v1"
    vlm_model: str = "nvidia/Cosmos-Reason2-8B"
    vlm_gpu: str = "1"
    vlm_gpu_mem: float = 0.85

    # ClickHouse
    ch_host: str = "localhost"
    ch_port: int = 8123
    ch_user: str = "default"
    ch_password: str = ""

    # API
    api_port: int = 8080

    # Entrypoint tuning
    ollama_retries: int = 60
    vlm_retries: int = 120
    llm_vllm_retries: int = 60
    ch_retries: int = 30
    poll_interval: int = 5

    # Requirements installation
    skip_requirements: bool = False

    @classmethod
    def from_env(cls) -> "Config":
        """Load configuration from environment variables."""
        def _bool(val: str) -> bool:
            return val.lower() in ("true", "1", "yes")

        return cls(
            llm_endpoint=os.environ.get("VSS_LLM__ENDPOINT", "http://localhost:11434"),
            llm_model=os.environ.get("VSS_LLM__MODEL_NAME", "llama3.1:8b"),
            llm_gpu=os.environ.get("VSS_LLM__GPU_DEVICE", "0"),
            llm_use_vllm=_bool(os.environ.get("VSS_LLM__USE_VLLM", "false")),
            llm_vllm_endpoint=os.environ.get("VSS_LLM__VLLM_ENDPOINT", "http://localhost:8001/v1"),
            vlm_enabled=_bool(os.environ.get("VSS_VLM_ENABLED", "false")),
            vlm_endpoint=os.environ.get("VSS_VLM__ENDPOINT", "http://localhost:8000/v1"),
            vlm_model=os.environ.get("VSS_VLM__MODEL_NAME", "nvidia/Cosmos-Reason2-8B"),
            vlm_gpu=os.environ.get("VSS_VLM__GPU_DEVICE", "1"),
            vlm_gpu_mem=float(os.environ.get("VSS_VLM__GPU_MEMORY_UTILIZATION", "0.85")),
            ch_host=os.environ.get("VSS_CLICKHOUSE__HOST", "clickhouse-db"),
            ch_port=int(os.environ.get("VSS_CLICKHOUSE__PORT", "8123")),
            ch_user=os.environ.get("VSS_CLICKHOUSE__USER", "matrice"),
            ch_password=os.environ.get("VSS_CLICKHOUSE__PASSWORD", ""),
            api_port=int(os.environ.get("VSS_API_PORT", "8080")),
            ollama_retries=int(os.environ.get("ENTRYPOINT_OLLAMA_RETRIES", "60")),
            vlm_retries=int(os.environ.get("ENTRYPOINT_VLM_RETRIES", "120")),
            llm_vllm_retries=int(os.environ.get("ENTRYPOINT_LLM_VLLM_RETRIES", "60")),
            ch_retries=int(os.environ.get("ENTRYPOINT_CH_RETRIES", "30")),
            poll_interval=int(os.environ.get("ENTRYPOINT_POLL_INTERVAL", "5")),
            skip_requirements=_bool(os.environ.get("VSS_SKIP_REQUIREMENTS", "false")),
        )


# ─── Stage 0: Install Requirements ────────────────────────────────────────────


def stage_requirements(config: Config) -> bool:
    """Install dependencies from bundled requirements.txt using uv or pip."""
    log_section("Stage 0 — Install Dependencies")

    if config.skip_requirements:
        log_info("Skipping requirements installation (VSS_SKIP_REQUIREMENTS=true)")
        return True

    # Find requirements.txt in package
    req_path = SCRIPT_DIR / "requirements.txt"

    if not req_path.exists():
        log_warn(f"requirements.txt not found at {req_path}")
        log_warn("Skipping dependency installation. Ensure dependencies are pre-installed.")
        return True

    log_info(f"Installing dependencies from {req_path}")

    # Try uv first (faster)
    uv_available = subprocess.run(
        ["which", "uv"],
        capture_output=True,
    ).returncode == 0

    if uv_available:
        log_info("Using uv for package installation (faster)")
        cmd = ["uv", "pip", "install", "-r", str(req_path)]
    else:
        log_info("Using pip for package installation (uv not available)")
        cmd = [sys.executable, "-m", "pip", "install", "-r", str(req_path)]

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
        )

        if result.returncode == 0:
            log_info("  ✓ Dependencies installed successfully.")
            return True
        else:
            log_error(f"Dependency installation failed (exit code {result.returncode})")
            log_error(result.stderr[:500] if result.stderr else "No error output")
            return False

    except Exception as e:
        log_error(f"Failed to run package installer: {e}")
        return False


# ─── Stage 1-A: Ollama Server ─────────────────────────────────────────────────


def stage_ollama(config: Config) -> bool:
    """Start Ollama server and pull LLM model."""
    log_section("Stage 1-A — Ollama LLM Server")

    ollama_host = url_host(config.llm_endpoint)
    ollama_port = url_port(config.llm_endpoint) or 11434
    ollama_base_url = url_base(config.llm_endpoint)

    if is_local(ollama_host):
        # Check if Ollama is installed
        ollama_installed = subprocess.run(
            ["which", "ollama"],
            capture_output=True,
        ).returncode == 0

        if not ollama_installed:
            log_info("Ollama binary not found — running official installer...")
            result = subprocess.run(
                ["bash", "-c", "curl -fsSL https://ollama.com/install.sh | sh"],
                capture_output=True,
                text=True,
            )
            if result.returncode != 0:
                log_error("Ollama installation failed.")
                return False
            log_info("Ollama installed successfully.")

        # Check if already running
        health_url = f"{ollama_base_url}/api/tags"
        try:
            with urllib.request.urlopen(health_url, timeout=3):
                log_info(f"Ollama server already responding on port {ollama_port}.")
        except Exception:
            # Start Ollama server
            log_info(f"Starting Ollama server (port={ollama_port}, CUDA_VISIBLE_DEVICES={config.llm_gpu})...")

            env = os.environ.copy()
            env["CUDA_VISIBLE_DEVICES"] = config.llm_gpu
            env["OLLAMA_HOST"] = f"0.0.0.0:{ollama_port}"
            env["OLLAMA_ORIGINS"] = "*"

            log_file = LOG_DIR / "ollama.log"
            LOG_DIR.mkdir(parents=True, exist_ok=True)

            with open(log_file, "a") as f:
                proc = subprocess.Popen(
                    ["ollama", "serve"],
                    env=env,
                    stdout=f,
                    stderr=subprocess.STDOUT,
                )

            register_pid(proc.pid)
            log_info(f"Ollama serve started (PID {proc.pid}). Logs → {log_file}")
            time.sleep(2)
    else:
        log_info(f"Ollama endpoint is remote ({ollama_host}) — skipping local start.")

    # Wait for Ollama API
    if not wait_for_http(
        "Ollama API",
        f"{ollama_base_url}/api/tags",
        config.ollama_retries,
        config.poll_interval,
    ):
        return False

    # Pull model if needed
    log_info(f"Checking whether model '{config.llm_model}' is cached in Ollama...")

    env = os.environ.copy()
    env["OLLAMA_HOST"] = f"{ollama_host}:{ollama_port}"

    result = subprocess.run(
        ["ollama", "list"],
        capture_output=True,
        text=True,
        env=env,
    )

    if config.llm_model in result.stdout:
        log_info(f"Model '{config.llm_model}' already cached — skipping pull.")
    else:
        log_info(f"Pulling model '{config.llm_model}' (this may take several minutes)...")
        subprocess.run(
            ["ollama", "pull", config.llm_model],
            env=env,
        )
        log_info(f"Model '{config.llm_model}' ready.")

    return True


# ─── Stage 1-B: vLLM LLM Server ───────────────────────────────────────────────


def stage_vllm_llm(config: Config) -> bool:
    """Start vLLM server for LLM backend."""
    log_section("Stage 1-B — vLLM LLM Backend")

    host = url_host(config.llm_vllm_endpoint)
    port = url_port(config.llm_vllm_endpoint) or 8001
    health_url = f"{config.llm_vllm_endpoint}/models"

    if is_local(host):
        # Check if already running
        try:
            with urllib.request.urlopen(health_url, timeout=3):
                log_info(f"vLLM LLM server already responding on port {port}.")
        except Exception:
            log_info(f"Starting vLLM LLM server: {config.llm_model} on port {port}")

            env = os.environ.copy()
            env["CUDA_VISIBLE_DEVICES"] = config.llm_gpu

            log_file = LOG_DIR / "vllm_llm.log"
            LOG_DIR.mkdir(parents=True, exist_ok=True)

            cmd = [
                sys.executable, "-m", "vllm.entrypoints.openai.api_server",
                "--model", config.llm_model,
                "--port", str(port),
                "--host", "0.0.0.0",
                "--gpu-memory-utilization", str(config.vlm_gpu_mem),
                "--max-model-len", "8192",
                "--trust-remote-code",
                "--dtype", "auto",
            ]

            with open(log_file, "a") as f:
                proc = subprocess.Popen(cmd, env=env, stdout=f, stderr=subprocess.STDOUT)

            register_pid(proc.pid)
            log_info(f"vLLM LLM server started (PID {proc.pid}). Logs → {log_file}")
    else:
        log_info(f"LLM vLLM endpoint is remote ({host}) — skipping local start.")

    return wait_for_http("vLLM-LLM API", health_url, config.llm_vllm_retries, config.poll_interval)


# ─── Stage 2: vLLM VLM Server ─────────────────────────────────────────────────


def stage_vllm_vlm(config: Config) -> bool:
    """Start vLLM server for VLM inference."""
    log_section(f"Stage 2 — vLLM VLM Server ({config.vlm_model})")

    host = url_host(config.vlm_endpoint)
    port = url_port(config.vlm_endpoint) or 8000
    health_url = f"{config.vlm_endpoint}/models"

    if is_local(host):
        # Check if already running
        try:
            with urllib.request.urlopen(health_url, timeout=3):
                log_info(f"vLLM VLM server already responding on port {port}.")
        except Exception:
            log_info(f"Starting vLLM VLM server: {config.vlm_model}")
            log_info(f"  port={port} | CUDA_VISIBLE_DEVICES={config.vlm_gpu} | gpu_mem={config.vlm_gpu_mem}")

            env = os.environ.copy()
            env["CUDA_VISIBLE_DEVICES"] = config.vlm_gpu

            log_file = LOG_DIR / "vllm_vlm.log"
            LOG_DIR.mkdir(parents=True, exist_ok=True)

            cmd = [
                sys.executable, "-m", "vllm.entrypoints.openai.api_server",
                "--model", config.vlm_model,
                "--port", str(port),
                "--host", "0.0.0.0",
                "--gpu-memory-utilization", str(config.vlm_gpu_mem),
                "--max-model-len", "8192",
                "--trust-remote-code",
                "--dtype", "auto",
            ]

            with open(log_file, "a") as f:
                proc = subprocess.Popen(cmd, env=env, stdout=f, stderr=subprocess.STDOUT)

            register_pid(proc.pid)
            log_info(f"vLLM VLM server started (PID {proc.pid}). Logs → {log_file}")
    else:
        log_info(f"VLM endpoint is remote ({host}) — skipping local start.")

    log_info(f"Note: loading {config.vlm_model} may take 3-10 min on first start.")
    return wait_for_http("vLLM-VLM API", health_url, config.vlm_retries, config.poll_interval)


# ─── Stage 3: ClickHouse ──────────────────────────────────────────────────────


def stage_clickhouse(config: Config) -> bool:
    """Verify ClickHouse reachability (non-fatal)."""
    log_section("Stage 3 — ClickHouse Reachability")

    # Build health URL (simple query endpoint, auth via header)
    health_url = f"http://{config.ch_host}:{config.ch_port}/?query=SELECT%201"

    log_info(f"Checking ClickHouse at {config.ch_host}:{config.ch_port}...")
    log_debug(f"  Auth: user={config.ch_user}, password={'*' * len(config.ch_password) if config.ch_password else '(none)'}")

    for attempt in range(1, config.ch_retries + 1):
        if http_get_with_auth(
            url=health_url,
            username=config.ch_user,
            password=config.ch_password,
            timeout=5,
        ):
            log_info(f"  ✓ ClickHouse is reachable at {config.ch_host}:{config.ch_port}.")
            return True

        log_info(f"  [{attempt}/{config.ch_retries}] ClickHouse not ready — retrying in {config.poll_interval}s...")
        time.sleep(config.poll_interval)

    log_warn(f"ClickHouse not reachable after {config.ch_retries * config.poll_interval}s.")
    log_warn("The VSS API will start; SQL queries will fail until ClickHouse is available.")
    return True  # Non-fatal


# ─── Stage 4: Launch Application ──────────────────────────────────────────────


def stage_application(config: Config) -> None:
    """Launch the VSS FastAPI application (replaces this process)."""
    log_section("Stage 4 — Launching VSS FastAPI Application")

    log_info(f"Working directory : {os.getcwd()}")
    log_info(f"Python            : {sys.version.split()[0]}")
    log_info(f"API bind          : http://0.0.0.0:{config.api_port}")
    log_info(f"Docs URL          : http://0.0.0.0:{config.api_port}/docs")
    log_info(f"Health URL        : http://0.0.0.0:{config.api_port}/api/v1/health")

    # exec replaces this process with uvicorn
    os.execv(
        sys.executable,
        [sys.executable, "-m", "matrice_vss.main"],
    )


# ─── Main ─────────────────────────────────────────────────────────────────────


def main() -> int:
    """Main entrypoint."""
    # Setup signal handlers
    signal.signal(signal.SIGTERM, signal_handler)
    signal.signal(signal.SIGINT, signal_handler)

    # Ensure log directory exists
    LOG_DIR.mkdir(parents=True, exist_ok=True)

    log_banner(f"VSS — Production Boot Sequence ({datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ')})")

    # Load configuration
    config = Config.from_env()

    # Display configuration
    log_info("Configuration:")
    log_info(f"  LLM backend      : {'vLLM' if config.llm_use_vllm else 'Ollama'}")
    log_info(f"  LLM endpoint     : {config.llm_endpoint}  (model: {config.llm_model}, GPU: {config.llm_gpu})")
    if config.llm_use_vllm:
        log_info(f"  LLM vLLM endpoint: {config.llm_vllm_endpoint}")
    log_info(f"  VLM enabled      : {config.vlm_enabled}")
    if config.vlm_enabled:
        log_info(f"  VLM endpoint     : {config.vlm_endpoint}  (model: {config.vlm_model}, GPU: {config.vlm_gpu})")
    log_info(f"  ClickHouse       : {config.ch_host}:{config.ch_port}")
    log_info(f"  API port         : {config.api_port}")
    log_info(f"  Log directory    : {LOG_DIR}")

    # Stage 0: Install requirements
    if not stage_requirements(config):
        log_error("FATAL: Dependency installation failed. Aborting.")
        return 1

    # Stage 1: LLM inference server
    if config.llm_use_vllm:
        if not stage_vllm_llm(config):
            log_error("FATAL: vLLM LLM backend failed to become ready. Aborting.")
            return 1
    else:
        if not stage_ollama(config):
            log_error("FATAL: Ollama LLM server failed to become ready. Aborting.")
            return 1

    # Stage 2: VLM inference server (optional)
    if config.vlm_enabled:
        if not stage_vllm_vlm(config):
            log_error("FATAL: vLLM VLM server failed to become ready. Aborting.")
            return 1
    else:
        log_section("Stage 2 — VLM Server (SKIPPED — VSS_VLM_ENABLED=false)")

    # Stage 3: ClickHouse (non-fatal)
    stage_clickhouse(config)

    # Summary
    log_banner("All required services are ready — starting VSS application")
    log_info("")
    log_info(f"  LLM  : {config.llm_endpoint}  ({config.llm_model})")
    if config.vlm_enabled:
        log_info(f"  VLM  : {config.vlm_endpoint}  ({config.vlm_model})")
    log_info(f"  DB   : {config.ch_host}:{config.ch_port}")
    log_info(f"  API  : http://0.0.0.0:{config.api_port}")
    log_info("")

    # Stage 4: Launch application
    stage_application(config)

    return 0  # Never reached due to exec


if __name__ == "__main__":
    sys.exit(main())