from __future__ import annotations

import json
import sys
import threading
import time

from rich.console import Console
from rich.live import Live

from hytop.core.history import SlidingHistory
from hytop.core.ssh import SSHOptions, collect_python_from_host
from hytop.net.collector import REMOTE_COLLECTOR_PY, collect_local_counters, parse_counter_payload
from hytop.net.models import HostSnapshot, MonitorState, NetKind, NodeCounterSnapshot, RateSample
from hytop.net.render import build_renderable

LOCAL_HOSTS = {"localhost", "127.0.0.1", "::1"}


def _build_remote_python_script(payload: dict[str, object]) -> str:
    return f"INPUT_JSON={json.dumps(payload)!r}\n{REMOTE_COLLECTOR_PY}"


def _collect_remote_counters(
    host: str,
    ssh_timeout: float,
    cmd_timeout: float,
    kind_filter: set[NetKind],
    include: set[str],
    ssh_options: SSHOptions | None = None,
) -> NodeCounterSnapshot:
    payload = {
        "kind_filter": sorted(kind_filter),
        "include": sorted(include),
    }
    raw = collect_python_from_host(
        host=host,
        ssh_timeout=ssh_timeout,
        cmd_timeout=cmd_timeout,
        python_code=_build_remote_python_script(payload),
        ssh_options=ssh_options,
    )
    if raw.error:
        return NodeCounterSnapshot(host=host, counters={}, error=raw.error)

    try:
        counters = parse_counter_payload(raw.stdout)
    except ValueError as exc:
        return NodeCounterSnapshot(host=host, counters={}, error=str(exc))
    return NodeCounterSnapshot(host=host, counters=counters, sample_ts=time.monotonic())


def collect_node(
    host: str,
    ssh_timeout: float,
    cmd_timeout: float,
    kind_filter: set[NetKind],
    include: set[str],
    ssh_options: SSHOptions | None = None,
) -> NodeCounterSnapshot:
    if host in LOCAL_HOSTS:
        counters = collect_local_counters(
            kind_filter=kind_filter,
            include=include,
        )
        return NodeCounterSnapshot(host=host, counters=counters, sample_ts=time.monotonic())
    return _collect_remote_counters(
        host=host,
        ssh_timeout=ssh_timeout,
        cmd_timeout=cmd_timeout,
        kind_filter=kind_filter,
        include=include,
        ssh_options=ssh_options,
    )


def host_collector_loop(
    host: str,
    ssh_timeout: float,
    cmd_timeout: float,
    interval: float,
    kind_filter: set[NetKind],
    include: set[str],
    ssh_options: SSHOptions | None,
    state: dict[str, HostSnapshot],
    state_lock: threading.Lock,
    stop_event: threading.Event,
) -> None:
    while not stop_event.is_set():
        started = time.monotonic()
        result = collect_node(
            host=host,
            ssh_timeout=ssh_timeout,
            cmd_timeout=cmd_timeout,
            kind_filter=kind_filter,
            include=include,
            ssh_options=ssh_options,
        )
        with state_lock:
            snapshot = state[host]
            snapshot.seq += 1
            snapshot.result = result
        sleep_s = max(0.0, interval - (time.monotonic() - started))
        if stop_event.wait(sleep_s):
            break


def init_monitor_state(hosts: list[str], max_window: float) -> MonitorState:
    return MonitorState(
        max_window=max_window,
        histories={},
        monitored_keys=set(),
        latest_counter_by_key={},
        previous_counter_by_key={},
        previous_sample_ts_by_key={},
        errors={},
        host_state={host: HostSnapshot() for host in hosts},
        processed_seq={host: 0 for host in hosts},
        state_lock=threading.Lock(),
        stop_event=threading.Event(),
    )


def start_collectors(
    hosts: list[str],
    ssh_timeout: float,
    cmd_timeout: float,
    interval: float,
    kind_filter: set[NetKind],
    include: set[str],
    ssh_options: SSHOptions | None,
    state: MonitorState,
) -> list[threading.Thread]:
    workers: list[threading.Thread] = []
    for host in hosts:
        worker = threading.Thread(
            target=host_collector_loop,
            args=(
                host,
                ssh_timeout,
                cmd_timeout,
                interval,
                kind_filter,
                include,
                ssh_options,
                state.host_state,
                state.state_lock,
                state.stop_event,
            ),
            daemon=True,
            name=f"net-collector-{host}",
        )
        worker.start()
        workers.append(worker)
    return workers


