wled-screen-controller-mixed/custom_components/wled_screen_controller/config_flow.py

"""Config flow for WLED Screen Controller integration."""
from __future__ import annotations

import logging
from typing import Any
from urllib.parse import urlparse, urlunparse

import aiohttp
import voluptuous as vol

from homeassistant import config_entries
from homeassistant.const import CONF_NAME
from homeassistant.core import HomeAssistant
from homeassistant.data_entry_flow import FlowResult
from homeassistant.helpers.aiohttp_client import async_get_clientsession

from .const import DOMAIN, CONF_SERVER_URL, DEFAULT_TIMEOUT

_LOGGER = logging.getLogger(__name__)

STEP_USER_DATA_SCHEMA = vol.Schema(
    {
        vol.Required(CONF_NAME, default="WLED Screen Controller"): str,
        vol.Required(CONF_SERVER_URL, default="http://localhost:8080"): str,
    }
)


def normalize_url(url: str) -> str:
    """Normalize URL to ensure port is an integer."""
    parsed = urlparse(url)

    # If port is specified, ensure it's an integer
    if parsed.port is not None:
        # Reconstruct URL with integer port
        netloc = parsed.hostname or "localhost"
        port = int(parsed.port)  # Cast to int to avoid float
        if port != (443 if parsed.scheme == "https" else 80):
            netloc = f"{netloc}:{port}"

        parsed = parsed._replace(netloc=netloc)

    return urlunparse(parsed)


async def validate_server_connection(
    hass: HomeAssistant, server_url: str
) -> dict[str, Any]:
    """Validate the server URL by checking the health endpoint."""
    session = async_get_clientsession(hass)

    try:
        async with session.get(
            f"{server_url}/health",
            timeout=aiohttp.ClientTimeout(total=DEFAULT_TIMEOUT),
        ) as response:
            if response.status == 200:
                data = await response.json()
                return {
                    "version": data.get("version", "unknown"),
                    "status": data.get("status", "unknown"),
                }
            raise ConnectionError(f"Server returned status {response.status}")

    except aiohttp.ClientError as err:
        raise ConnectionError(f"Cannot connect to server: {err}")
    except Exception as err:
        raise ConnectionError(f"Unexpected error: {err}")


class WLEDScreenControllerConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
    """Handle a config flow for WLED Screen Controller."""

    VERSION = 1

    async def async_step_user(
        self, user_input: dict[str, Any] | None = None
    ) -> FlowResult:
        """Handle the initial step."""
        errors: dict[str, str] = {}

        if user_input is not None:
            server_url = normalize_url(user_input[CONF_SERVER_URL].rstrip("/"))

            try:
                info = await validate_server_connection(self.hass, server_url)

                # Set unique ID based on server URL
                await self.async_set_unique_id(server_url)
                self._abort_if_unique_id_configured()

                return self.async_create_entry(
                    title=user_input[CONF_NAME],
                    data={
                        CONF_SERVER_URL: server_url,
                        "version": info["version"],
                    },
                )

            except ConnectionError as err:
                _LOGGER.error("Connection error: %s", err)
                errors["base"] = "cannot_connect"
            except Exception as err:  # pylint: disable=broad-except
                _LOGGER.exception("Unexpected exception: %s", err)
                errors["base"] = "unknown"

        return self.async_show_form(
            step_id="user",
            data_schema=STEP_USER_DATA_SCHEMA,
            errors=errors,
        )