"""Sidecar process: owns HTTP I/O, system monitoring, and health watching.

Launched via multiprocessing.Process(target=run_sidecar, kwargs={"config": ...}).
Uses asyncio internally — no framework conflicts with the training code.

Architecture:
- _socket_reader: accepts Unix socket connections, parses newline-JSON, enqueues messages
- _api_sender: drains queue, batches log entries, sends HTTP to Tyko API
- _system_monitor: polls psutil/pynvml at interval, enqueues system/ metrics
- _health_watcher: polls main PID; on death, marks run failed and shuts down
- _command_poller: polls GET /runs/{id}/commands every 5s; writes stop sentinel file
"""

import asyncio
import json
import os
import signal
from dataclasses import dataclass, field
from typing import Any

import httpx


@dataclass
class SidecarConfig:
    """Configuration passed to the sidecar at launch time.

    All fields are primitives so this is picklable under the 'spawn'
    multiprocessing start method (default on macOS).
    """

    client_uuid: str
    api_url: str
    api_key: str | None
    main_pid: int
    socket_path: str
    # Set per-run via "set_run" message; None until first run starts
    run_id: str | None = field(default=None)
    system_monitor_interval: float | None = field(default=15.0)
    batch_window_s: float = field(default=0.5)
    command_poll_interval: float = field(default=5.0)


def stop_sentinel_path(*, client_uuid: str) -> str:
    """Path to the stop sentinel file for this client."""
    return f"/tmp/tyko-{client_uuid}.stop"


def run_sidecar(*, config: SidecarConfig) -> None:
    """Entrypoint for multiprocessing.Process. Runs the asyncio event loop."""
    # Ignore SIGINT in the sidecar — the main process handles it.
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    asyncio.run(_main(config=config))


async def _main(*, config: SidecarConfig) -> None:
    """Coordinate the four asyncio tasks."""
    queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
    shutdown_event = asyncio.Event()
    # Mutable dict shared between tasks (avoids per-task config copies)
    state: dict[str, Any] = {
        "run_id": config.run_id,
        "system_monitor_interval": config.system_monitor_interval,
    }

    tasks = [
        asyncio.create_task(_socket_reader(config=config, queue=queue, shutdown_event=shutdown_event, state=state)),
        asyncio.create_task(_api_sender(config=config, queue=queue, shutdown_event=shutdown_event, state=state)),
        asyncio.create_task(_health_watcher(config=config, queue=queue, shutdown_event=shutdown_event)),
        asyncio.create_task(_system_monitor(config=config, queue=queue, shutdown_event=shutdown_event, state=state)),
        asyncio.create_task(_command_poller(config=config, shutdown_event=shutdown_event, state=state)),
    ]

    await asyncio.gather(*tasks, return_exceptions=True)


# ---------------------------------------------------------------------------
# Task 1: Socket reader
# ---------------------------------------------------------------------------


async def _socket_reader(
    *,
    config: SidecarConfig,
    queue: asyncio.Queue[dict[str, Any]],
    shutdown_event: asyncio.Event,
    state: dict[str, Any],
) -> None:
    """Accept Unix socket connections, parse newline-delimited JSON, enqueue messages.

    Sends {"ok": true} to the client for "set_run", "close", and "shutdown".
    """
    server: asyncio.Server | None = None
    try:
        # Remove stale socket file if it exists
        try:
            os.unlink(config.socket_path)
        except FileNotFoundError:
            pass

        server = await asyncio.start_unix_server(
            lambda r, w: _handle_client(r, w, queue=queue, state=state),
            path=config.socket_path,
        )

        async with server:
            await shutdown_event.wait()

    finally:
        try:
            os.unlink(config.socket_path)
        except FileNotFoundError:
            pass


