"""Stage 5 — bring up the runtime-only stack from the baked images and smoke it.

Follows the bring-up/poll/teardown shape of upstream's
`ci/quickstart/run_tests_quickstart.sh`: `up -d --wait`, poll /search until the
API answers, run the smoke script, tear down.
"""

from __future__ import annotations

import subprocess
import time
from typing import Dict, List

import httpx

from .config import REPO_ROOT, Config
from .enumerate_assets import load_manifest, save_manifest

COMPOSE_PROJECT = "simready-runtime"
SMOKE = REPO_ROOT / "scripts" / "smoke.sh"


def _compose_env(cfg: Config, manifest: Dict) -> Dict[str, str]:
    env = cfg.compose_env()
    env["RUN_ID"] = manifest["run_id"]
    env["ES_NAME"] = manifest["es_name"]
    # The exact tags `bake` built (manifest["images"]), not reassembled from
    # REGISTRY/IMAGE_NAME/RUN_ID here — see docker-compose.runtime.yml's
    # image: comment for why that duplication is the bug this avoids.
    # `run()` below already guarantees "images" is present before calling this.
    env["OPENSEARCH_IMAGE_TAG"] = manifest["images"]["opensearch"]
    env["NEO4J_IMAGE_TAG"] = manifest["images"]["neo4j"]
    return env


def _compose(cfg: Config, args: List[str], env: Dict[str, str]) -> subprocess.CompletedProcess:
    cmd = ["docker", "compose", "-p", COMPOSE_PROJECT, *cfg.runtime_compose_files(), *args]
    print("  $ " + " ".join(cmd[:8]) + (" ..." if len(cmd) > 8 else ""))
    return subprocess.run(cmd, env=env, check=False)


def down(cfg: Config) -> None:
    """Tear down the runtime stack a `run(keep_up=True)` left running."""
    _compose(cfg, ["down", "-v"], _compose_env(cfg, load_manifest()))


def wait_for_api(base: str, timeout: float) -> bool:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            if httpx.get(f"{base}/health", timeout=5.0).status_code == 200:
                return True
        except httpx.HTTPError:
            pass
        time.sleep(3)
    return False


def run(cfg: Config, keep_up: bool = False) -> int:
    manifest = load_manifest()
    if "images" not in manifest:
        raise SystemExit("manifest has no baked images — run `bake` first.")

    env = _compose_env(cfg, manifest)
    base = cfg.get("GATEWAY_URL", "http://localhost:8080").rstrip("/")

    print(f"Starting the runtime stack from:")
    for name, tag in manifest["images"].items():
        print(f"  {name}: {tag}")
    if cfg.siglip2_nvcf:
        print(f"siglip2 backend: nvcf ({cfg.get('SIGLIP2_NVCF_ENDPOINT', 'grpc.nvcf.nvidia.com:443')})")

    result = _compose(cfg, ["up", "-d", "--wait"], env)
    if result.returncode != 0:
        print("\nstack failed to come up; logs follow")
        _compose(cfg, ["logs", "--tail", "50"], env)
        if not keep_up:
            _compose(cfg, ["down", "-v"], env)
        return 1

    try:
        print(f"\nWaiting for the API at {base} ...")
        if not wait_for_api(base, cfg.get_float("VERIFY_TIMEOUT", 300.0)):
            print("API never became ready")
            _compose(cfg, ["logs", "--tail", "50", "deepsearch-api"], env)
            return 1

        print("\nRunning smoke tests ...\n")
        smoke_env = dict(env)
        smoke_env["BASE"] = base
        smoke = subprocess.run(
            ["bash", str(SMOKE)], cwd=str(REPO_ROOT), env=smoke_env, check=False
        )
        if smoke.returncode != 0:
            print("\nSmoke tests FAILED — the images will not be pushed.")
            return 1
        print("\nSmoke tests passed.")
        # `push` checks this — set here, not only when `verify` runs as part
        # of `all`, so running it standalone (e.g. after a manual re-run)
        # also unlocks push.
        manifest["smoke_passed"] = True
        save_manifest(manifest)
        return 0
    finally:
        if keep_up:
            print(f"\nStack left running ({base}). Tear down with:")
            print(f"  docker compose -p {COMPOSE_PROJECT} {' '.join(cfg.runtime_compose_files())} down -v")
        else:
            print("\nTearing down the runtime stack ...")
            _compose(cfg, ["down", "-v"], env)