"""Hybrid field-device gateway demo featuring batching and heartbeat monitoring.

This example expands upon the basic TCP forwarding demo by simulating an
industrial PLC (programmable logic controller) that accepts JSON command batches
from OmniLink via the TCP adapter while acknowledgements are reported back to
MQTT subscribers. The scenario highlights how the OmniLink engine can drive
legacy field devices without sacrificing structured feedback for UI clients.

Key capabilities demonstrated
-----------------------------
* Mock PLC listener with structured acknowledgements.
* Command batching with retry semantics for failed operations.
* Periodic heartbeat publishing to monitor link health.
* OmniLink MQTT bridge integration that reports queue depth and PLC responses.

To run the example locally you will need a Mosquitto broker (or any broker that
speaks the MQTT/WebSocket defaults used by :class:`OmniLinkMQTTBridge`). Once the
broker is running::

    python -m omnilink.examples.industrial_gateway_demo

The script starts the OmniLink engine, MQTT bridge, TCP connection, and mock PLC.
Sample commands are published automatically once the MQTT connection is
established. Watch the terminal output for batched payloads, retry attempts, and
structured acknowledgements from the PLC.
"""

from __future__ import annotations

import json
import socket
import socketserver
import threading
import time
import uuid
from collections import deque
from dataclasses import dataclass, field
from typing import Any, Deque, Dict, Iterable, List, Optional, Tuple

from omnilink import OmniLinkEngine, OmniLinkMQTTBridge, give_context

# OmniLink templates recognised by the demo engine.
DEMO_PATTERNS: Iterable[str] = (
    "forward_[payload:any]",
    "set_line_speed_to_[speed:int]",
    "trigger_emergency_stop",
    "ping",
)

# Behavioural tuning knobs for the gateway controller.
BATCH_SIZE = 3
BATCH_INTERVAL_SECONDS = 3.0
HEARTBEAT_INTERVAL_SECONDS = 5.0
MAX_ATTEMPTS = 3


class _ThreadedTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    daemon_threads = True
    allow_reuse_address = True


@dataclass
class _MockPLCState:
    """Track basic state for the simulated PLC."""

    last_heartbeat: float = 0.0
    processed_batches: int = 0
    retry_memory: Dict[str, int] = field(default_factory=dict)

    def record_heartbeat(self) -> Dict[str, Any]:
        self.last_heartbeat = time.time()
        return {
            "status": "heartbeat_ack",
            "received_at": self.last_heartbeat,
        }

    def mark_batch(self) -> None:
        self.processed_batches += 1


