"""Async Modbus TCP client for Victron devices via Cerbo GX."""

from __future__ import annotations

import asyncio
import logging
from types import TracebackType

from pymodbus.client import AsyncModbusTcpClient
from pymodbus.exceptions import ConnectionException, ModbusIOException

from .const import ALL_UNIT_IDS, DEFAULT_PORT, DEFAULT_RETRIES, DEFAULT_TIMEOUT
from .exceptions import (
    VictronConnectionError,
    VictronModbusError,
    VictronNoDevicesError,
    VictronTimeoutError,
)
from .models import (
    DeviceIdentity,
    SolarChargerData,
    SystemData,
    VEBusData,
    VictronData,
)
from .registers.base import DeviceParser
from .registers.solar_charger import SolarChargerParser
from .registers.system import SystemParser
from .registers.vebus import VEBusParser

_LOGGER = logging.getLogger(__name__)

# Map unit IDs to their parsers
_PARSER_MAP: dict[int, DeviceParser] = {
    100: SystemParser(),
    223: SolarChargerParser(),
    227: VEBusParser(),
}


class VictronClient:
    """Async client for reading Victron devices via Modbus TCP."""

    def __init__(
        self,
        host: str,
        *,
        port: int = DEFAULT_PORT,
        timeout: int = DEFAULT_TIMEOUT,
        retries: int = DEFAULT_RETRIES,
    ) -> None:
        self.host = host
        self.port = port
        self._client = AsyncModbusTcpClient(
            host,
            port=port,
            timeout=timeout,
            retries=retries,
        )
        self._lock = asyncio.Lock()
        self._discovered: dict[int, DeviceParser] = {}

    async def __aenter__(self) -> VictronClient:
        await self.connect()
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        await self.close()

    async def connect(self) -> None:
        """Establish Modbus TCP connection."""
        try:
            connected = await self._client.connect()
        except (ConnectionException, OSError) as err:
            raise VictronConnectionError(
                f"Cannot connect to {self.host}:{self.port}"
            ) from err

        if not connected:
            raise VictronConnectionError(
                f"Cannot connect to {self.host}:{self.port}"
            )

    async def close(self) -> None:
        """Close the Modbus TCP connection."""
        self._client.close()

    @property
    def connected(self) -> bool:
        """Return True if the Modbus client is connected."""
        return self._client.connected

    async def discover_devices(self) -> tuple[DeviceIdentity, ...]:
        """Probe known unit IDs and return identities of responding devices."""
        devices: list[DeviceIdentity] = []
        self._discovered.clear()

        for unit_id in ALL_UNIT_IDS:
            parser = _PARSER_MAP.get(unit_id)
            if parser is None:
                continue

            try:
                # Try reading first register to check if device responds
                first_reg = parser.registers[0]
                result = await self._client.read_holding_registers(
                    first_reg.address, count=1, device_id=unit_id
                )
                if result.isError():
                    _LOGGER.debug(
                        "Unit %d (%s) not responding", unit_id, parser.product_name
                    )
                    continue

                self._discovered[unit_id] = parser
                devices.append(
                    DeviceIdentity(
                        unit_id=unit_id,
                        device_type=parser.device_type,
                        product_name=parser.product_name,
                    )
                )
                _LOGGER.debug("Discovered unit %d: %s", unit_id, parser.product_name)

            except (ConnectionException, ModbusIOException, OSError) as err:
                _LOGGER.debug(
                    "Error probing unit %d: %s", unit_id, err
                )
                continue

        return tuple(devices)

    async def validate_connection(self) -> tuple[DeviceIdentity, ...]:
        """Connect, discover devices, raise if none found.

        Used by HA config flow to validate user input.
        """
        if not self.connected:
            await self.connect()

        devices = await self.discover_devices()
        if not devices:
            raise VictronNoDevicesError(
                f"No Victron devices found at {self.host}:{self.port}"
            )
        return devices

    async def get_data(self) -> VictronData:
        """Read all discovered devices and return combined data."""
        async with self._lock:
            return await self._read_all()

    async def _read_all(self) -> VictronData:
        """Read registers from all discovered devices."""
        if not self._discovered:
            await self.discover_devices()
            if not self._discovered:
                raise VictronNoDevicesError(
                    f"No Victron devices found at {self.host}:{self.port}"
                )

        system_data: SystemData | None = None
        solar_data: SolarChargerData | None = None
        vebus_data: VEBusData | None = None
        identities: list[DeviceIdentity] = []

        for unit_id, parser in self._discovered.items():
            try:
                register_map = await self._read_device(unit_id, parser)
                parsed = parser.parse(register_map)

                identity = DeviceIdentity(
                    unit_id=unit_id,
                    device_type=parser.device_type,
                    product_name=parser.product_name,
                )
                identities.append(identity)

                if parser.device_type == "system":
                    system_data = SystemData(**parsed)
                elif parser.device_type == "solar_charger":
                    solar_data = SolarChargerData(**parsed)
                elif parser.device_type == "vebus":
                    vebus_data = VEBusData(**parsed)

            except (ConnectionException, OSError) as err:
                raise VictronConnectionError(
                    f"Lost connection to {self.host}:{self.port}"
                ) from err
            except TimeoutError as err:
                raise VictronTimeoutError(
                    f"Timeout reading unit {unit_id}"
                ) from err
            except ModbusIOException as err:
                _LOGGER.warning("Modbus error reading unit %d: %s", unit_id, err)
                continue

        return VictronData(
            system=system_data,
            solar_charger=solar_data,
            vebus=vebus_data,
            discovered_devices=tuple(identities),
        )

    async def _read_device(
        self, unit_id: int, parser: DeviceParser
    ) -> dict[int, int]:
        """Read all register ranges for a device, return address→value map."""
        register_map: dict[int, int] = {}

        for start_addr, count in parser.register_ranges:
            try:
                result = await self._client.read_holding_registers(
                    start_addr, count=count, device_id=unit_id
                )
            except (ConnectionException, OSError):
                raise
            except Exception as err:
                raise VictronModbusError(
                    f"Error reading registers {start_addr}-{start_addr + count} "
                    f"from unit {unit_id}"
                ) from err

            if result.isError():
                raise VictronModbusError(
                    f"Modbus error reading registers {start_addr}-{start_addr + count} "
                    f"from unit {unit_id}: {result}"
                )

            for i, value in enumerate(result.registers):
                register_map[start_addr + i] = value

        return register_map