def drain_pending_nodes(hosts: list[str], state: MonitorState) -> list[NodeCounterSnapshot]:
    nodes: list[NodeCounterSnapshot] = []
    with state.state_lock:
        for host in hosts:
            snapshot = state.host_state[host]
            if snapshot.seq <= state.processed_seq[host]:
                continue
            state.processed_seq[host] = snapshot.seq
            if snapshot.result is not None:
                nodes.append(snapshot.result)
    return nodes


def apply_node_results(
    nodes: list[NodeCounterSnapshot], interval: float, state: MonitorState
) -> None:
    for node in nodes:
        if node.error:
            state.errors[node.host] = node.error
            continue
        state.errors.pop(node.host, None)
        for iface_key, counter in node.counters.items():
            key = (node.host, iface_key)
            state.monitored_keys.add(key)
            state.latest_counter_by_key[key] = counter

            prev = state.previous_counter_by_key.get(key)
            prev_ts = state.previous_sample_ts_by_key.get(key)
            state.previous_counter_by_key[key] = counter
            state.previous_sample_ts_by_key[key] = node.sample_ts
            if prev is None or prev_ts is None:
                continue
            delta_rx = counter.rx_bytes - prev.rx_bytes
            delta_tx = counter.tx_bytes - prev.tx_bytes
            if delta_rx < 0 or delta_tx < 0:
                continue
            # Use actual elapsed time between samples rather than the configured
            # interval, so SSH latency spikes don't inflate the rate.
            actual_delta = node.sample_ts - prev_ts
            if actual_delta <= 0:
                continue

            history = state.histories.get(key)
            if history is None:
                history = SlidingHistory(max_window_s=state.max_window)
                state.histories[key] = history
            history.add(
                RateSample(
                    ts=node.sample_ts,
                    rx_bps=delta_rx / actual_delta,
                    tx_bps=delta_tx / actual_delta,
                )
            )


def run_monitor(
    hosts: list[str],
    kind_filter: set[NetKind],
    include: set[str],
    window: float,
    interval: float,
    timeout: float | None,
    iec: bool = False,
    ssh_options: SSHOptions | None = None,
) -> int:
    if interval <= 0:
        print("argument error: --interval must be > 0", file=sys.stderr)
        return 2
    if interval > window:
        print("argument error: --interval must be <= --window value", file=sys.stderr)
        return 2

    ssh_timeout = min(max(5 * interval, 2.0), 5.0)
    cmd_timeout = min(max(10 * interval, 5.0), 10.0)

    console = Console()
    err_console = Console(stderr=True)
    started = time.monotonic()
    render_interval = min(interval, 0.5)
    state = init_monitor_state(hosts=hosts, max_window=window)

    try:
        with Live(console=console, auto_refresh=False, screen=True) as live:
            workers = start_collectors(
                hosts=hosts,
                ssh_timeout=ssh_timeout,
                cmd_timeout=cmd_timeout,
                interval=interval,
                kind_filter=kind_filter,
                include=include,
                ssh_options=ssh_options,
                state=state,
            )
            try:
                while True:
                    loop_started = time.monotonic()
                    apply_node_results(
                        nodes=drain_pending_nodes(hosts=hosts, state=state),
                        interval=interval,
                        state=state,
                    )
                    live.update(
                        build_renderable(
                            window=window,
                            hosts=hosts,
                            histories=state.histories,
                            monitored_keys=state.monitored_keys,
                            latest_counter_by_key=state.latest_counter_by_key,
                            errors=state.errors,
                            poll_interval=interval,
                            elapsed_since_start=time.monotonic() - started,
                            iec=iec,
                        ),
                        refresh=True,
                    )
                    if timeout is not None and (time.monotonic() - started) >= timeout:
                        return 124
                    time.sleep(max(0.0, render_interval - (time.monotonic() - loop_started)))
            finally:
                state.stop_event.set()
                for worker in workers:
                    worker.join(timeout=min(0.2, interval))
    except KeyboardInterrupt:
        err_console.print("status: interrupted by user", style="yellow")
        return 130
    return 0
