#!/usr/bin/env python3
"""
Modbus Probe Tool for ICS-SimLab Diagnostics.

Probes all monitor targets (HMI→PLC, PLC→Sensor) to diagnose:
- TCP connectivity
- Modbus exceptions (illegal address/function)
- Register type/address mismatches

Usage:
    python3 tools/probe_modbus.py [--config path/to/configuration.json]
    python3 tools/probe_modbus.py --docker  # Run from host via docker exec
"""

import argparse
import json
import socket
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional

# Optional pymodbus import (may not be available on host)
try:
    from pymodbus.client import ModbusTcpClient
    from pymodbus.exceptions import ModbusException
    PYMODBUS_AVAILABLE = True
except ImportError:
    PYMODBUS_AVAILABLE = False


@dataclass
class ProbeTarget:
    """A Modbus read target to probe."""
    source: str          # e.g., "operator_hmi", "plc1"
    monitor_id: str      # e.g., "water_tank_level_reg"
    target_ip: str
    target_port: int
    value_type: str      # coil, discrete_input, holding_register, input_register
    slave_id: int
    address: int
    count: int


@dataclass
class ProbeResult:
    """Result of probing a target."""
    target: ProbeTarget
    tcp_ok: bool
    modbus_ok: bool
    value: Optional[Any] = None
    error: Optional[str] = None


def parse_config(config_path: Path) -> Dict[str, Any]:
    """Load and parse configuration.json."""
    with open(config_path) as f:
        return json.load(f)


def extract_probe_targets(config: Dict[str, Any]) -> List[ProbeTarget]:
    """Extract all monitor targets from config (HMIs and PLCs)."""
    targets = []

    # HMI monitors -> PLCs
    for hmi in config.get("hmis", []):
        hmi_name = hmi.get("name", "unknown_hmi")
        outbound_map = {
            conn["id"]: (conn["ip"], conn["port"])
            for conn in hmi.get("outbound_connections", [])
        }

        for mon in hmi.get("monitors", []):
            conn_id = mon.get("outbound_connection_id")
            if conn_id not in outbound_map:
                continue
            ip, port = outbound_map[conn_id]
            targets.append(ProbeTarget(
                source=hmi_name,
                monitor_id=mon["id"],
                target_ip=ip,
                target_port=port,
                value_type=mon["value_type"],
                slave_id=mon.get("slave_id", 1),
                address=mon["address"],
                count=mon.get("count", 1),
            ))

    # PLC monitors -> Sensors/other PLCs
    for plc in config.get("plcs", []):
        plc_name = plc.get("name", "unknown_plc")
        outbound_map = {
            conn["id"]: (conn["ip"], conn["port"])
            for conn in plc.get("outbound_connections", [])
        }

        for mon in plc.get("monitors", []):
            conn_id = mon.get("outbound_connection_id")
            if conn_id not in outbound_map:
                continue
            ip, port = outbound_map[conn_id]
            targets.append(ProbeTarget(
                source=plc_name,
                monitor_id=mon["id"],
                target_ip=ip,
                target_port=port,
                value_type=mon["value_type"],
                slave_id=mon.get("slave_id", 1),
                address=mon["address"],
                count=mon.get("count", 1),
            ))

    return targets


def check_tcp(ip: str, port: int, timeout: float = 2.0) -> bool:
    """Check if TCP port is reachable."""
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.settimeout(timeout)
            sock.connect((ip, port))
        return True
    except (socket.timeout, ConnectionRefusedError, OSError):
        return False


def probe_modbus(target: ProbeTarget, timeout: float = 3.0) -> ProbeResult:
    """Probe a single Modbus target."""
    # First check TCP connectivity
    tcp_ok = check_tcp(target.target_ip, target.target_port, timeout)
    if not tcp_ok:
        return ProbeResult(
            target=target,
            tcp_ok=False,
            modbus_ok=False,
            error=f"TCP connection refused: {target.target_ip}:{target.target_port}"
        )

    if not PYMODBUS_AVAILABLE:
        return ProbeResult(
            target=target,
            tcp_ok=True,
            modbus_ok=False,
            error="pymodbus not installed (TCP OK)"
        )

    # Try Modbus read
    client = ModbusTcpClient(
        host=target.target_ip,
        port=target.target_port,
        timeout=timeout
    )

    try:
        if not client.connect():
            return ProbeResult(
                target=target,
                tcp_ok=True,
                modbus_ok=False,
                error="Modbus connect() failed"
            )

        # Select appropriate read function
        read_funcs = {
            "coil": client.read_coils,
            "discrete_input": client.read_discrete_inputs,
            "holding_register": client.read_holding_registers,
            "input_register": client.read_input_registers,
        }

        func = read_funcs.get(target.value_type)
        if not func:
            return ProbeResult(
                target=target,
                tcp_ok=True,
                modbus_ok=False,
                error=f"Unknown value_type: {target.value_type}"
            )

        # Perform Modbus read
        # ICS-SimLab containers use pymodbus that accepts positional args
        addr = target.address
        cnt = target.count

        # Simple call - positional args work with ICS-SimLab's pymodbus
        result = func(addr, cnt)

        if result.isError():
            return ProbeResult(
                target=target,
                tcp_ok=True,
                modbus_ok=False,
                error=f"Modbus error: {result}"
            )

        # Extract value
        if target.value_type in ("coil", "discrete_input"):
            value = result.bits[:target.count]
        else:
            value = result.registers[:target.count]

        return ProbeResult(
            target=target,
            tcp_ok=True,
            modbus_ok=True,
            value=value
        )

    except ModbusException as e:
        return ProbeResult(
            target=target,
            tcp_ok=True,
            modbus_ok=False,
            error=f"ModbusException: {e}"
        )
    except Exception as e:
        return ProbeResult(
            target=target,
            tcp_ok=True,
            modbus_ok=False,
            error=f"Exception: {type(e).__name__}: {e}"
        )
    finally:
        client.close()


