#!/usr/bin/env python3
"""
Modbus Slave Simulator - A configurable Modbus TCP/RTU slave server for testing.
"""

import argparse
import logging
import signal
import sys
from dataclasses import dataclass
from typing import Optional

from pymodbus.datastore import (
    ModbusSequentialDataBlock,
    ModbusServerContext,
    ModbusSlaveContext,
)
from pymodbus.device import ModbusDeviceIdentification
from pymodbus.server import StartTcpServer, StartTlsServer, StartSerialServer
import pymodbus

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger("modbus-slave")


@dataclass
class ServerConfig:
    """Configuration for the Modbus slave server."""

    host: str = "127.0.0.1"
    port: int = 5020
    slave_id: int = 1
    type: str = "tcp"

    # RTU/Serial specific
    serial_port: str = "/dev/ttyUSB0"
    serial_baudrate: int = 9600
    serial_bytesize: int = 8
    serial_parity: str = "N"
    serial_stopbits: int = 1

    # Data store sizes
    coils_size: int = 100
    discrete_inputs_size: int = 100
    holding_registers_size: int = 100
    input_registers_size: int = 100

    # Initial values
    coils_default: bool = False
    discrete_inputs_default: bool = False
    holding_registers_default: int = 0
    input_registers_default: int = 0

    # TLS
    certfile: Optional[str] = None
    keyfile: Optional[str] = None

    # Identity
    vendor_name: str = "ModbusSimulator"
    product_code: str = "MBS"
    model_name: str = "Simulator"
    version: str = "1.0"


def create_data_store(config: ServerConfig) -> ModbusSlaveContext:
    """Create and configure the Modbus data store."""
    def make_block(start: int, size: int, default) -> ModbusSequentialDataBlock:
        return ModbusSequentialDataBlock(start, [default] * size)

    return ModbusSlaveContext(
        co=make_block(0, config.coils_size, config.coils_default),
        di=make_block(0, config.discrete_inputs_size, config.discrete_inputs_default),
        hr=make_block(0, config.holding_registers_size, config.holding_registers_default),
        ir=make_block(0, config.input_registers_size, config.input_registers_default),
    )


def create_server_context(slave_id: int, store: ModbusSlaveContext) -> ModbusServerContext:
    """Create the server context with single slave."""
    return ModbusServerContext(slaves={slave_id: store}, single=False)


def create_identity(config: ServerConfig) -> ModbusDeviceIdentification:
    """Create Modbus device identification."""
    return ModbusDeviceIdentification(
        info_name={
            "VendorName": config.vendor_name,
            "ProductCode": config.product_code,
            "ModelName": config.model_name,
            "MajorMinorRevision": config.version,
        }
    )


