"""SQLAlchemy repository for usage event persistence and aggregation."""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass
from datetime import datetime
from uuid import UUID

from sqlalchemy import func, select

from noteflow.application.observability.ports import UsageEvent
from noteflow.infrastructure.persistence.models.observability.usage_event import (
    UsageEventModel,
)
from noteflow.infrastructure.persistence.repositories._base import BaseRepository

# Column label constants for aggregate queries
_COL_EVENT_COUNT = "event_count"
_COL_TOTAL_TOKENS_INPUT = "total_tokens_input"
_COL_TOTAL_TOKENS_OUTPUT = "total_tokens_output"
_COL_AVG_LATENCY_MS = "avg_latency_ms"
_COL_SUCCESS_COUNT = "success_count"
_COL_ERROR_COUNT = "error_count"


@dataclass(frozen=True, slots=True)
class UsageAggregate:
    """Aggregated usage statistics for a time period or grouping.

    Provides summary metrics for analytics dashboards and billing.
    """

    event_count: int
    """Total number of events."""

    total_tokens_input: int
    """Sum of input tokens across all events."""

    total_tokens_output: int
    """Sum of output tokens across all events."""

    avg_latency_ms: float | None
    """Average latency in milliseconds."""

    success_count: int
    """Number of successful events."""

    error_count: int
    """Number of failed events."""


@dataclass(frozen=True, slots=True)
class ProviderUsageAggregate:
    """Usage statistics grouped by provider.

    Enables per-provider analytics and cost allocation.
    """

    provider_name: str
    """Provider identifier (e.g., 'openai', 'anthropic', 'ollama')."""

    model_name: str | None
    """Most commonly used model for this provider."""

    event_count: int
    """Total number of events for this provider."""

    total_tokens_input: int
    """Sum of input tokens."""

    total_tokens_output: int
    """Sum of output tokens."""

    avg_latency_ms: float | None
    """Average latency in milliseconds."""