async def _handle_client(
    reader: asyncio.StreamReader,
    writer: asyncio.StreamWriter,
    *,
    queue: asyncio.Queue[dict[str, Any]],
    state: dict[str, Any],
) -> None:
    """Handle one client connection: read messages until EOF."""
    try:
        while True:
            line = await reader.readline()
            if not line:
                break
            try:
                message = json.loads(line.decode().strip())
            except json.JSONDecodeError:
                continue

            msg_type = message.get("type")

            if msg_type == "set_run":
                state["run_id"] = message.get("run_id")
                state["system_monitor_interval"] = message.get("system_monitor_interval")
                await queue.put(message)
                writer.write(b'{"ok": true}\n')
                await writer.drain()

            elif msg_type == "close":
                await queue.put(message)
                writer.write(b'{"ok": true}\n')
                await writer.drain()

            elif msg_type == "shutdown":
                await queue.put(message)
                writer.write(b'{"ok": true}\n')
                await writer.drain()
                break

            else:
                # fire-and-forget: log, params, environment, status
                await queue.put(message)
    except (asyncio.IncompleteReadError, ConnectionResetError, BrokenPipeError):
        pass
    finally:
        try:
            writer.close()
            await writer.wait_closed()
        except Exception:
            pass


# ---------------------------------------------------------------------------
# Task 2: API sender
# ---------------------------------------------------------------------------


async def _api_sender(
    *,
    config: SidecarConfig,
    queue: asyncio.Queue[dict[str, Any]],
    shutdown_event: asyncio.Event,
    state: dict[str, Any],
) -> None:
    """Drain the queue and send HTTP requests to the Tyko API.

    Log entries are batched within a window (batch_window_s) for efficiency.
    """
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            pending_entries: list[dict[str, Any]] = []
            pending_console_logs: list[dict[str, Any]] = []
            batch_deadline: float | None = None

            while True:
                # Determine how long to wait for the next message
                now = asyncio.get_event_loop().time()
                has_pending = bool(pending_entries or pending_console_logs)
                if batch_deadline is not None and has_pending:
                    wait = max(0.0, batch_deadline - now)
                else:
                    wait = None  # block until next message

                try:
                    message = await asyncio.wait_for(queue.get(), timeout=wait)
                except asyncio.TimeoutError:
                    # Batch window expired — flush pending batches
                    await _flush_entries(pending_entries, state=state, config=config, client=client)
                    await _flush_console_logs(pending_console_logs, state=state, config=config, client=client)
                    pending_entries = []
                    pending_console_logs = []
                    batch_deadline = None
                    continue

                msg_type = message.get("type")

                if msg_type == "log":
                    entry = {
                        "timestamp": message["timestamp"],
                        "relative_time_ns": message["relative_time_ns"],
                        "values": message["values"],
                    }
                    pending_entries.append(entry)
                    if batch_deadline is None:
                        batch_deadline = asyncio.get_event_loop().time() + config.batch_window_s

                elif msg_type == "console":
                    log_entry = {
                        "stream": message["stream"],
                        "text": message["text"],
                        "timestamp": message["timestamp"],
                        "relative_time_ns": message.get("relative_time_ns"),
                    }
                    pending_console_logs.append(log_entry)
                    if batch_deadline is None:
                        batch_deadline = asyncio.get_event_loop().time() + config.batch_window_s

                elif msg_type == "params":
                    run_id = state.get("run_id")
                    if run_id:
                        await _patch_run(
                            {"params": message["values"]},
                            run_id=run_id,
                            config=config,
                            client=client,
                        )

                elif msg_type == "environment":
                    run_id = state.get("run_id")
                    if run_id:
                        await _patch_run(
                            {"environment": message["values"]},
                            run_id=run_id,
                            config=config,
                            client=client,
                        )

                elif msg_type == "status":
                    run_id = state.get("run_id")
                    if run_id:
                        await _patch_run(
                            {"status": message["status"]},
                            run_id=run_id,
                            config=config,
                            client=client,
                        )

                elif msg_type == "set_run":
                    # run_id already set in state by socket reader
                    pass

                elif msg_type == "close":
                    # Flush pending batches first, then send final status
                    await _flush_entries(pending_entries, state=state, config=config, client=client)
                    await _flush_console_logs(pending_console_logs, state=state, config=config, client=client)
                    pending_entries = []
                    pending_console_logs = []
                    batch_deadline = None
                    run_id = state.get("run_id")
                    if run_id:
                        payload: dict[str, Any] = {"status": message["status"]}
                        if message.get("at"):
                            payload["at"] = message["at"]
                        await _patch_run(payload, run_id=run_id, config=config, client=client)

                elif msg_type == "shutdown":
                    # Final flush
                    await _flush_entries(pending_entries, state=state, config=config, client=client)
                    await _flush_console_logs(pending_console_logs, state=state, config=config, client=client)
                    shutdown_event.set()
                    return

                queue.task_done()

                # Flush if batch window expired while processing other messages
                now = asyncio.get_event_loop().time()
                if batch_deadline is not None and now >= batch_deadline and (pending_entries or pending_console_logs):
                    await _flush_entries(pending_entries, state=state, config=config, client=client)
                    await _flush_console_logs(pending_console_logs, state=state, config=config, client=client)
                    pending_entries = []
                    pending_console_logs = []
                    batch_deadline = None

    except Exception:
        shutdown_event.set()


