#!/usr/bin/env python3

# Base code is Copyright (c) 2017 Dennis Mellican
# Original licence bellow
#
# =============================
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# =============================
# Sections removed by blog.mqbx.nl were not required for our needs.
# Modification added to print output to terminal with --print switch.
# Modification also added to do some basic calculations for import/export and pv active power + estimated losses
# Added extra switches for one-shot and to reduce unnecessary requests.

from SungrowModbusTcpClient import SungrowModbusTcpClient
from pymodbus.payload import BinaryPayloadDecoder
from pymodbus.client.sync import ModbusTcpClient
from pymodbus.constants import Endian
from importlib import import_module
from threading import Thread

import datetime
import requests
import argparse
import logging
import dweepy
import json
import time
import sys
import re


MIN_SIGNED   = -2147483648
MAX_UNSIGNED =  4294967295

requests.packages.urllib3.disable_warnings()

# Load in the config module
parser = argparse.ArgumentParser()
parser.add_argument("-c", "--config", default="config", help="Python module to load as our config")
parser.add_argument("-v", "--verbose", action="count", default=0, help="Level of verbosity 0=ERROR 1=INFO 2=DEBUG")
parser.add_argument("--one-shot", action="store_true", help="Run solariot just once then exit, useful for cron based execution")
parser.add_argument("--print", action="store_true", help="Print one-shot to terminal stdout")

args = parser.parse_args()

if args.verbose == 0:
    log_level = logging.WARNING
elif args.verbose == 1:
    log_level = logging.INFO
else:
    log_level = logging.DEBUG

logging.basicConfig(level=log_level)

try:
    config = import_module(args.config)
    logging.info(f"Loaded config {config.model}")
except ModuleNotFoundError:
    parser.error(f"Unable to locate {args.config}.py")

# SMA datatypes and their register lengths
# S = Signed Number, U = Unsigned Number, STR = String
sma_moddatatype = {
  "S16": 1,
  "U16": 1,
  "S32": 2,
  "U32": 2,
  "U64": 4,
  "STR16": 8,
  "STR32": 16,
}

# Load the modbus register map for the inverter
modmap_file = f"modbus-{config.model}"

try:
    modmap = import_module(modmap_file)
except ModuleNotFoundError:
    logging.error(f"Unable to locate {modmap_file}.py")
    sys.exit(1)

# This will try the Sungrow client otherwise will default to the standard library.
client_payload = {
    "host": config.inverter_ip,
    "timeout": config.timeout,
    "RetryOnEmpty": True,
    "retries": 3,
    "port": config.inverter_port,
}

if "sungrow-" in config.model:
    logging.info("Creating SungrowModbusTcpClient")
    client = SungrowModbusTcpClient.SungrowModbusTcpClient(**client_payload)
else:
    logging.info("Creating ModbusTcpClient")
    client = ModbusTcpClient(**client_payload)

logging.info("Connecting")
client.connect()
client.close()
logging.info("Connected")

# Inverter Scanning
inverter = {}
bus = json.loads(modmap.scan)

def load_registers(register_type, start, count=100):
    try:
        if register_type == "read":
            rr = client.read_input_registers(
                int(start),
                count=count,
                unit=config.slave,
            )
        elif register_type == "holding":
            rr = client.read_holding_registers(
                int(start),
                count=count,
                unit=config.slave,
            )
        else:
            raise RuntimeError(f"Unsupported register type: {type}")
    except Exception as err:
        logging.warning("No data. Try increasing the timeout or scan interval.")
        return False

    if rr.isError():
        logging.warning("Modbus connection failed")
        return False

    if len(rr.registers) != count:
        logging.warning(f"Mismatched number of registers read {len(rr.registers)} != {count}")
        return

    divide_regex = re.compile(r"(?P<register_name>[a-zA-Z0-9_]+)_(?P<divide_by>[0-9\.]+)$")

    for num in range(0, count):
        run = int(start) + num + 1

        if register_type == "read" and modmap.read_register.get(str(run)):
            # Check if the modbus map has an '_10' or '_100' etc on the end
            # If so, we divide by that and drop it from the name
            should_divide = divide_regex.match(modmap.read_register.get(str(run)))

            if should_divide:
                inverter[should_divide["register_name"]] = float(rr.registers[num]) / float(should_divide["divide_by"])
            else:
                inverter[modmap.read_register.get(str(run))] = rr.registers[num]
        elif register_type == "holding" and modmap.holding_register.get(str(run)):
            inverter[modmap.holding_register.get(str(run))] = rr.registers[num]

    return True

