# custom_components/ai_automation_suggester/sensor.py
"""Sensor platform for AI Automation Suggester."""

from __future__ import annotations

import logging
from typing import cast

from homeassistant.components.sensor import (
    SensorEntity,
    SensorEntityDescription,
    SensorStateClass,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import STATE_UNKNOWN, EntityCategory
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.device_registry import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import (
    CoordinatorEntity,
    DataUpdateCoordinator,
)

from .const import (
    DOMAIN,
    INTEGRATION_NAME,
    CONF_PROVIDER,
    PROVIDER_STATUS_CONNECTED,
    PROVIDER_STATUS_DISCONNECTED,
    PROVIDER_STATUS_ERROR,
    PROVIDER_STATUS_INITIALIZING,
    CONF_MAX_INPUT_TOKENS,
    DEFAULT_MAX_INPUT_TOKENS,
    CONF_MAX_OUTPUT_TOKENS,
    DEFAULT_MAX_OUTPUT_TOKENS,
    # Model configuration keys (used to display current model)
    CONF_OPENAI_MODEL,
    CONF_ANTHROPIC_MODEL,
    CONF_GOOGLE_MODEL,
    CONF_GROQ_MODEL,
    CONF_LOCALAI_MODEL,
    CONF_OLLAMA_MODEL,
    CONF_CUSTOM_OPENAI_MODEL,
    CONF_MISTRAL_MODEL,
    CONF_PERPLEXITY_MODEL,
    CONF_OPENROUTER_MODEL,
    CONF_OPENAI_AZURE_DEPLOYMENT_ID,
    CONF_GENERIC_OPENAI_MODEL,
    DEFAULT_MODELS,
    # Sensor Keys from const.py
    SENSOR_KEY_SUGGESTIONS,
    SENSOR_KEY_STATUS,
    SENSOR_KEY_INPUT_TOKENS,
    SENSOR_KEY_OUTPUT_TOKENS,
    SENSOR_KEY_MODEL,
    SENSOR_KEY_LAST_ERROR,
    SENSOR_KEY_HISTORY_COUNT,
)

_LOGGER = logging.getLogger(__name__)

PROVIDER_TO_MODEL_KEY_MAP: dict[str, str] = {
    "OpenAI": CONF_OPENAI_MODEL,
    "Anthropic": CONF_ANTHROPIC_MODEL,
    "Google": CONF_GOOGLE_MODEL,
    "Groq": CONF_GROQ_MODEL,
    "LocalAI": CONF_LOCALAI_MODEL,
    "Ollama": CONF_OLLAMA_MODEL,
    "Custom OpenAI": CONF_CUSTOM_OPENAI_MODEL,
    "Mistral AI": CONF_MISTRAL_MODEL,
    "Perplexity AI": CONF_PERPLEXITY_MODEL,
    "OpenRouter": CONF_OPENROUTER_MODEL,
    "OpenAI Azure": CONF_OPENAI_AZURE_DEPLOYMENT_ID,
    "Generic OpenAI": CONF_GENERIC_OPENAI_MODEL,
}

SENSOR_DESCRIPTIONS: tuple[SensorEntityDescription, ...] = (
    SensorEntityDescription(
        key=SENSOR_KEY_SUGGESTIONS,
        name="AI Automation Suggestions",
        icon="mdi:robot-happy-outline",
    ),
    SensorEntityDescription(
        key=SENSOR_KEY_STATUS,
        name="AI Provider Status",
        icon="mdi:lan-check",
        entity_category=EntityCategory.DIAGNOSTIC,
    ),
    SensorEntityDescription(
        key=SENSOR_KEY_INPUT_TOKENS,
        name="Max Input Tokens",
        icon="mdi:format-letter-starts-with",
        entity_category=EntityCategory.DIAGNOSTIC,
        native_unit_of_measurement="tokens",
        state_class=SensorStateClass.MEASUREMENT,
    ),
    SensorEntityDescription(
        key=SENSOR_KEY_OUTPUT_TOKENS,
        name="Max Output Tokens",
        icon="mdi:format-letter-ends-with",
        entity_category=EntityCategory.DIAGNOSTIC,
        native_unit_of_measurement="tokens",
        state_class=SensorStateClass.MEASUREMENT,
    ),
    SensorEntityDescription(
        key=SENSOR_KEY_MODEL,
        name="AI Model In Use",
        icon="mdi:brain",
        entity_category=EntityCategory.DIAGNOSTIC,
    ),
    SensorEntityDescription(
        key=SENSOR_KEY_LAST_ERROR,
        name="Last Error Message",
        icon="mdi:alert-circle-outline",
        entity_category=EntityCategory.DIAGNOSTIC,
    ),
    SensorEntityDescription(
        key=SENSOR_KEY_HISTORY_COUNT,
        name="Suggestion History Count",
        icon="mdi:history",
        entity_category=EntityCategory.DIAGNOSTIC,
        state_class=SensorStateClass.MEASUREMENT,
    ),
)

async def async_setup_entry(
    hass: HomeAssistant,
    entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up AI Automation Suggester sensors from a config entry."""
    coordinator = cast(DataUpdateCoordinator, hass.data[DOMAIN][entry.entry_id])
    provider_name = entry.data.get(CONF_PROVIDER, "Unknown Provider")

    entities: list[SensorEntity] = []
    for description in SENSOR_DESCRIPTIONS:
        formatted_name = f"{description.name} ({provider_name})"
        specific_description = SensorEntityDescription(
            key=description.key,
            name=formatted_name,
            icon=description.icon,
            entity_category=description.entity_category,
            native_unit_of_measurement=description.native_unit_of_measurement,
            state_class=description.state_class,
            device_class=description.device_class,
        )

        if description.key == SENSOR_KEY_SUGGESTIONS:
            entities.append(AISuggestionsSensor(coordinator, entry, specific_description))
        elif description.key == SENSOR_KEY_STATUS:
            entities.append(AIProviderStatusSensor(coordinator, entry, specific_description))
        elif description.key == SENSOR_KEY_INPUT_TOKENS:
            entities.append(MaxInputTokensSensor(coordinator, entry, specific_description))
        elif description.key == SENSOR_KEY_OUTPUT_TOKENS:
            entities.append(MaxOutputTokensSensor(coordinator, entry, specific_description))
        elif description.key == SENSOR_KEY_MODEL:
            entities.append(AIModelSensor(coordinator, entry, specific_description))
        elif description.key == SENSOR_KEY_LAST_ERROR:
            entities.append(AILastErrorSensor(coordinator, entry, specific_description))
        elif description.key == SENSOR_KEY_HISTORY_COUNT:
            entities.append(AIHistoryCountSensor(coordinator, entry, specific_description))
        else:
            entities.append(AIBaseSensor(coordinator, entry, specific_description))


    # update_before_add must stay False: in HA 2025.x+, passing True calls
    # CoordinatorEntity.async_update() during platform setup, which triggers a
    # full LLM inference and exceeds the setup timeout (issue #166). Entities
    # populate from coordinator data and refresh on generate_suggestions.
    async_add_entities(entities, False)
    _LOGGER.debug("Sensor platform setup complete for provider: %s", provider_name)

# ─────────────────────────────────────────────────────────────
# Base sensor
# ─────────────────────────────────────────────────────────────
class AIBaseSensor(CoordinatorEntity[DataUpdateCoordinator], SensorEntity):
    """Base class for AI Automation Suggester sensors."""

    def __init__(
        self,
        coordinator: DataUpdateCoordinator,
        entry: ConfigEntry,
        description: SensorEntityDescription,
    ) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator)
        self.entity_description = description
        self._attr_unique_id = f"{entry.entry_id}_{description.key}"
        self._entry = entry
        self._provider_name = entry.data.get(CONF_PROVIDER, "Unknown Provider")

        # Common device info for all sensors of this config entry
        self._attr_device_info = DeviceInfo(
            identifiers={(DOMAIN, entry.entry_id)},
            name=f"{INTEGRATION_NAME} ({self._provider_name})",
            manufacturer="Community",
            model=self._provider_name,
            sw_version=str(entry.version) if entry.version else "N/A",
            configuration_url=None, # Link Github?
        )

    @property
    def available(self) -> bool:
        """Return True if coordinator is available and has data."""
        return super().available and self.coordinator.last_update_success

    @callback
    def _handle_coordinator_update(self) -> None:
        """Handle updated data from the coordinator."""
        if self.coordinator.last_update_success:
            self._update_state_and_attributes()
        super()._handle_coordinator_update()

    def _update_state_and_attributes(self) -> None:
        """Update the sensor's state and attributes based on coordinator data.

        This method should be overridden by subclasses.
        """
        self._attr_native_value = STATE_UNKNOWN
        _LOGGER.debug(
            "Sensor %s._update_state_and_attributes not fully implemented for key %s",
            self.__class__.__name__,
            self.entity_description.key
        )

# ─────────────────────────────────────────────────────────────
# Suggestions sensor
# ─────────────────────────────────────────────────────────────
class AISuggestionsSensor(AIBaseSensor):
    """Shows the availability of new AI suggestions."""
    _attr_should_poll = False

    def __init__(
        self,
        coordinator: DataUpdateCoordinator,
        entry: ConfigEntry,
        description: SensorEntityDescription,
    ) -> None:
        super().__init__(coordinator, entry, description)
        self._previous_suggestions_timestamp: float | None = None

        # Initialize state with default values
        self._attr_native_value = "No Suggestions"
        self._attr_extra_state_attributes = {
            "suggestions": "No suggestions yet",
            "description": None,
            "yaml_block": None,
            "last_update": None,
            "entities_processed": [],
            "provider": self._entry.data.get(CONF_PROVIDER, "unknown"),
            "entities_processed_count": 0,
            "suggestion_count": 0,
            "warnings": [],
        }

    async def async_added_to_hass(self) -> None:
        """Handle added to Hass."""
        await super().async_added_to_hass()

        # Update initial state from coordinator if data exists
        if self.coordinator.data:
            self._update_state_and_attributes()


    def _update_state_and_attributes(self) -> None:
        """Update sensor state and attributes."""
        data = self.coordinator.data or {}
        suggestions = data.get("suggestions")
        last_update_timestamp = data.get("last_update")

        if suggestions and suggestions not in ("No suggestions available", "No suggestions yet"):
            if last_update_timestamp and (self._previous_suggestions_timestamp is None or last_update_timestamp > self._previous_suggestions_timestamp):
                self._attr_native_value = "New Suggestions Available"
                self._previous_suggestions_timestamp = last_update_timestamp
            else:
                self._attr_native_value = "Suggestions Available"
        else:
            self._attr_native_value = "No Suggestions"

        self._attr_extra_state_attributes = {
            "suggestions": suggestions,
            "description": data.get("description"),
            "yaml_block": data.get("yaml_block"),
            "last_update": data.get("last_update"),
            "entities_processed": data.get("entities_processed", []),
            "provider": self._entry.data.get(CONF_PROVIDER, "unknown"),
            "model": data.get("model"),
            "entities_processed_count": len(data.get("entities_processed", [])),
            "suggestion": data.get("suggestion"),
            "suggestion_count": data.get("suggestion_count", 0),
            "warnings": data.get("warnings", []),
        }

# ─────────────────────────────────────────────────────────────
# Provider‑status sensor
# ─────────────────────────────────────────────────────────────
class AIProviderStatusSensor(AIBaseSensor):
    """Indicates whether the configured provider is reachable."""
    _attr_should_poll = False

    def __init__(
        self,
        coordinator: DataUpdateCoordinator,
        entry: ConfigEntry,
        description: SensorEntityDescription,
    ) -> None:
        super().__init__(coordinator, entry, description)
        self._update_state_and_attributes()

    def _update_state_and_attributes(self) -> None:
        """Update sensor state and attributes."""
        data = self.coordinator.data or {}
        if not self.coordinator.last_update_success:
            self._attr_native_value = PROVIDER_STATUS_ERROR
        elif not data:
            self._attr_native_value = PROVIDER_STATUS_INITIALIZING
        elif data.get("last_error"):
            self._attr_native_value = PROVIDER_STATUS_ERROR
        elif "suggestions" in data:
             self._attr_native_value = PROVIDER_STATUS_CONNECTED
        else:
            self._attr_native_value = PROVIDER_STATUS_DISCONNECTED

        self._attr_extra_state_attributes = {
            "last_error_message": data.get("last_error", None),
            "last_attempted_update": data.get("last_update"),
            "provider": data.get("provider", self._provider_name),
            "model": data.get("model"),
            "warnings": data.get("warnings", []),
            "response_metadata": data.get("response_metadata", {}),
        }

# ─────────────────────────────────────────────────────────────
# Max Input Token Sensors
# ─────────────────────────────────────────────────────────────
class MaxInputTokensSensor(AIBaseSensor):
    """Shows the configured maximum input tokens."""
    _attr_should_poll = False

    def __init__(
        self,
        coordinator: DataUpdateCoordinator,
        entry: ConfigEntry,
        description: SensorEntityDescription,
    ) -> None:
        super().__init__(coordinator, entry, description)
        self._update_state_and_attributes() # Initial update

    def _update_state_and_attributes(self) -> None:
        """Update sensor state from config entry options or data."""
        self._attr_native_value = self._entry.options.get(
            CONF_MAX_INPUT_TOKENS,
            self._entry.data.get(CONF_MAX_INPUT_TOKENS, DEFAULT_MAX_INPUT_TOKENS)
        )

# ─────────────────────────────────────────────────────────────
# Max Output Token Sensors
# ─────────────────────────────────────────────────────────────
class MaxOutputTokensSensor(AIBaseSensor):
    """Shows the configured maximum output tokens."""
    _attr_should_poll = False

    def __init__(
        self,
        coordinator: DataUpdateCoordinator,
        entry: ConfigEntry,
        description: SensorEntityDescription,
    ) -> None:
        super().__init__(coordinator, entry, description)
        self._update_state_and_attributes() # Initial update

    def _update_state_and_attributes(self) -> None:
        """Update sensor state from config entry options or data."""
        self._attr_native_value = self._entry.options.get(
            CONF_MAX_OUTPUT_TOKENS,
            self._entry.data.get(CONF_MAX_OUTPUT_TOKENS, DEFAULT_MAX_OUTPUT_TOKENS)
        )

# ─────────────────────────────────────────────────────────────
# Model Sensor
# ─────────────────────────────────────────────────────────────
class AIModelSensor(AIBaseSensor):
    """Shows the currently configured AI model."""
    _attr_should_poll = False

    def __init__(
        self,
        coordinator: DataUpdateCoordinator,
        entry: ConfigEntry,
        description: SensorEntityDescription,
    ) -> None:
        super().__init__(coordinator, entry, description)
        self._update_state_and_attributes()

    def _update_state_and_attributes(self) -> None:
        """Update sensor state with the configured model."""
        provider = self._entry.data.get(CONF_PROVIDER)
        if not provider:
            self._attr_native_value = STATE_UNKNOWN
            return

        model_key = PROVIDER_TO_MODEL_KEY_MAP.get(provider)
        if not model_key:
            self._attr_native_value = "Unknown Model Key"
            _LOGGER.warning("No model key found for provider: %s", provider)
            return

        self._attr_native_value = self._entry.options.get(
            model_key,
            self._entry.data.get(model_key, DEFAULT_MODELS.get(provider, "unknown"))
        ) if model_key else "unknown"

# ─────────────────────────────────────────────────────────────
# Last Error sensor
# ─────────────────────────────────────────────────────────────
class AILastErrorSensor(AIBaseSensor):
    """Shows the last error message from the AI provider."""
    _attr_should_poll = False

    def __init__(
        self,
        coordinator: DataUpdateCoordinator,
        entry: ConfigEntry,
        description: SensorEntityDescription,
    ) -> None:
        super().__init__(coordinator, entry, description)
        self._update_state_and_attributes() # Initial update

    def _update_state_and_attributes(self) -> None:
        """Update sensor state with the last error message."""
        data = self.coordinator.data or {}
        last_error = data.get("last_error")
        self._attr_native_value = str(last_error) if last_error else "No Error"
        self._attr_extra_state_attributes = {
             "last_error_timestamp": data.get("last_update") if last_error else None,
        }


class AIHistoryCountSensor(AIBaseSensor):
    """Shows how many suggestions are retained in history."""
    _attr_should_poll = False

    def __init__(
        self,
        coordinator: DataUpdateCoordinator,
        entry: ConfigEntry,
        description: SensorEntityDescription,
    ) -> None:
        super().__init__(coordinator, entry, description)
        self._update_state_and_attributes()

    def _update_state_and_attributes(self) -> None:
        """Update state from coordinator history data."""
        data = self.coordinator.data or {}
        self._attr_native_value = data.get("suggestion_count", 0)
        self._attr_extra_state_attributes = {
            "latest_suggestion_id": (data.get("suggestion") or {}).get("id"),
            "latest_suggestion_status": (data.get("suggestion") or {}).get("status"),
        }