class SqlAlchemyUsageEventRepository(BaseRepository):
    """SQLAlchemy implementation of usage event repository.

    Provides persistence and aggregation queries for usage analytics.
    """

    async def add(self, event: UsageEvent) -> UsageEvent:
        """Persist a usage event.

        Args:
            event: Usage event to persist.

        Returns:
            Persisted event (unchanged, as events are immutable).
        """
        model = UsageEventModel(
            event_type=event.event_type,
            timestamp=event.timestamp,
            meeting_id=UUID(event.meeting_id) if event.meeting_id else None,
            workspace_id=UUID(event.workspace_id) if event.workspace_id else None,
            project_id=UUID(event.project_id) if event.project_id else None,
            provider_name=event.provider_name,
            model_name=event.model_name,
            tokens_input=event.tokens_input,
            tokens_output=event.tokens_output,
            latency_ms=event.latency_ms,
            success=event.success,
            error_code=event.error_code,
            attributes=dict(event.attributes),
        )
        await self._add_and_flush(model)
        return event

    async def add_batch(self, events: Sequence[UsageEvent]) -> int:
        """Persist multiple usage events efficiently.

        Args:
            events: Usage events to persist.

        Returns:
            Number of events persisted.
        """
        if not events:
            return 0

        models = [
            UsageEventModel(
                event_type=e.event_type,
                timestamp=e.timestamp,
                meeting_id=UUID(e.meeting_id) if e.meeting_id else None,
                workspace_id=UUID(e.workspace_id) if e.workspace_id else None,
                project_id=UUID(e.project_id) if e.project_id else None,
                provider_name=e.provider_name,
                model_name=e.model_name,
                tokens_input=e.tokens_input,
                tokens_output=e.tokens_output,
                latency_ms=e.latency_ms,
                success=e.success,
                error_code=e.error_code,
                attributes=dict(e.attributes),
            )
            for e in events
        ]
        await self._add_all_and_flush(models)
        return len(models)

    async def get_by_meeting(
        self,
        meeting_id: UUID,
        event_type: str | None = None,
        limit: int = 100,
    ) -> Sequence[UsageEvent]:
        """Retrieve usage events for a meeting.

        Args:
            meeting_id: Meeting to get events for.
            event_type: Optional filter by event type.
            limit: Maximum number of events.

        Returns:
            Sequence of usage events, newest first.
        """
        stmt = (
            select(UsageEventModel)
            .where(UsageEventModel.meeting_id == meeting_id)
            .order_by(UsageEventModel.timestamp.desc())
            .limit(limit)
        )
        if event_type:
            stmt = stmt.where(UsageEventModel.event_type == event_type)

        models = await self._execute_scalars(stmt)
        return [self._to_domain(m) for m in models]

    async def get_in_time_range(
        self,
        start_time: datetime,
        end_time: datetime,
        event_type: str | None = None,
        provider_name: str | None = None,
        workspace_id: UUID | None = None,
        limit: int = 1000,
    ) -> Sequence[UsageEvent]:
        """Retrieve usage events within a time range.

        Args:
            start_time: Start of time range (inclusive).
            end_time: End of time range (exclusive).
            event_type: Optional filter by event type.
            provider_name: Optional filter by provider.
            workspace_id: Optional filter by workspace.
            limit: Maximum number of events.

        Returns:
            Sequence of usage events, newest first.
        """
        stmt = (
            select(UsageEventModel)
            .where(UsageEventModel.timestamp >= start_time)
            .where(UsageEventModel.timestamp < end_time)
            .order_by(UsageEventModel.timestamp.desc())
            .limit(limit)
        )

        if event_type:
            stmt = stmt.where(UsageEventModel.event_type == event_type)
        if provider_name:
            stmt = stmt.where(UsageEventModel.provider_name == provider_name)
        if workspace_id:
            stmt = stmt.where(UsageEventModel.workspace_id == workspace_id)

        models = await self._execute_scalars(stmt)
        return [self._to_domain(m) for m in models]

    async def aggregate(
        self,
        start_time: datetime,
        end_time: datetime,
        event_type: str | None = None,
        provider_name: str | None = None,
        workspace_id: UUID | None = None,
    ) -> UsageAggregate:
        """Calculate aggregate statistics for a time range.

        Args:
            start_time: Start of time range (inclusive).
            end_time: End of time range (exclusive).
            event_type: Optional filter by event type.
            provider_name: Optional filter by provider.
            workspace_id: Optional filter by workspace.

        Returns:
            Aggregated usage statistics.
        """
        stmt = select(
            func.count(UsageEventModel.id).label(_COL_EVENT_COUNT),
            func.coalesce(func.sum(UsageEventModel.tokens_input), 0).label(
                _COL_TOTAL_TOKENS_INPUT
            ),
            func.coalesce(func.sum(UsageEventModel.tokens_output), 0).label(
                _COL_TOTAL_TOKENS_OUTPUT
            ),
            func.avg(UsageEventModel.latency_ms).label(_COL_AVG_LATENCY_MS),
            func.count(UsageEventModel.id)
            .filter(UsageEventModel.success.is_(True))
            .label(_COL_SUCCESS_COUNT),
            func.count(UsageEventModel.id)
            .filter(UsageEventModel.success.is_(False))
            .label(_COL_ERROR_COUNT),
        ).where(
            UsageEventModel.timestamp >= start_time,
            UsageEventModel.timestamp < end_time,
        )

        if event_type:
            stmt = stmt.where(UsageEventModel.event_type == event_type)
        if provider_name:
            stmt = stmt.where(UsageEventModel.provider_name == provider_name)
        if workspace_id:
            stmt = stmt.where(UsageEventModel.workspace_id == workspace_id)

        result = await self._session.execute(stmt)
        row = result.one()

        return UsageAggregate(
            event_count=int(row.event_count),
            total_tokens_input=int(row.total_tokens_input),
            total_tokens_output=int(row.total_tokens_output),
            avg_latency_ms=float(row.avg_latency_ms) if row.avg_latency_ms else None,
            success_count=int(row.success_count),
            error_count=int(row.error_count),
        )

    async def aggregate_by_provider(
        self,
        start_time: datetime,
        end_time: datetime,
        event_type: str | None = None,
        workspace_id: UUID | None = None,
    ) -> Sequence[ProviderUsageAggregate]:
        """Calculate aggregate statistics grouped by provider.

        Args:
            start_time: Start of time range (inclusive).
            end_time: End of time range (exclusive).
            event_type: Optional filter by event type.
            workspace_id: Optional filter by workspace.

        Returns:
            List of per-provider aggregated statistics.
        """
        stmt = (
            select(
                UsageEventModel.provider_name,
                func.mode()
                .within_group(UsageEventModel.model_name)
                .label("model_name"),
                func.count(UsageEventModel.id).label(_COL_EVENT_COUNT),
                func.coalesce(func.sum(UsageEventModel.tokens_input), 0).label(
                    _COL_TOTAL_TOKENS_INPUT
                ),
                func.coalesce(func.sum(UsageEventModel.tokens_output), 0).label(
                    _COL_TOTAL_TOKENS_OUTPUT
                ),
                func.avg(UsageEventModel.latency_ms).label(_COL_AVG_LATENCY_MS),
            )
            .where(
                UsageEventModel.timestamp >= start_time,
                UsageEventModel.timestamp < end_time,
                UsageEventModel.provider_name.isnot(None),
            )
            .group_by(UsageEventModel.provider_name)
            .order_by(func.count(UsageEventModel.id).desc())
        )

        if event_type:
            stmt = stmt.where(UsageEventModel.event_type == event_type)
        if workspace_id:
            stmt = stmt.where(UsageEventModel.workspace_id == workspace_id)

        result = await self._session.execute(stmt)
        rows = result.all()

        return [
            ProviderUsageAggregate(
                provider_name=row.provider_name,
                model_name=row.model_name,
                event_count=int(row.event_count),
                total_tokens_input=int(row.total_tokens_input),
                total_tokens_output=int(row.total_tokens_output),
                avg_latency_ms=float(row.avg_latency_ms) if row.avg_latency_ms else None,
            )
            for row in rows
        ]

    async def aggregate_by_event_type(
        self,
        start_time: datetime,
        end_time: datetime,
        workspace_id: UUID | None = None,
    ) -> dict[str, UsageAggregate]:
        """Calculate aggregate statistics grouped by event type.

        Args:
            start_time: Start of time range (inclusive).
            end_time: End of time range (exclusive).
            workspace_id: Optional filter by workspace.

        Returns:
            Dictionary mapping event types to aggregated statistics.
        """
        stmt = (
            select(
                UsageEventModel.event_type,
                func.count(UsageEventModel.id).label(_COL_EVENT_COUNT),
                func.coalesce(func.sum(UsageEventModel.tokens_input), 0).label(
                    _COL_TOTAL_TOKENS_INPUT
                ),
                func.coalesce(func.sum(UsageEventModel.tokens_output), 0).label(
                    _COL_TOTAL_TOKENS_OUTPUT
                ),
                func.avg(UsageEventModel.latency_ms).label(_COL_AVG_LATENCY_MS),
                func.count(UsageEventModel.id)
                .filter(UsageEventModel.success.is_(True))
                .label(_COL_SUCCESS_COUNT),
                func.count(UsageEventModel.id)
                .filter(UsageEventModel.success.is_(False))
                .label(_COL_ERROR_COUNT),
            )
            .where(
                UsageEventModel.timestamp >= start_time,
                UsageEventModel.timestamp < end_time,
            )
            .group_by(UsageEventModel.event_type)
        )

        if workspace_id:
            stmt = stmt.where(UsageEventModel.workspace_id == workspace_id)

        result = await self._session.execute(stmt)
        rows = result.all()

        return {
            row.event_type: UsageAggregate(
                event_count=int(row.event_count),
                total_tokens_input=int(row.total_tokens_input),
                total_tokens_output=int(row.total_tokens_output),
                avg_latency_ms=float(row.avg_latency_ms) if row.avg_latency_ms else None,
                success_count=int(row.success_count),
                error_count=int(row.error_count),
            )
            for row in rows
        }

    async def count_by_event_type(
        self,
        start_time: datetime,
        end_time: datetime,
    ) -> dict[str, int]:
        """Count events by type within a time range.

        Args:
            start_time: Start of time range (inclusive).
            end_time: End of time range (exclusive).

        Returns:
            Dictionary mapping event types to counts.
        """
        stmt = (
            select(
                UsageEventModel.event_type,
                func.count(UsageEventModel.id).label("count"),
            )
            .where(
                UsageEventModel.timestamp >= start_time,
                UsageEventModel.timestamp < end_time,
            )
            .group_by(UsageEventModel.event_type)
        )

        result = await self._session.execute(stmt)
        rows = result.mappings().all()
        return {str(row["event_type"]): int(row["count"]) for row in rows}

    async def delete_before(self, before: datetime) -> int:
        """Delete events older than a given timestamp.

        Useful for data retention policies.

        Args:
            before: Delete events before this timestamp.

        Returns:
            Number of events deleted.
        """
        from sqlalchemy import delete

        stmt = delete(UsageEventModel).where(UsageEventModel.timestamp < before)
        return await self._execute_delete(stmt)

    @staticmethod
    def _to_domain(model: UsageEventModel) -> UsageEvent:
        """Convert ORM model to domain entity.

        Args:
            model: ORM model instance.

        Returns:
            Domain UsageEvent entity.
        """
        return UsageEvent(
            event_type=model.event_type,
            timestamp=model.timestamp,
            meeting_id=str(model.meeting_id) if model.meeting_id else None,
            workspace_id=str(model.workspace_id) if model.workspace_id else None,
            project_id=str(model.project_id) if model.project_id else None,
            provider_name=model.provider_name,
            model_name=model.model_name,
            tokens_input=model.tokens_input,
            tokens_output=model.tokens_output,
            latency_ms=model.latency_ms,
            success=model.success,
            error_code=model.error_code,
            attributes=dict(model.attributes) if model.attributes else {},
        )
