"""
Async MQTT Power Logger with InfluxDB Integration

This module provides an asyncio-based MQTT power logger that:
- Uses asyncio-mqtt for non-blocking MQTT operations
- Employs InfluxDBConnectionPool for efficient database writes
- Integrates heartbeat monitoring for service health tracking
- Handles graceful shutdown and error recovery
- Maintains window-based power aggregation per topic

Architecture follows the patterns from supervisor.py and evse_poll.py:
- Main event loop managing multiple async tasks
- Pooled InfluxDB connections with automatic reconnection
- Heartbeat loop running in parallel
- Proper cancellation and cleanup handling
"""

from __future__ import annotations

import asyncio
from importlib.metadata import PackageNotFoundError, version
import logging
import os
import signal
import socket
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, Optional

import yaml
import aiomqtt

# Import the InfluxDB connection pool (similar to supervisor/evse_poll pattern)
log = logging.getLogger(__name__)

from ..mqtt.interface import MQTTConfig, TopicState
from ..mqtt.parsers import parse_mqtt_message
from ..influx.influx_client import pool #as InfluxDBConnectionPool
from ..tools.config import setup_logging, load_secrets, read_env_variable
from ..ops.heartbeat import Heartbeat



class PowerLogger:
    """
    Async MQTT power logger with InfluxDB backend.
    
    Follows the supervisor.py pattern:
    - Manages lifecycle of MQTT connection
    - Spawns multiple async tasks (mqtt listener, heartbeat)
    - Handles graceful shutdown
    """
    
    def __init__(self, config: MQTTConfig):
        self.config = config
        self.state: Dict[str, TopicState] = {}
        self.stop_event = asyncio.Event()
        self.mqtt_task: Optional[asyncio.Task] = None
        self.heartbeat_task: Optional[asyncio.Task] = None
        
    async def initialize(self,secrets: dict):
        """Initialize connections and resources."""
        # Initialize InfluxDB pool with secrets
        if secrets is None:
            log.info("No secrets provided for InfluxDB connection pool")
        if pool._secrets is None:
            pool.initialize(secrets)
        
    def init_topic_state(self, topic: str, now: float, power: float, energy: Optional[float] = None, phase: Optional[str] = None):
        """Initialize state tracking for a new topic."""
        self.state[topic] = TopicState(
            last_write_time=now,
            window_start_time=now,
            last_sample_time=now,
            last_v=power,
            power_integral=0.0,
            max_power=power,
            energy_start=energy if energy is not None else 0.0,
            phase=phase if phase is not None else "unknown",
        )
        log.debug("Initialized state for topic: %s", topic)
        
    async def write_bucket(
        self,
        asset: str,
        avg_power: float,
        phase: Optional[str] = None
    ):
        """
        Write aggregated power data to InfluxDB.
        
        Uses the connection pool pattern from evse_poll.py with asyncio.to_thread
        for non-blocking writes.
        
        Args:
            asset: Asset identifier
            avg_power: Average power in watts
            phase: Phase identifier for Shelly (A, B, C) or None for Frient
        """
        tags = {
            "asset": asset, 
            "flow": "consumption",
        }
        
        # Add phase tag for Shelly messages
        if phase is not None:
            tags["phase"] = phase

        payload = [{
            "measurement": "power",
            "tags": tags,
            "fields": {
                "power_kW": round(avg_power / 1000.0, 3),
            },
        }]
        
        try:
            # Non-blocking write using thread pool (like evse_poll pattern)
            success = await asyncio.to_thread(
                 pool.write_mons,
                 self.config.influx_database,
                 payload
            )
            #log.info("Writing bucket to InfluxDB: %s", payload)
            #success=True
            if success:
                phase_str = f" phase={phase}" if phase else ""
                log.debug(
                    "Wrote bucket for %s%s: avg=%.1f W",
                    asset, phase_str, avg_power
                )
            else:
                log.warning("Failed to write bucket for %s", asset)
                
        except Exception as e:
            log.exception("Error writing bucket for %s: %r", asset, e)
            
    async def process_message(self, topic: str, payload: bytes):
        """
        Process a single MQTT message and update state.
        
        Implements the windowed aggregation logic.
        Uses mqtt.parsers to handle different device formats.
        """
        now = time.time()
        
        # Parse message using device-agnostic parser
        result = parse_mqtt_message(topic, payload)
        if result is None:
            return
        
        power_data = result.power_data
        energy = result.energy

        log.debug("Received message on %s: power_data=%s", topic, power_data)
        
        # Process each power value (Frient has 1, Shelly has 3)
        for phase, power in power_data:
            if power is None:
                continue
                
            # Create a state key: for Frient use topic, for Shelly use topic_phase_X
            state_key = f"{topic}_phase_{phase}" if phase else topic
            
            # Initialize state if this is the first message for this state_key
            if state_key not in self.state:
                self.init_topic_state(state_key, now, power, energy, phase)
                continue
                
            s = self.state[state_key]
            
            # Update power integral (trapezoidal approximation)
            dt = now - s.last_sample_time
            if dt > 0 and s.last_v is not None:
                s.power_integral += s.last_v * dt
                
            s.last_sample_time = now
            s.last_v = power
            
            # Track maximum power in window
            if s.max_power is None or power > s.max_power:
                s.max_power = power
                
            # Check if it's time to write aggregated bucket
            if now - s.last_write_time >= self.config.mqtt_interval:
                window_duration = now - s.window_start_time
                
                if window_duration > 0:
                    avg_power = s.power_integral / window_duration
                else:
                    avg_power = power
                    
                # Calculate energy delta (only for Frient)
                if energy is not None and s.energy_start is not None:
                    delta_energy = energy - s.energy_start
                    if delta_energy < 0:
                        delta_energy = None
                else:
                    delta_energy = None
                    
                # Get asset name from topics mapping
                asset = self.config.topics.get(topic, topic)
                
                await self.write_bucket(
                    asset,
                    avg_power,
                    phase=phase
                )
                
                # Reset window
                s.last_write_time = now
                s.window_start_time = now
                s.power_integral = 0.0
                s.max_power = power
                s.energy_start = energy if energy is not None else s.energy_start
            
    async def mqtt_listener(self):
        """
        Main MQTT listener loop.
        
        Follows asyncio-mqtt pattern with automatic reconnection.
        Similar to the polling pattern in evse_poll.py but for MQTT.
        """
        reconnect_interval = 5.0
        
        while not self.stop_event.is_set():
            try:
                async with aiomqtt.Client(
                    hostname=self.config.broker,
                    port=self.config.port,
                    username=self.config.mqtt_user,
                    password=self.config.mqtt_password,
                    identifier=f"power_logger_{socket.gethostname()}",
                ) as client:
                    
                    log.info("Connected to MQTT broker at %s:%d", self.config.broker, self.config.port)
                    
                    # Subscribe to all configured topics
                    for topic in self.config.topics.keys():
                        await client.subscribe(topic)
                        log.info("Subscribed to topic: %s", topic)
                        
                    # Message processing loop
                    async for message in client.messages:
                        if self.stop_event.is_set():
                            break
                        await self.process_message(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(
                        self.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)
                
    async def heartbeat_loop(self):
        """
        Background heartbeat loop for service health monitoring.
        
        Follows the pattern from ops/heartbeat.py.
        """

        SERVICE_NAME = os.getenv("SERVICE_NAME", "mqtt-power")
        HEARTBEAT_INTERVAL_S = float(os.getenv("HEARTBEAT_INTERVAL_S", "5.0"))
        try:    
            SERVICE_VERSION = version("evse-hub")
        except PackageNotFoundError:
            SERVICE_VERSION = "unknown"
        secrets_file=read_env_variable("MQTT_SECRETS_FILE", "secrets.yml")
        secrets = load_secrets(secrets_file)
    


        
        hb = Heartbeat.create_with_pool(
            service=SERVICE_NAME,
            version=SERVICE_VERSION,
            target="mons",
            database=self.config.ops_db_name,
            secrets=secrets
        )
        
        try:
            while not self.stop_event.is_set():
                await hb.write_async(alive=1)
                log.debug("Heartbeat written (alive=1)")
                
                try:
                    await asyncio.wait_for(
                        self.stop_event.wait(),
                        timeout=HEARTBEAT_INTERVAL_S
                    )
                    break  # stop_event was set
                except asyncio.TimeoutError:
                    pass  # Continue loop
                    
        except asyncio.CancelledError:
            log.info("Heartbeat loop cancelled, writing final heartbeat (alive=0)")
            await hb.write_async(alive=0)
            raise
        finally:
            hb.close()
            
    async def run(self,secrets: dict):
        """
        Main orchestration loop.
        
        Follows the supervisor.py pattern:
        - Spawn multiple tasks (MQTT, heartbeat)
        - Wait for completion or cancellation
        - Ensure proper cleanup
        """
        await self.initialize(secrets=secrets)
                                      
        # Spawn tasks
        self.mqtt_task = asyncio.create_task(
            self.mqtt_listener(),
            name="mqtt_listener"
        )
        
        self.heartbeat_task = asyncio.create_task(
            self.heartbeat_loop(),
            name="heartbeat"
        )
        
        log.info("Power logger started, managing %d topics", len(self.config.topics))
        
        try:
            # Wait for stop signal or task completion
            await self.stop_event.wait()
            
        except asyncio.CancelledError:
            log.info("Power logger cancelled")
            raise
            
        finally:
            log.info("Shutting down power logger...")
            
            # Cancel tasks
            if self.mqtt_task and not self.mqtt_task.done():
                self.mqtt_task.cancel()
                
            if self.heartbeat_task and not self.heartbeat_task.done():
                self.heartbeat_task.cancel()
                
            # Wait for tasks to complete
            tasks = [t for t in [self.mqtt_task, self.heartbeat_task] if t]
            if tasks:
                await asyncio.gather(*tasks, return_exceptions=True)
                
            # Close connections
            pool.close()
            log.info("Power logger shutdown complete")
            

def load_config() -> tuple[MQTTConfig, dict]:
    """Load configuration from YAML files."""

    secrets_file=read_env_variable("MQTT_SECRETS_FILE", "secrets.yml")
    topics_file=read_env_variable("MQTT_TOPICS_FILE", "topics.yml")
    secrets = load_secrets(secrets_file)
    topics = load_secrets(topics_file)

    return MQTTConfig(
        broker=secrets['ha_ip'],
        port=secrets['mqtt_port'],
        mqtt_user=secrets['mqtt_user'],
        mqtt_password=secrets['mqtt_user_pwd'],
        mqtt_interval=secrets['mqtt_intervall'],
        topics=topics,
        mons_ip=secrets['mons_ip'],
        mons_influx_username=secrets['mons_influx_username'],
        mons_influx_pwd=secrets['mons_influx_pwd'],
        influx_database=secrets.get('influx_database', 'building'),
        ops_db_name=secrets.get('ops_db_name', 'ops'),
        heartbeat_interval_s=secrets.get('heartbeat_interval_s', 5.0),
        service_name=secrets.get('service_name', 'mqtt_power_logger'),
        service_version=secrets.get('service_version', '1.0.0'),
    ),secrets




async def async_main():
    """Async entry point."""
    config, secrets = load_config()
    logger = PowerLogger(config)
    
    
    # Setup signal handlers for graceful shutdown
    def signal_handler():
        log.info("Received shutdown signal")
        logger.stop_event.set()
        
    loop = asyncio.get_running_loop()
    for sig in (signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, signal_handler)
        
    await logger.run(secrets=secrets)


def main():
    """Synchronous entry point."""
    setup_logging()

    
    try:
        asyncio.run(async_main())
    except KeyboardInterrupt:
        log.info("Interrupted by user")
    except Exception:
        log.exception("Fatal error in main")
        raise


if __name__ == "__main__":
    main()