async def _flush_entries(
    entries: list[dict[str, Any]],
    *,
    state: dict[str, Any],
    config: SidecarConfig,
    client: Any,
) -> None:
    if not entries:
        return
    run_id = state.get("run_id")
    if not run_id:
        return
    url = f"{config.api_url}/runs/{run_id}/entries"
    headers = _auth_headers(config)
    try:
        await client.post(url, json={"entries": entries}, headers=headers)
    except Exception:
        pass


async def _flush_console_logs(
    logs: list[dict[str, Any]],
    *,
    state: dict[str, Any],
    config: SidecarConfig,
    client: Any,
) -> None:
    if not logs:
        return
    run_id = state.get("run_id")
    if not run_id:
        return
    url = f"{config.api_url}/runs/{run_id}/logs"
    headers = _auth_headers(config)
    try:
        await client.post(url, json={"logs": logs}, headers=headers)
    except Exception:
        pass


async def _patch_run(
    payload: dict[str, Any],
    *,
    run_id: str,
    config: SidecarConfig,
    client: Any,
) -> None:
    url = f"{config.api_url}/runs/{run_id}"
    headers = _auth_headers(config)
    try:
        await client.patch(url, json=payload, headers=headers)
    except Exception:
        pass


def _auth_headers(config: SidecarConfig) -> dict[str, str]:
    if config.api_key:
        return {"Authorization": f"Bearer {config.api_key}"}
    return {}


# ---------------------------------------------------------------------------
# Task 3: System monitor
# ---------------------------------------------------------------------------


async def _system_monitor(
    *,
    config: SidecarConfig,
    queue: asyncio.Queue[dict[str, Any]],
    shutdown_event: asyncio.Event,
    state: dict[str, Any],
) -> None:
    """Poll system metrics and enqueue log messages at the configured interval.

    Waits until a run is active (state["run_id"] is set) before sampling.
    Stops when shutdown_event is set.
    """
    loop = asyncio.get_event_loop()
    while not shutdown_event.is_set():
        interval = state.get("system_monitor_interval")
        run_id = state.get("run_id")

        if interval is None or run_id is None:
            # Not yet configured or monitoring disabled — check again soon
            try:
                await asyncio.wait_for(shutdown_event.wait(), timeout=0.1)
            except asyncio.TimeoutError:
                pass
            continue

        # Wait for the interval, then sample
        try:
            await asyncio.wait_for(shutdown_event.wait(), timeout=interval)
            return  # shutdown_event set during wait
        except asyncio.TimeoutError:
            pass

        metrics = await loop.run_in_executor(None, _collect_system_metrics)
        if metrics and state.get("run_id"):
            import time

            now_ns = time.monotonic_ns()
            await queue.put(
                {
                    "type": "log",
                    "values": metrics,
                    "timestamp": now_ns,
                    "relative_time_ns": now_ns,  # sidecar doesn't know run start time
                }
            )


