import time
from threading import Thread
from pymodbus.server import StartSerialServer
from pymodbus.datastore import (
    ModbusServerContext,
    ModbusSequentialDataBlock,
    ModbusSlaveContext
)
from librpiplc import rpiplc

# Define the slave address
SLAVE_ADDR = 1

VERSION = "RPIPLC_V4"
MODEL = "RPIPLC_21"

# Define the configuration for the PLC models (Define your PLC pinouts here)
PLC_CONFIG = {
    'RPIPLC_21': {'Q': 5, 'A': 3, 'AI': 6, 'I': 7}
}

def init():
    # Initialize the RPI PLC
    rpiplc.init(VERSION, MODEL)
    
    # Set pin modes for digital and analog outputs, as well as inputs
    for i in range(PLC_CONFIG[MODEL]['Q']):
        pin_name = "Q0.{}".format(i)
        rpiplc.pin_mode(pin_name, rpiplc.OUTPUT)

    for i in range(PLC_CONFIG[MODEL]['A']):
        pin_name = "A0.{}".format(i + 5)
        rpiplc.pin_mode(pin_name, rpiplc.OUTPUT)

    for i in range(PLC_CONFIG[MODEL]['I']):
        pin_name = "I0.{}".format(i)
        rpiplc.pin_mode(pin_name, rpiplc.INPUT)

    for i in range(PLC_CONFIG[MODEL]['AI']):
        pin_name = "I0.{}".format(i + 7)
        rpiplc.pin_mode(pin_name, rpiplc.INPUT)

def update_plc_datastore(store):
    # Update digital outputs
    coils = store.getValues(1, 0, PLC_CONFIG[MODEL]['Q'])
    for i in range(PLC_CONFIG[MODEL]['Q']):
        pin_name = "Q0.{}".format(i)
        rpiplc.digital_write(pin_name, coils[i])

    # Update analog outputs
    registers = store.getValues(3, 0, PLC_CONFIG[MODEL]['A'])
    for i in range(PLC_CONFIG[MODEL]['A']):
        pin_name = "A0.{}".format(i + 5)
        rpiplc.analog_write(pin_name, registers[i])

    # Update digital inputs
    di_values = []
    for i in range(PLC_CONFIG[MODEL]['I']):
        pin_name = "I0.{}".format(i)
        value = rpiplc.digital_read(pin_name)
        di_values.append(value)

    store.setValues(2, 0, di_values)

    # Update analog inputs
    ai_values = []
    for i in range(PLC_CONFIG[MODEL]['AI']):
        pin_name = "I0.{}".format(i + 7)
        value = rpiplc.analog_read(pin_name)
        ai_values.append(value)

    store.setValues(4, 0, ai_values)
    
def start_modbus_server(context):
    StartSerialServer(context=context,
                      port='/dev/ttySC0',
                      baudrate=38400,
                      bytesize=8,
                      parity='E',
                      stopbits=1,
                      )

def run_server():
    # Define data blocks with 100 registers
    di = ModbusSequentialDataBlock(0, [0] * 100)
    co = ModbusSequentialDataBlock(0, [0] * 100)
    hr = ModbusSequentialDataBlock(0, [0] * 100)
    ir = ModbusSequentialDataBlock(0, [0] * 100)

    # Create a Modbus slave context
    store = ModbusSlaveContext(di=di, co=co, hr=hr, ir=ir)
    context = ModbusServerContext(slaves={SLAVE_ADDR: store}, single=False)

    # Start the Modbus RTU server in a separate thread
    server_thread = Thread(target=start_modbus_server, args=(context,), daemon=True)
    server_thread.start()

    # Main loop to update PLC data based on Modbus register values
    try:
        while True:
            update_plc_datastore(store)
            time.sleep(0.1)
    except KeyboardInterrupt:
        print("Server stopped")

if __name__ == '__main__':
    init()
    run_server()