import asyncio
import os
import matplotlib.pyplot as plt
import sys
from dotenv import load_dotenv
from huawei_solar_tools.devices.battery import Battery
import huawei_solar_tools.registers.battery_register_names as brn

async def main():
    # Load environment variables
    load_dotenv("huawei_conf.env")
    host = os.getenv("SMARTLOGGER_IP")
    port = int(os.getenv("SMARTLOGGER_PORT"))
    battery_ids = [int(i) for i in os.getenv("BATTERY_IDS").split(",")]
    for device_id in battery_ids:
        # Set up individual report file for each battery
        script_dir = os.path.dirname(os.path.abspath(__file__))
        reports_dir = os.path.join(script_dir, "battery_reports")
        os.makedirs(reports_dir, exist_ok=True)
        report_file = os.path.join(reports_dir, f"battery_{device_id}_report.txt")

        with open(report_file, "w") as report:
            sys.stdout = report

            print(f"Connecting to battery {device_id}...")

            # Initialize the Inverter
            battery = Battery(host, port, device_id)

            try:
                # Test connection
                await battery.connect()
                print("Connection successful.")

                # Test connection status
                print(f"Is battery {device_id} connected? {battery.is_connected()}")

                # Test getting device ID
                print(f"Battery ID: {battery.get_device_id()}")

                # Test reading all registers
                results = await battery.read_all_registers()
                print("All Registers:")
                for i in results:
                    print(i)

                # Test reading a specific register (grid frequency)
                freq = await battery.read_raw_register(37700, 10)
                print(f"Grid Frequency: {freq[0] / 100} Hz")

                # Test reading a specific register by name (Serial Number)
                serial_number = await battery.read_register_by_name(brn.ENERGY_STORAGE_SOC)
                print(f"Serial Number: {serial_number}")

                #Test getting the running status of the units
                print(await battery.get_unit_status())
                
                #Test getting the running status of the system
                print(await battery.get_system_status())
                
                #Test getting the working mode
                working_mode = await battery.get_working_mode()
                print(f"Working Mode: {working_mode}")
                
                #Test getting the product mode
                product_mode = await battery.get_product_mode()
                print(f"Product Mode: {product_mode}")

            except Exception as e:
                print(f"Error processing battery {device_id}: {e}")

            finally:
                # Test disconnection
                battery.disconnect()
                print(f"Disconnected from battery {device_id}.")

def plot_and_save_curve(curve, title, x_label, y_label, x_key, y_key, filename):
    """
    Plot and save a curve to the tests/plots directory.
    """
    if not curve:
        print(f"No data to plot for {title}.")
        return

    x = [point[x_key] for point in curve]
    y = [point[y_key] for point in curve]

    plt.figure(figsize=(8, 6))
    plt.plot(x, y, marker='o', linestyle='-', label=title)
    plt.title(title)
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    plt.grid(True)
    plt.legend()
    plt.tight_layout()

    # Save the plot to the tests/plots directory
    plots_dir = os.path.join(os.path.dirname(__file__), "plots")
    os.makedirs(plots_dir, exist_ok=True)
    plot_path = os.path.join(plots_dir, filename)

    plt.savefig(plot_path)
    plt.close()
    print(f"Plot saved to: {plot_path}")



if __name__ == "__main__":
    asyncio.run(main())
