"""Usage Metering — tracks device usage for billing.

Records:
- Device uptime (connection duration)
- AI task count and VLM token usage
- Screenshot count
- WebRTC streaming duration
- Data transfer volume

Periodically flushes metrics to the platform via MQTT telemetry.
"""

from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)


@dataclass
class UsageMetrics:
    """Accumulated usage metrics for a billing period."""
    device_id: str = ""
    owner_id: str = ""

    # Uptime
    connect_time: float = 0.0
    total_uptime_seconds: float = 0.0

    # AI tasks
    task_count: int = 0
    task_steps: int = 0
    vlm_input_tokens: int = 0
    vlm_output_tokens: int = 0

    # Screenshots
    screenshot_count: int = 0
    screenshot_bytes: int = 0

    # WebRTC streaming
    stream_sessions: int = 0
    stream_duration_seconds: float = 0.0

    # Data transfer
    data_sent_bytes: int = 0
    data_received_bytes: int = 0

    # Period
    period_start: float = field(default_factory=time.time)

    def to_dict(self) -> dict:
        return {
            "device_id": self.device_id,
            "owner_id": self.owner_id,
            "total_uptime_seconds": round(self.total_uptime_seconds, 1),
            "task_count": self.task_count,
            "task_steps": self.task_steps,
            "vlm_input_tokens": self.vlm_input_tokens,
            "vlm_output_tokens": self.vlm_output_tokens,
            "screenshot_count": self.screenshot_count,
            "screenshot_bytes": self.screenshot_bytes,
            "stream_sessions": self.stream_sessions,
            "stream_duration_seconds": round(self.stream_duration_seconds, 1),
            "data_sent_bytes": self.data_sent_bytes,
            "data_received_bytes": self.data_received_bytes,
            "period_start": self.period_start,
            "period_end": time.time(),
        }

    def reset(self) -> dict:
        """Reset metrics and return the previous period's data."""
        snapshot = self.to_dict()
        self.task_count = 0
        self.task_steps = 0
        self.vlm_input_tokens = 0
        self.vlm_output_tokens = 0
        self.screenshot_count = 0
        self.screenshot_bytes = 0
        self.stream_sessions = 0
        self.stream_duration_seconds = 0.0
        self.data_sent_bytes = 0
        self.data_received_bytes = 0
        self.period_start = time.time()
        return snapshot


class UsageMeter:
    """Tracks and reports device usage metrics."""

    def __init__(
        self,
        device_id: str,
        owner_id: str = "",
        flush_interval: float = 300.0,
    ) -> None:
        self._metrics = UsageMetrics(device_id=device_id, owner_id=owner_id)
        self._flush_interval = flush_interval
        self._flush_callback = None
        self._running = False
        self._task: asyncio.Task | None = None

    @property
    def metrics(self) -> UsageMetrics:
        self._metrics.total_uptime_seconds = time.time() - self._metrics.connect_time
        return self._metrics

    def set_flush_callback(self, callback) -> None:
        """Set callback for periodic flush: async fn(dict) -> None."""
        self._flush_callback = callback

    def record_connect(self) -> None:
        self._metrics.connect_time = time.time()

    def record_task_started(self) -> None:
        self._metrics.task_count += 1

    def record_task_step(self) -> None:
        self._metrics.task_steps += 1

    def record_vlm_tokens(self, input_tokens: int, output_tokens: int) -> None:
        self._metrics.vlm_input_tokens += input_tokens
        self._metrics.vlm_output_tokens += output_tokens

    def record_screenshot(self, size_bytes: int) -> None:
        self._metrics.screenshot_count += 1
        self._metrics.screenshot_bytes += size_bytes

    def record_stream_start(self) -> None:
        self._metrics.stream_sessions += 1

    def record_stream_duration(self, seconds: float) -> None:
        self._metrics.stream_duration_seconds += seconds

    def record_data_transfer(self, sent: int = 0, received: int = 0) -> None:
        self._metrics.data_sent_bytes += sent
        self._metrics.data_received_bytes += received

    async def start(self) -> None:
        if self._running:
            return
        self._running = True
        self._task = asyncio.create_task(self._flush_loop())

    async def stop(self) -> None:
        self._running = False
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        # Final flush
        await self._flush()

    async def _flush_loop(self) -> None:
        while self._running:
            await asyncio.sleep(self._flush_interval)
            await self._flush()

    async def _flush(self) -> None:
        if not self._flush_callback:
            return
        try:
            data = self._metrics.reset()
            await self._flush_callback(data)
        except Exception:
            logger.exception("Usage meter flush failed")