# Function for polling data from the target and triggering writing to log file if set
def load_sma_register(registers):
    # Request each register from datasets, omit first row which contains only column headers
    for thisrow in registers:
        name = thisrow[0]
        startPos = thisrow[1]
        type = thisrow[2]
        format = thisrow[3]

        # If the connection is somehow not possible (e.g. target not responding)
        # show a error message instead of excepting and stopping
        try:
            received = client.read_input_registers(
                address=startPos,
                count=sma_moddatatype[type],
                unit=config.slave
            )
        except Exception:
            thisdate = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            logging.error(f"{thisdate}: Connection not possible, check settings or connection")
            return

        message = BinaryPayloadDecoder.fromRegisters(received.registers, endian=Endian.Big)

        # Provide the correct result depending on the defined datatype
        if type == "S32":
            interpreted = message.decode_32bit_int()
        elif type == "U32":
            interpreted = message.decode_32bit_uint()
        elif type == "U64":
            interpreted = message.decode_64bit_uint()
        elif type == "STR16":
            interpreted = message.decode_string(16)
        elif type == "STR32":
            interpreted = message.decode_string(32)
        elif type == "S16":
            interpreted = message.decode_16bit_int()
        elif type == "U16":
            interpreted = message.decode_16bit_uint()
        else:
            # If no data type is defined do raw interpretation of the delivered data
            interpreted = message.decode_16bit_uint()

        # Check for "None" data before doing anything else
        if ((interpreted == MIN_SIGNED) or (interpreted == MAX_UNSIGNED)):
            displaydata = None
        else:
            # Put the data with correct formatting into the data table
            if format == "FIX3":
                displaydata = float(interpreted) / 1000
            elif format == "FIX2":
                displaydata = float(interpreted) / 100
            elif format == "FIX1":
                displaydata = float(interpreted) / 10
            else:
                displaydata = interpreted

        logging.debug(f"************** {name} = {displaydata}")
        inverter[name] = displaydata

    # Add timestamp
    inverter["00000 - Timestamp"] = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

def save_json(inverter):
    try:
        f = open(config.json_file,'w')
        f.write(json.dumps(inverter))
        f.close()
    except Exception as err:
        logging.error("Error writing telemetry to file: %s" % err)
        return
    logging.info("Inverter telemetry written to %s file." % config.json_file)

# Core monitoring loop
def scrape_inverter():
    """ Connect to the inverter and scrape the metrics """
    client.connect()

    if "sungrow-" in config.model:
        for i in bus["read"]:
            if not load_registers("read", i["start"], int(i["range"])):
                return False

        for i in bus["holding"]:
            if not load_registers("holding", i["start"], int(i["range"])):
                return False

        # Sungrow inverter specifics:
        # Work out if the grid power is being imported or exported
        if config.model == "sungrow-sh5k":
            try:
                if inverter["grid_import_or_export"] == 65535:
                    export_power = (65535 - inverter["export_power"]) * -1
                    inverter["export_power"] = export_power
            except Exception:
                pass

        try:
            inverter["timestamp"] = "%s/%s/%s %s:%02d:%02d" % (
                inverter["day"],
                inverter["month"],
                inverter["year"],
                inverter["hour"],
                inverter["minute"],
                inverter["second"],
            )
        except Exception:
            pass
    elif "sma-" in config.model:
        load_sma_register(modmap.sma_registers)
    else:
        raise RuntimeError(f"Unsupported inverter model detected: {config.model}")

    client.close()

    logging.info(inverter)
    return True

while True:
    # Scrape the inverter
    success = scrape_inverter()

    if not success:
      if args.one_shot:
          logging.warning("Failed to scrape Inverter, exiting due to --one-shot")
          break
      else:
          logging.warning("Failed to scrape inverter, sleeping until next scan")
          time.sleep(config.scan_interval)
          continue

    if hasattr(config, "json_file"):
        t = Thread(target=save_json, args=(inverter,))
        t.start()

    if args.one_shot:
        logging.info("Exiting due to --one-shot")
        if args.print:
            #Mod Calculated Items Extras
            inverter['export_to_grid'] = inverter['total_active_power'] - inverter['house_loads']
            inverter['import_from_grid'] = inverter['house_loads'] - inverter['total_active_power']
            inverter['export_import'] = inverter['import_from_grid']

            inverter['pv1_active_power'] = inverter['pv1_current'] * inverter['pv1_voltage']
            inverter['pv2_active_power'] = inverter['pv2_current'] * inverter['pv2_voltage']
            inverter['total_pv_active_power'] = inverter['pv1_active_power'] + inverter['pv2_active_power']
            try:
                inverter['pv_power_losses'] = 100 * ((inverter['total_pv_active_power'] - inverter['total_active_power']) / inverter['total_pv_active_power'])
            except ZeroDivisionError:
                inverter['pv_power_losses'] = 0.0

            print(json.dumps(inverter))
        break

    # Sleep until the next scan
    #time.sleep(config.scan_interval)
    break