import os
import socket
import signal
import subprocess
import sys
import time
import uuid
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from pathlib import Path
from typing import Final

import httpx
import psutil

from integration._support.client import Gateway


def in_group(process: psutil.Process, group: int) -> bool:
    try:
        return os.getpgid(process.pid) == group
    except ProcessLookupError:
        return False


def group_members(group: int) -> tuple[psutil.Process, ...]:
    return tuple(process for process in psutil.process_iter() if in_group(process, group))


def signal_group(group: int, action: int) -> None:
    try:
        os.killpg(group, action)
    except ProcessLookupError:
        pass


@contextmanager
def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str]) -> Iterator[Gateway]:
    with socket.socket() as reserve:
        reserve.bind(("127.0.0.1", 0))
        port: Final = reserve.getsockname()[1]
    root: Final = Path(__file__).resolve().parents[3]
    environment: Final = {
        **os.environ,
        "LITELLM_MASTER_KEY": gateway.key,
        "LITELLM_SALT_KEY": os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt"),
        "STORE_MODEL_IN_DB": "True",
        **overrides,
    }
    output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory)))
    output.mkdir(parents=True, exist_ok=True)
    with (output / f"owned-proxy-{uuid.uuid4().hex}.log").open("w") as log:
        process: Final = subprocess.Popen(
            [sys.executable, "-m", "integration._support.proxy", "--config", "tests/integration/proxy_config.yaml",
             "--host", "127.0.0.1", "--port", str(port), "--num_workers", "1", "--telemetry", "False",
             "--use_prisma_db_push", "--enforce_prisma_migration_check"],
            cwd=root, env=environment, stdout=log, stderr=subprocess.STDOUT, start_new_session=True,
        )
        try:
            with httpx.Client(base_url=f"http://127.0.0.1:{port}", timeout=15, trust_env=False) as client:
                deadline: Final = time.monotonic() + 70
                while True:
                    assert process.poll() is None, "Owned proxy exited before readiness"
                    try:
                        if client.get("/health/readiness", timeout=2).status_code == 200:
                            break
                    except httpx.TransportError:
                        pass
                    assert time.monotonic() < deadline, "Owned proxy readiness deadline exceeded"
                    time.sleep(0.1)
                yield Gateway(client, gateway.key, gateway.upstream_url)
        finally:
            forced = False
            if process.poll() is None:
                process.terminate()
                try:
                    process.wait(timeout=30)
                except subprocess.TimeoutExpired:
                    forced = True
            residual: Final = group_members(process.pid)
            if residual:
                signal_group(process.pid, signal.SIGTERM)
                psutil.wait_procs(residual, timeout=5)
            remaining: Final = group_members(process.pid)
            if remaining:
                forced = True
                signal_group(process.pid, signal.SIGKILL)
                psutil.wait_procs(remaining, timeout=3)
            process.wait(timeout=3)
            survivors: Final = group_members(process.pid)
            assert not survivors, "Owned proxy child survived cleanup"
            assert not forced, "Owned proxy required forced cleanup"