def _collect_system_metrics() -> dict[str, float]:
    """Collect all available system metrics using psutil and pynvml."""
    metrics: dict[str, float] = {}
    try:
        import psutil  # type: ignore[import-untyped]

        metrics["system/cpu_percent"] = psutil.cpu_percent(interval=None)
        vm = psutil.virtual_memory()
        metrics["system/ram_used_gb"] = round(vm.used / (1024**3), 3)
        metrics["system/ram_percent"] = vm.percent
        metrics["system/process_ram_gb"] = round(psutil.Process().memory_info().rss / (1024**3), 3)
    except Exception:
        pass

    metrics.update(_collect_gpu_metrics())
    return metrics


# ---------------------------------------------------------------------------
# GPU via pynvml
# ---------------------------------------------------------------------------


def _collect_gpu_metrics() -> dict[str, float]:
    """Query GPU metrics via pynvml."""
    try:
        import pynvml  # type: ignore[import-untyped]

        pynvml.nvmlInit()
        try:
            count = pynvml.nvmlDeviceGetCount()
            gpu_metrics: dict[str, float] = {}
            for i in range(count):
                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
                suffix = f"_{i}" if count > 1 else ""
                util = pynvml.nvmlDeviceGetUtilizationRates(handle)
                gpu_metrics[f"system/gpu_util{suffix}"] = float(util.gpu)
                mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
                gpu_metrics[f"system/gpu_mem_used_gb{suffix}"] = round(mem.used / (1024**3), 3)
                if mem.total > 0:
                    gpu_metrics[f"system/gpu_mem_percent{suffix}"] = round(mem.used / mem.total * 100, 1)
            return gpu_metrics
        finally:
            pynvml.nvmlShutdown()
    except Exception:
        return {}


# ---------------------------------------------------------------------------
# Task 4: Health watcher
# ---------------------------------------------------------------------------


async def _health_watcher(
    *,
    config: SidecarConfig,
    queue: asyncio.Queue[dict[str, Any]],
    shutdown_event: asyncio.Event,
) -> None:
    """Poll the main process PID every 5 seconds.

    If the main process is dead, enqueue close(failed) + shutdown so the
    sidecar marks the run as failed and exits cleanly.
    """
    import datetime

    while not shutdown_event.is_set():
        try:
            await asyncio.wait_for(shutdown_event.wait(), timeout=5.0)
            return  # shutdown_event set — normal exit
        except asyncio.TimeoutError:
            pass

        try:
            os.kill(config.main_pid, 0)  # signal 0 = existence check
        except ProcessLookupError:
            # Main process is dead
            at = datetime.datetime.now(datetime.timezone.utc).isoformat()
            await queue.put({"type": "close", "status": "failed", "at": at})
            await queue.put({"type": "shutdown"})
            return
        except PermissionError:
            # Process exists but we don't have permission — treat as alive
            pass


# ---------------------------------------------------------------------------
# Task 5: Command poller
# ---------------------------------------------------------------------------


async def _command_poller(
    *,
    config: SidecarConfig,
    shutdown_event: asyncio.Event,
    state: dict[str, Any],
) -> None:
    """Poll GET /runs/{id}/commands every command_poll_interval seconds.

    When a 'stop' command is received, writes a sentinel file that
    run.should_stop() checks from the main process.
    """
    async with httpx.AsyncClient(timeout=10.0) as client:
        while not shutdown_event.is_set():
            try:
                await asyncio.wait_for(
                    shutdown_event.wait(),
                    timeout=config.command_poll_interval,
                )
                return  # shutdown_event set
            except asyncio.TimeoutError:
                pass

            run_id = state.get("run_id")
            if not run_id:
                continue

            try:
                url = f"{config.api_url}/runs/{run_id}/commands"
                headers = _auth_headers(config)
                response = await client.get(url, headers=headers)
                if response.status_code == 200:
                    data = response.json()
                    for cmd in data.get("commands", []):
                        if cmd.get("type") == "stop":
                            sentinel = stop_sentinel_path(client_uuid=config.client_uuid)
                            try:
                                with open(sentinel, "w") as f:
                                    f.write("stop")
                            except OSError:
                                pass
            except Exception:
                pass