class _MockPLCHandler(socketserver.BaseRequestHandler):
    """Simple PLC emulator that expects JSON payloads terminated by newlines."""

    server: "_MockPLCServer"

    def handle(self) -> None:  # type: ignore[override]
        buffer = b""
        peer = self.client_address
        encoding = self.server.encoding
        delimiter = self.server.delimiter
        while True:
            chunk = self.request.recv(4096)
            if not chunk:
                break
            buffer += chunk
            while True:
                if delimiter:
                    marker = buffer.find(delimiter)
                    if marker == -1:
                        break
                    raw, buffer = buffer[:marker], buffer[marker + len(delimiter) :]
                else:
                    raw, buffer = buffer, b""
                if not raw.strip():
                    if not delimiter:
                        break
                    continue
                try:
                    message = json.loads(raw.decode(encoding))
                except json.JSONDecodeError as exc:
                    print(f"[plc] Invalid payload from {peer}: {raw!r} ({exc})")
                    continue
                response = self._handle_message(message)
                response_bytes = json.dumps(response).encode(encoding)
                if delimiter:
                    response_bytes += delimiter
                self.request.sendall(response_bytes)
        print(f"[plc] Connection closed for {peer}")

    def _handle_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
        msg_type = message.get("type")
        if msg_type == "heartbeat":
            ack = self.server.state.record_heartbeat()
            print(
                "[plc] Heartbeat received",
                json.dumps({"age": time.time() - ack["received_at"]}),
            )
            return ack

        if msg_type == "command_batch":
            return self._handle_batch(message)

        print(f"[plc] Unexpected message type: {msg_type!r}")
        return {"status": "error", "error": "unknown_message", "received": message}

    def _handle_batch(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        commands: List[Dict[str, Any]] = []
        raw_commands = payload.get("commands")
        if isinstance(raw_commands, list):
            commands = [c for c in raw_commands if isinstance(c, dict)]

        statuses: List[Dict[str, Any]] = []
        for cmd in commands:
            command_id = str(cmd.get("id"))
            command_text = str(cmd.get("command", ""))
            meta = cmd.get("meta") if isinstance(cmd.get("meta"), dict) else {}
            needs_retry = bool(meta.get("simulate_retry"))
            attempts = self.server.state.retry_memory.get(command_id, 0)

            if needs_retry and attempts == 0:
                # Ask the gateway to retry once to demonstrate the retry flow.
                self.server.state.retry_memory[command_id] = attempts + 1
                statuses.append(
                    {
                        "command_id": command_id,
                        "status": "retry",
                        "details": "Simulated PLC validation failure",
                    }
                )
                print(f"[plc] Requested retry for {command_text!r}")
                continue

            self.server.state.retry_memory.pop(command_id, None)
            statuses.append(
                {
                    "command_id": command_id,
                    "status": "completed",
                    "details": f"Applied: {command_text}",
                }
            )
            print(
                "[plc] Executed command",
                json.dumps(
                    {
                        "command": command_text,
                        "meta": meta,
                        "vars": cmd.get("vars"),
                    }
                ),
            )

        self.server.state.mark_batch()
        ack_payload = {
            "type": "batch_ack",
            "batch_id": payload.get("batch_id"),
            "received_at": time.time(),
            "statuses": statuses,
            "processed_batches": self.server.state.processed_batches,
            "heartbeat_age": (
                time.time() - self.server.state.last_heartbeat
                if self.server.state.last_heartbeat
                else None
            ),
        }
        return ack_payload


class _MockPLCServer(_ThreadedTCPServer):
    """Custom server class to attach encoder/decoder settings."""

    def __init__(self, server_address: Tuple[str, int], handler_cls: Any, encoding: str, delimiter: bytes):
        super().__init__(server_address, handler_cls)
        self.encoding = encoding
        self.delimiter = delimiter
        self.state = _MockPLCState()


def start_mock_plc_server(
    host: str = "localhost",
    port: int = 8766,
    *,
    encoding: str = "utf-8",
    delimiter: str = "\n",
) -> Tuple[_MockPLCServer, threading.Thread]:
    """Launch the mock PLC server in a background thread."""

    delim_bytes = delimiter.encode(encoding) if delimiter else b""
    server = _MockPLCServer((host, port), _MockPLCHandler, encoding, delim_bytes)
    thread = threading.Thread(target=server.serve_forever, name="mock-plc", daemon=True)
    thread.start()
    print(f"[plc] Mock PLC listening on {host}:{port}")
    return server, thread


def stop_mock_plc_server(server: _MockPLCServer) -> None:
    """Stop the running mock PLC server."""

    server.shutdown()
    server.server_close()


@dataclass
class _QueuedCommand:
    """Representation of a queued OmniLink command awaiting dispatch."""

    id: str
    command: str
    template: Optional[str]
    vars: Dict[str, Any]
    meta: Dict[str, Any]
    attempts: int = 0
    enqueued_at: float = field(default_factory=time.time)


@dataclass
class _TCPConnectionConfig:
    """Minimal TCP connection settings formerly provided by OmniLinkTCPAdapter."""

    host: str = "localhost"
    port: int = 8766
    timeout: float = 5.0
    encoding: str = "utf-8"
    delimiter: str = "\n"


class IndustrialGatewayController:
    """Coordinate batching, retries, and heartbeats for a TCP endpoint."""

    def __init__(
        self,
        tcp_config: _TCPConnectionConfig,
        *,
        batch_size: int = BATCH_SIZE,
        flush_interval: float = BATCH_INTERVAL_SECONDS,
        heartbeat_interval: float = HEARTBEAT_INTERVAL_SECONDS,
        max_attempts: int = MAX_ATTEMPTS,
    ) -> None:
        self.tcp_config = tcp_config
        self.batch_size = max(1, batch_size)
        self.flush_interval = max(0.5, flush_interval)
        self.heartbeat_interval = max(1.0, heartbeat_interval)
        self.max_attempts = max(1, max_attempts)

        self._pending: Deque[_QueuedCommand] = deque()
        self._lock = threading.Lock()
        self._batch_counter = 0

        self._stop_event = threading.Event()
        self._flush_event = threading.Event()
        self._flush_thread: Optional[threading.Thread] = None
        self._heartbeat_thread: Optional[threading.Thread] = None

    # ------------------------------------------------------------------
    # Lifecycle management
    # ------------------------------------------------------------------
    def start(self) -> None:
        if self._flush_thread and self._flush_thread.is_alive():
            return
        self._stop_event.clear()
        self._flush_thread = threading.Thread(target=self._flush_loop, name="gateway-flush", daemon=True)
        self._heartbeat_thread = threading.Thread(
            target=self._heartbeat_loop,
            name="gateway-heartbeat",
            daemon=True,
        )
        self._flush_thread.start()
        self._heartbeat_thread.start()

    def stop(self) -> None:
        self._stop_event.set()
        self._flush_event.set()
        if self._flush_thread:
            self._flush_thread.join(timeout=2.0)
        if self._heartbeat_thread:
            self._heartbeat_thread.join(timeout=2.0)

    # ------------------------------------------------------------------
    # Event handling
    # ------------------------------------------------------------------
    def handle_event(self, event: Dict[str, Any]) -> Dict[str, Any]:
        command = event.get("command")
        if not command:
            return {"ack": False, "error": "missing_command"}

        queued = _QueuedCommand(
            id=uuid.uuid4().hex,
            command=str(command),
            template=event.get("template"),
            vars=event.get("vars") if isinstance(event.get("vars"), dict) else {},
            meta=event.get("meta") if isinstance(event.get("meta"), dict) else {},
        )
        with self._lock:
            self._pending.append(queued)
            pending_count = len(self._pending)
        self._flush_event.set()
        return {
            "ack": True,
            "queued": True,
            "command_id": queued.id,
            "pending_commands": pending_count,
        }

    # ------------------------------------------------------------------
    # Background loops
    # ------------------------------------------------------------------
    def _flush_loop(self) -> None:
        while not self._stop_event.is_set():
            triggered = self._flush_event.wait(self.flush_interval)
            self._flush_event.clear()
            batch = self._collect_batch(force=triggered)
            if batch:
                try:
                    self._dispatch_batch(batch)
                except Exception as exc:  # pragma: no cover - best effort logging
                    print(f"[gateway] Failed to dispatch batch: {exc}")
            if not triggered:
                # Time-based flush; there might still be commands waiting for interval expiry.
                continue
        # Final drain when stopping
        batch = self._collect_batch(force=True)
        if batch:
            try:
                self._dispatch_batch(batch)
            except Exception as exc:
                print(f"[gateway] Failed to dispatch final batch: {exc}")

    def _heartbeat_loop(self) -> None:
        while not self._stop_event.wait(self.heartbeat_interval):
            try:
                payload = {
                    "type": "heartbeat",
                    "sent_at": time.time(),
                }
                response = self._send_payload(payload)
                print("[gateway] Heartbeat acknowledgement", json.dumps(response))
                try:
                    give_context(json.dumps({"gateway": "heartbeat", "response": response}))
                except Exception:
                    # Context publishing is optional; ignore if the MQTT bridge is absent.
                    pass
            except Exception as exc:
                print(f"[gateway] Heartbeat failed: {exc}")

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------
    def _collect_batch(self, *, force: bool) -> List[_QueuedCommand]:
        now = time.time()
        with self._lock:
            if not self._pending:
                return []
            if force or len(self._pending) >= self.batch_size:
                count = min(len(self._pending), self.batch_size)
            else:
                oldest = self._pending[0]
                if now - oldest.enqueued_at < self.flush_interval:
                    return []
                count = min(len(self._pending), self.batch_size)
            batch = [self._pending.popleft() for _ in range(count)]
        return batch

    def _dispatch_batch(self, batch: List[_QueuedCommand]) -> None:
        if not batch:
            return
        self._batch_counter += 1
        batch_id = f"batch-{self._batch_counter}"

        payload = {
            "type": "command_batch",
            "batch_id": batch_id,
            "commands": [
                {
                    "id": cmd.id,
                    "command": cmd.command,
                    "template": cmd.template,
                    "vars": cmd.vars,
                    "meta": cmd.meta,
                    "attempt": cmd.attempts + 1,
                }
                for cmd in batch
            ],
            "sent_at": time.time(),
        }

        print("[gateway] Dispatching batch", json.dumps(payload, default=str))
        response = self._send_payload(payload)
        print("[gateway] PLC acknowledgement", json.dumps(response, default=str))
        self._handle_ack(batch, response)

    def _handle_ack(self, batch: List[_QueuedCommand], ack: Dict[str, Any]) -> None:
        statuses = ack.get("statuses") if isinstance(ack, dict) else None
        if not isinstance(statuses, list):
            return

        remaining: List[_QueuedCommand] = []
        lookup = {cmd.id: cmd for cmd in batch}
        for status in statuses:
            if not isinstance(status, dict):
                continue
            command_id = status.get("command_id")
            state = status.get("status")
            cmd = lookup.get(command_id)
            if not cmd:
                continue
            if state == "completed":
                continue
            cmd.attempts += 1
            if cmd.attempts >= self.max_attempts:
                print(
                    "[gateway] Command permanently failed",
                    json.dumps({
                        "command_id": cmd.id,
                        "status": state,
                        "details": status.get("details"),
                    }),
                )
                continue
            remaining.append(cmd)
            print(
                "[gateway] Retrying command",
                json.dumps(
                    {
                        "command_id": cmd.id,
                        "attempt": cmd.attempts + 1,
                        "details": status.get("details"),
                    }
                ),
            )

        if remaining:
            with self._lock:
                # Requeue at the front so retries happen promptly.
                for cmd in reversed(remaining):
                    cmd.enqueued_at = time.time()
                    self._pending.appendleft(cmd)
            self._flush_event.set()

        # Push the acknowledgement status to the UI via the context channel.
        summary = {
            "gateway": "industrial",
            "batch_id": ack.get("batch_id"),
            "processed_batches": ack.get("processed_batches"),
            "statuses": statuses,
        }
        try:
            give_context(json.dumps(summary))
        except Exception:
            pass

    def _send_payload(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        cfg = self.tcp_config
        raw = json.dumps(payload).encode(cfg.encoding)
        delimiter = cfg.delimiter.encode(cfg.encoding) if cfg.delimiter else None
        if delimiter and not raw.endswith(delimiter):
            raw += delimiter

        with socket.create_connection((cfg.host, cfg.port), timeout=cfg.timeout) as sock:
            sock.sendall(raw)
            sock.shutdown(socket.SHUT_WR)
            response = b""
            sock.settimeout(cfg.timeout)
            while True:
                chunk = sock.recv(4096)
                if not chunk:
                    break
                response += chunk
                if delimiter and response.endswith(delimiter):
                    break

        if delimiter and response.endswith(delimiter):
            response = response[: -len(delimiter)]
        text = response.decode(cfg.encoding).strip()
        if not text:
            return {}
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            return {"raw": text}


def build_demo_engine() -> Tuple[OmniLinkEngine, _TCPConnectionConfig, IndustrialGatewayController]:
    """Construct the OmniLink engine and supporting controller for the demo."""

    engine = OmniLinkEngine(DEMO_PATTERNS)
    tcp_config = _TCPConnectionConfig()
    controller = IndustrialGatewayController(tcp_config)

    def _handle(evt: Dict[str, Any]) -> Dict[str, Any]:
        return controller.handle_event(evt)

    engine.on(lambda evt: True, _handle)
    return engine, tcp_config, controller


def publish_sample_commands(bridge: OmniLinkMQTTBridge) -> None:
    """Publish a handful of commands once the MQTT bridge connects."""

    deadline = time.time() + 15
    while not bridge.client.is_connected():
        if time.time() > deadline:
            raise TimeoutError("MQTT client failed to connect within 15 seconds")
        time.sleep(0.1)

    samples = [
        {
            "command": "set line speed to 42",
            "meta": {"source": "demo", "id": "cmd-001"},
        },
        {
            "command": "forward calibrate station a",
            "meta": {"source": "demo", "id": "cmd-002"},
        },
        {
            "command": "forward staging queue retry",
            "meta": {"source": "demo", "id": "cmd-003", "simulate_retry": True},
        },
        {
            "command": "trigger emergency stop",
            "meta": {"source": "demo", "id": "cmd-004"},
        },
    ]

    for payload in samples:
        message = json.dumps(payload)
        result = bridge.client.publish(bridge.command_topic, message)
        result.wait_for_publish()
        print(
            "[demo] Published command",
            json.dumps({"topic": bridge.command_topic, "payload": payload}),
        )
        time.sleep(0.5)


def run_demo() -> None:
    """Run the industrial gateway demo with batching and heartbeat support."""

    engine, tcp_config, controller = build_demo_engine()
    plc_server, _ = start_mock_plc_server(
        tcp_config.host,
        tcp_config.port,
        encoding=tcp_config.encoding,
        delimiter=tcp_config.delimiter,
    )

    bridge = OmniLinkMQTTBridge(engine)
    bridge.start()
    controller.start()

    print(
        "[demo] Industrial gateway running. Watch for batch dispatches, retries,",
        "and PLC acknowledgements. Press Ctrl+C to stop.",
    )

    try:
        publish_sample_commands(bridge)
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n[demo] Stopping demo…")
    finally:
        controller.stop()
        bridge.client.loop_stop()
        bridge.client.disconnect()
        stop_mock_plc_server(plc_server)


if __name__ == "__main__":
    run_demo()