def format_result(r: ProbeResult) -> str:
    """Format a probe result as a single line."""
    t = r.target
    status = "OK" if r.modbus_ok else "FAIL"
    tcp_status = "TCP_OK" if r.tcp_ok else "TCP_FAIL"

    line = f"[{status}] {t.source} -> {t.target_ip}:{t.target_port}"
    line += f"  {t.value_type}@{t.address} (id={t.monitor_id})"

    if r.modbus_ok:
        line += f"  value={r.value}"
    else:
        line += f"  ({tcp_status}) {r.error}"

    return line


def run_probe(config_path: Path, verbose: bool = False) -> List[ProbeResult]:
    """Run probe on all targets."""
    config = parse_config(config_path)
    targets = extract_probe_targets(config)

    if verbose:
        print(f"Found {len(targets)} probe targets")
        print("=" * 70)

    results = []
    for target in targets:
        result = probe_modbus(target)
        results.append(result)
        if verbose:
            print(format_result(result))

    return results


def generate_report(results: List[ProbeResult]) -> str:
    """Generate a full probe report."""
    lines = []
    lines.append("=" * 70)
    lines.append("MODBUS PROBE REPORT")
    lines.append("=" * 70)
    lines.append("")

    # Summary
    total = len(results)
    tcp_ok = sum(1 for r in results if r.tcp_ok)
    modbus_ok = sum(1 for r in results if r.modbus_ok)

    lines.append(f"Total targets: {total}")
    lines.append(f"TCP reachable: {tcp_ok}/{total}")
    lines.append(f"Modbus OK: {modbus_ok}/{total}")
    lines.append("")

    # Group by source
    by_source: Dict[str, List[ProbeResult]] = {}
    for r in results:
        src = r.target.source
        by_source.setdefault(src, []).append(r)

    for source, source_results in sorted(by_source.items()):
        lines.append(f"--- {source} monitors ---")
        for r in source_results:
            lines.append(format_result(r))
        lines.append("")

    # Diagnosis
    lines.append("=" * 70)
    lines.append("DIAGNOSIS")
    lines.append("=" * 70)

    tcp_fails = [r for r in results if not r.tcp_ok]
    modbus_fails = [r for r in results if r.tcp_ok and not r.modbus_ok]

    if tcp_fails:
        lines.append("")
        lines.append("TCP FAILURES (connection refused/timeout):")
        for r in tcp_fails:
            lines.append(f"  - {r.target.source} -> {r.target.target_ip}:{r.target.target_port}")
        lines.append("  Likely causes: container not running, network isolation, firewall")

    if modbus_fails:
        lines.append("")
        lines.append("MODBUS FAILURES (TCP OK but read failed):")
        for r in modbus_fails:
            lines.append(f"  - {r.target.source}.{r.target.monitor_id}: {r.error}")
        lines.append("  Likely causes: wrong address, wrong register type, device not serving")

    if modbus_ok == total:
        lines.append("")
        lines.append("All Modbus reads successful. Data flow issue is likely in:")
        lines.append("  - HIL not updating physical_values")
        lines.append("  - Sensors not receiving from HIL")
        lines.append("  - Value is legitimately 0")

    # Values that are 0 (potential issue)
    zero_values = [r for r in results if r.modbus_ok and r.value in ([0], [False], 0)]
    if zero_values:
        lines.append("")
        lines.append("REGISTERS WITH VALUE 0 (may indicate HIL not producing data):")
        for r in zero_values:
            lines.append(f"  - {r.target.source}.{r.target.monitor_id} = {r.value}")

    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description="Modbus probe for ICS-SimLab diagnostics")
    parser.add_argument(
        "--config",
        default="outputs/scenario_run/configuration.json",
        help="Path to configuration.json"
    )
    parser.add_argument(
        "--docker",
        action="store_true",
        help="Run probes via docker exec (from host)"
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Print results as they are collected"
    )
    parser.add_argument(
        "--output", "-o",
        help="Write report to file instead of stdout"
    )

    args = parser.parse_args()
    config_path = Path(args.config)

    if not config_path.exists():
        print(f"ERROR: Config not found: {config_path}", file=sys.stderr)
        sys.exit(1)

    if args.docker:
        # TODO: Implement docker exec wrapper
        print("--docker mode not yet implemented", file=sys.stderr)
        sys.exit(1)

    results = run_probe(config_path, verbose=args.verbose)
    report = generate_report(results)

    if args.output:
        Path(args.output).write_text(report)
        print(f"Report written to {args.output}")
    else:
        print(report)

    # Exit with error if any failures
    if not all(r.modbus_ok for r in results):
        sys.exit(1)


if __name__ == "__main__":
    main()