from __future__ import annotations

import asyncio
import json
import logging
import signal
import socket
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Coroutine, Dict, Optional

import yaml
import aiomqtt

# Import the InfluxDB connection pool (similar to supervisor/evse_poll pattern)

log = logging.getLogger(__name__)

@dataclass
class TopicState:
    """State tracking for a single MQTT topic"""
    last_write_time: float
    window_start_time: float
    last_sample_time: float
    last_v: float
    power_integral: float = 0.0
    max_power: float = 0.0
    energy_start: float = 0.0
    phase: str = "unknown"


@dataclass
class MQTTConfig:
    """Configuration loaded from YAML files."""
    # MQTT settings
    broker: str
    port: int
    mqtt_user: str
    mqtt_password: str
    
    # Aggregation settings
    mqtt_interval: float
    
    # Topics mapping (topic -> asset name)
    topics: Dict[str, str]
    
    # InfluxDB settings
    mons_ip: str
    mons_influx_username: str
    mons_influx_pwd: str
    influx_database: str = "building"
    
    # Heartbeat settings
    ops_db_name: str = "ops"
    heartbeat_interval_s: float = 30.0
    service_name: str = "mqtt_power_logger"
    service_version: str = "1.0.0"

async def mqtt_listener(
    config: MQTTConfig,
    stop_event: asyncio.Event,
    message_handler: Callable[[str, bytes], Coroutine]
) -> None:
    """
    Main MQTT listener loop.
    
    Follows asyncio-mqtt pattern with automatic reconnection.
    Similar to the polling pattern in evse_poll.py but for MQTT.
    
    Args:
        config: MQTT configuration including broker, credentials, and topics
        stop_event: Event to signal shutdown
        message_handler: Async callback that processes messages (topic, payload)
    """
    reconnect_interval = 5.0
    
    while not stop_event.is_set():
        try:
            async with aiomqtt.Client(
                hostname=config.broker,
                port=config.port,
                username=config.mqtt_user,
                password=config.mqtt_password,
                identifier=f"power_logger_{socket.gethostname()}",
            ) as client:
                
                log.info("Connected to MQTT broker at %s:%d", config.broker, config.port)
                
                # Subscribe to all configured topics
                for topic in config.topics.keys():
                    await client.subscribe(topic)
                    log.info("Subscribed to topic: %s", topic)
                    
                # Message processing loop
                async for message in client.messages:
                    if stop_event.is_set():
                        break
                    await message_handler(str(message.topic), message.payload)
                        
        except aiomqtt.MqttError as e:
            log.warning("MQTT connection error: %s, reconnecting in %.1fs", e, reconnect_interval)
            try:
                await asyncio.wait_for(
                    stop_event.wait(),
                    timeout=reconnect_interval
                )
                break  # stop_event was set
            except asyncio.TimeoutError:
                continue  # Reconnect
                
        except asyncio.CancelledError:
            log.info("MQTT listener cancelled")
            raise
            
        except Exception:
            log.exception("Unexpected error in MQTT listener, reconnecting")
            await asyncio.sleep(reconnect_interval)