"""Usage tracking and reporting for API metrics.

Services use UsageReporter to batch and send usage metrics to vikdata.
Metrics are flexible - services define their own metric types.
"""

import asyncio
import logging
from collections import defaultdict
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any

from .models.usage import UsageMetric, UsageReportRequest

if TYPE_CHECKING:
    from .client import VikdataClient

logger = logging.getLogger(__name__)


class UsageReporter:
    """Usage reporter for tracking and batching metrics.

    Supports two recording modes:
    1. record() - Legacy mode for HTTP request tracking (aggregated)
    2. record_metric() - Flexible mode for any metric type

    Metrics are batched in memory and flushed periodically or when
    batch size is reached.
    """

    def __init__(
        self,
        client: "VikdataClient",
        service: str,
        batch_size: int = 100,
        flush_interval: float = 60.0,
    ) -> None:
        """Initialize usage reporter.

        Args:
            client: VikdataClient instance
            service: Service name (used in metadata)
            batch_size: Number of metrics to batch before flushing
            flush_interval: Interval in seconds to flush metrics
        """
        self.client = client
        self.service = service
        self.batch_size = batch_size
        self.flush_interval = flush_interval

        # Buffer for aggregated API request metrics
        # Key: "{tenant_id}:{endpoint}:{method}"
        self._api_buffer: dict[str, dict[str, Any]] = defaultdict(
            lambda: {"count": 0, "latencies": [], "status_code": None}
        )

        # Buffer for generic metrics (not aggregated, stored as-is)
        self._metric_buffer: list[UsageMetric] = []

        self._lock = asyncio.Lock()
        self._task: asyncio.Task[None] | None = None
        self._running = False

    async def start(self) -> None:
        """Start background flush task."""
        if self._running:
            return

        self._running = True
        self._task = asyncio.create_task(self._flush_loop())
        logger.debug(f"UsageReporter started for service '{self.service}'")

    async def stop(self) -> None:
        """Stop and flush remaining metrics."""
        self._running = False
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        await self.flush()
        logger.debug(f"UsageReporter stopped for service '{self.service}'")

    async def record(
        self,
        tenant_id: str,
        endpoint: str,
        method: str,
        latency_ms: float | None = None,
        status_code: int | None = None,
    ) -> None:
        """Record a single API request (legacy method).

        This aggregates requests by tenant/endpoint/method for efficiency.
        For other metric types, use record_metric() instead.

        Args:
            tenant_id: Tenant ID making the request
            endpoint: API endpoint path
            method: HTTP method
            latency_ms: Request latency in milliseconds
            status_code: HTTP status code
        """
        key = f"{tenant_id}:{endpoint}:{method}"

        async with self._lock:
            self._api_buffer[key]["count"] += 1
            if latency_ms is not None:
                self._api_buffer[key]["latencies"].append(latency_ms)
            if status_code is not None:
                self._api_buffer[key]["status_code"] = status_code

            # Check if we should flush
            total_count = (
                sum(data["count"] for data in self._api_buffer.values())
                + len(self._metric_buffer)
            )
            if total_count >= self.batch_size:
                await self._flush_unlocked()

    async def record_metric(
        self,
        tenant_id: str,
        metric_type: str,
        quantity: float,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """Record a generic usage metric.

        Use this for non-API metrics like:
        - storage_gb_hour
        - telemetry_points
        - data_egress_mb
        - core_hour

        Args:
            tenant_id: Tenant ID
            metric_type: Service-defined metric type
            quantity: Metric quantity
            metadata: Additional context
        """
        metric = UsageMetric(
            tenant_id=tenant_id,
            metric_type=metric_type,
            quantity=quantity,
            timestamp=datetime.now(timezone.utc),
            metadata=metadata,
        )

        async with self._lock:
            self._metric_buffer.append(metric)

            # Check if we should flush
            total_count = (
                sum(data["count"] for data in self._api_buffer.values())
                + len(self._metric_buffer)
            )
            if total_count >= self.batch_size:
                await self._flush_unlocked()

    async def flush(self) -> None:
        """Flush buffered metrics to vikdata."""
        async with self._lock:
            await self._flush_unlocked()

    async def _flush_unlocked(self) -> None:
        """Internal flush without lock (assumes lock is held)."""
        metrics: list[UsageMetric] = []

        # Convert API buffer to metrics
        for key, data in self._api_buffer.items():
            if data["count"] == 0:
                continue

            tenant_id, endpoint, method = key.split(":", 2)

            # Calculate average latency if available
            avg_latency = None
            if data["latencies"]:
                avg_latency = sum(data["latencies"]) / len(data["latencies"])

            metrics.append(
                UsageMetric(
                    tenant_id=tenant_id,
                    metric_type="api_request",
                    quantity=float(data["count"]),
                    timestamp=datetime.now(timezone.utc),
                    metadata={
                        "service": self.service,
                        "endpoint": endpoint,
                        "method": method,
                        "avg_latency_ms": avg_latency,
                        "status_code": data.get("status_code"),
                    },
                )
            )

        # Add generic metrics
        metrics.extend(self._metric_buffer)

        if not metrics:
            return

        try:
            # Send metrics to vikdata
            request = UsageReportRequest(metrics=metrics)
            await self.client.request(
                "POST",
                "/api/v1/usage/report",
                json=request.model_dump(mode="json"),
            )

            # Clear buffers on success
            self._api_buffer.clear()
            self._metric_buffer.clear()

            logger.debug(f"Flushed {len(metrics)} metrics to vikdata")

        except Exception as e:
            # Keep buffers on failure, will retry on next flush
            logger.warning(f"Failed to flush metrics: {e}")

    async def _flush_loop(self) -> None:
        """Background loop for periodic flushing."""
        while self._running:
            await asyncio.sleep(self.flush_interval)
            await self.flush()