def create_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Modbus Slave Simulator - A configurable Modbus server for testing.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s                           # TCP on 127.0.0.1:5020
  %(prog)s --port 5020 --holding-registers 1000
  %(prog)s --type rtu --serial-port /dev/ttyUSB0 --slave-id 2
  %(prog)s --coils-size 50 --holding-registers-default 100
        """,
    )

    parser.add_argument("--type", "-t", choices=["tcp", "rtu", "tls"], default="tcp",
                        help="Server type (default: tcp)")
    parser.add_argument("--host", default="127.0.0.1", help="TCP bind host (default: 127.0.0.1)")
    parser.add_argument("--port", "-p", type=int, default=5020, help="TCP port (default: 5020)")
    parser.add_argument("--slave-id", "-s", type=int, default=1, help="Modbus slave ID (default: 1)")

    ser = parser.add_argument_group("Serial (RTU)")
    ser.add_argument("--serial-port", default="/dev/ttyUSB0", help="Serial port")
    ser.add_argument("--baudrate", type=int, default=9600, help="Baudrate (default: 9600)")
    ser.add_argument("--bytesize", type=int, choices=[5, 6, 7, 8], default=8)
    ser.add_argument("--parity", choices=["N", "E", "O"], default="N")
    ser.add_argument("--stopbits", type=int, choices=[1, 2], default=1)

    tls = parser.add_argument_group("TLS")
    tls.add_argument("--certfile", help="Certificate file")
    tls.add_argument("--keyfile", help="Key file")

    size = parser.add_argument_group("Data store sizes")
    size.add_argument("--coils-size", type=int, default=100)
    size.add_argument("--discrete-inputs-size", type=int, default=100)
    size.add_argument("--holding-registers-size", type=int, default=100)
    size.add_argument("--input-registers-size", type=int, default=100)

    init = parser.add_argument_group("Initial values")
    init.add_argument("--coils-default", type=lambda x: x.lower() == "true", default=False)
    init.add_argument("--discrete-inputs-default", type=lambda x: x.lower() == "true", default=False)
    init.add_argument("--holding-registers-default", type=int, default=0)
    init.add_argument("--input-registers-default", type=int, default=0)

    ident = parser.add_argument_group("Device identity")
    ident.add_argument("--vendor", default="ModbusSimulator")
    ident.add_argument("--product-code", default="MBS")
    ident.add_argument("--model", default="Simulator")
    ident.add_argument("--version", default="1.0")

    parser.add_argument("--verbose", "-v", action="count", default=0, help="Increase verbosity")
    return parser


def get_log_level(verbose: int) -> int:
    levels = [logging.WARNING, logging.INFO, logging.DEBUG]
    return levels[min(verbose, len(levels) - 1)]


def run_server(config: ServerConfig):
    """Create and start the server based on configuration."""
    store = create_data_store(config)
    context = create_server_context(config.slave_id, store)
    identity = create_identity(config)

    logger.info(f"Starting Modbus {config.type.upper()} server")
    logger.info(f"  Address: {config.host if config.type != 'rtu' else config.serial_port}")
    if config.type != "rtu":
        logger.info(f"  Port: {config.port}")
    logger.info(f"  Slave ID: {config.slave_id}")
    logger.info(f"  Data store: coils={config.coils_size}, di={config.discrete_inputs_size}, "
                f"hr={config.holding_registers_size}, ir={config.input_registers_size}")

    if config.type == "tcp":
        StartTcpServer(context=context, identity=identity, address=(config.host, config.port))
    elif config.type == "tls":
        if not config.certfile or not config.keyfile:
            raise ValueError("TLS requires --certfile and --keyfile")
        StartTlsServer(context=context, identity=identity, address=(config.host, config.port),
                       certfile=config.certfile, keyfile=config.keyfile)
    elif config.type == "rtu":
        StartSerialServer(context=context, identity=identity, port=config.serial_port,
                         baudrate=config.serial_baudrate, bytesize=config.serial_bytesize,
                         parity=config.serial_parity, stopbits=config.serial_stopbits)


def main():
    parser = create_parser()
    args = parser.parse_args()
    logger.setLevel(get_log_level(args.verbose))

    config = ServerConfig(
        host=args.host, port=args.port, slave_id=args.slave_id, type=args.type,
        serial_port=args.serial_port, serial_baudrate=args.baudrate,
        serial_bytesize=args.bytesize, serial_parity=args.parity, serial_stopbits=args.stopbits,
        coils_size=args.coils_size, discrete_inputs_size=args.discrete_inputs_size,
        holding_registers_size=args.holding_registers_size, input_registers_size=args.input_registers_size,
        coils_default=args.coils_default, discrete_inputs_default=args.discrete_inputs_default,
        holding_registers_default=args.holding_registers_default, input_registers_default=args.input_registers_default,
        certfile=args.certfile, keyfile=args.keyfile,
        vendor_name=args.vendor, product_code=args.product_code,
        model_name=args.model, version=args.version,
    )

    try:
        logger.info(f"Modbus Slave Simulator v{pymodbus.__version__}")
        run_server(config)
    except KeyboardInterrupt:
        logger.info("Server stopped by user")
    except Exception as e:
        logger.error(f"Server error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()