import os
import asyncio
from datetime import datetime, timezone
from dotenv import load_dotenv, find_dotenv
from clickhouse_connect import get_client
from influxdb_client import InfluxDBClient, Point, WritePrecision
from influxdb_client.client.write_api import SYNCHRONOUS
from huawei_solar_tools.devices.inverter import Inverter
from huawei_solar_tools.devices.battery import Battery
from huawei_solar_tools.devices.power_meter import PowerMeter
from huawei_solar_tools.devices.emi import EMI

async def main():
    
    clickhouse_client = None
    influx_client = None
    data_clickhouse = []
    points_influx = []
    
    # Create Clikchouse client
    print("Creating Clickhouse client")
    try:
        clickhouse_client = get_client(
            host=clickhouse_host,
            port=clickhouse_port,
            user=clickhouse_user,
            password=clickhouse_password
        )
    except Exception as e:
        print(f"Error creating ClickHouse client: {e}")

    # Create InfluxDB client
    try:
        influx_client = InfluxDBClient(
            url=influxdb_url_ua,
            token=influxdb_token_ua,
            org=influxdb_org_ua
        )
        influx_write_api = influx_client.write_api(write_options=SYNCHRONOUS)
    except Exception as e:
        print(f"Error creating Influx client: {e}")
    
    # Beginning reading data from devices
    for id in inverter_ids:
        try:
            # Initialize the Inverter
            inverter = Inverter(host, port, id)
            await inverter.connect()

            # Fetch all register data
            results = await inverter.read_all_registers()
            
            # Prepare data for ClickHouse
            data_to_send_inverter = [
                (stamptime, "inverter", id, result[0], result[1])
                for result in results if isinstance(result[1], float)
            ]
            data_clickhouse.extend(data_to_send_inverter)
            
            # Prepare points for InfluxDB
            influx_points_inverter = [
                Point("huawei_solar_system")
                .tag("inverter_id", id)
                .field(result[0], result[1])
                .time(stamptime, WritePrecision.S)
                for result in results if isinstance(result[1], float)
            ]
            points_influx.extend(influx_points_inverter)
            inverter.disconnect()
        except Exception as e:
            print(f"Error processing data from inverter {id}: {e}")
            continue


    for id in battery_ids:
        try:
            # Initialize the Inverter
            battery = Battery(host, port, id)
            await battery.connect()

            # Fetch all register data
            results = await battery.read_all_registers()
            
            # Prepare data for ClickHouse
            data_to_send_battery = [
                (stamptime, "battery", id, result[0], result[1])
                for result in results if isinstance(result[1], float)
            ]
            data_clickhouse.extend(data_to_send_battery)
            
            # Prepare points for InfluxDB
            influx_points_battery = [
                Point("huawei_solar_system")
                .tag("battery_id", id)
                .field(result[0], result[1])
                .time(stamptime, WritePrecision.S)
                for result in results if isinstance(result[1], float)
            ]
            points_influx.extend(influx_points_battery)
            battery.disconnect()
            
        except Exception as e:
            print(f"Error processing data from battery {id}: {e}")
            continue

    for id in power_meter_ids:
        try:
            # Initialize the Power Meter
            power_meter = PowerMeter(host, port, id)
            await power_meter.connect()

            # Fetch all register data
            results = await power_meter.read_all_registers()
            
            # Prepare data for ClickHouse
            data_to_send_power_meter = [
                (stamptime, "power_meter", id, result[0], result[1])
                for result in results if isinstance(result[1], float)
            ]
            data_clickhouse.extend(data_to_send_power_meter)
            
            # Prepare points for InfluxDB
            influx_points_power_meter = [
                Point("huawei_solar_system")
                .tag("power_meter_id", id)
                .field(result[0], result[1])
                .time(stamptime, WritePrecision.S)
                for result in results if isinstance(result[1], float)
            ]
            points_influx.extend(influx_points_power_meter)
            power_meter.disconnect()
            
        except Exception as e:
            print(f"Error processing data from battery {id}: {e}")
            continue
                   
    for id in emi_ids:
        try:
            # Initialize the EMI
            emi = EMI(host, port, id)
            await emi.connect()

            # Fetch all register data
            results = await emi.read_all_registers()
            
            # Prepare data for ClickHouse
            data_to_send_emi = [
                (stamptime, "emi", id, result[0], result[1])
                for result in results if isinstance(result[1], float)
            ]
            data_clickhouse.extend(data_to_send_emi)
            
            # Prepare points for InfluxDB
            influx_points_emi = [
                Point("huawei_solar_system")
                .tag("emi_id", id)
                .field(result[0], result[1])
                .time(stamptime, WritePrecision.S)
                for result in results if isinstance(result[1], float)
            ]
            points_influx.extend(influx_points_emi)
            emi.disconnect()
            
        except Exception as e:
            print(f"Error processing data from battery {id}: {e}")
            continue



    # Write to Clickhouse
    clickhouse_client.insert(clickhouse_table, data_clickhouse)
    print(f"Data written to Clickhouse")
    # Write to influx
    influx_write_api.write(bucket=influxdb_bucket_ua, record=points_influx)
    print(f"Data written to InfluxDB")
    
    if influx_client:
        influx_client.close()
    if clickhouse_client:
        clickhouse_client.close()
            
if __name__ == "__main__":
    # Load environment variables
    load_dotenv(find_dotenv("huawei_conf.env"))

    # Environment variables
    host = os.getenv("SMARTLOGGER_IP")
    port = int(os.getenv("SMARTLOGGER_PORT"))
    inverter_ids = [int(i) for i in os.getenv("INVERTER_IDS").split(",")]
    battery_ids = [int(i) for i in os.getenv("BATTERY_IDS").split(",")]
    power_meter_ids = [int(i) for i in os.getenv("POWER_METER_IDS").split(",")]
    emi_ids = [int(i) for i in os.getenv("EMI_IDS").split(",")]

    # InfluxDB
    influxdb_url_ua = os.getenv("INFLUXDB_URL_UA")
    influxdb_token_ua = os.getenv("INFLUXDB_TOKEN_UA")
    influxdb_org_ua = os.getenv("INFLUXDB_ORG_UA")
    influxdb_bucket_ua = os.getenv("INFLUXDB_BUCKET_UA")

    # ClickHouse
    clickhouse_host = os.getenv("HOST_CLICKHOUSE")
    clickhouse_port = int(os.getenv("PORT_CLICKHOUSE"))
    clickhouse_user = os.getenv("USER_CLICKHOUSE")
    clickhouse_password = os.getenv("PASSWORD_CLICKHOUSE")
    clickhouse_timeout = int(os.getenv("TIMEOUT_CLICKHOUSE"))
    clickhouse_table = os.getenv("TABLE_CLICKHOUSE")

    # Current timestamp
    stamptime = datetime.now(timezone.utc).replace(second=0, microsecond=0)

    asyncio.run(main())
