"""Azure Monitor telemetry client for fetching application insights data."""

import logging
from datetime import timedelta
from typing import Any, Dict, List, Optional

from azure.core.credentials import TokenCredential
from azure.monitor.query import LogsQueryClient, LogsQueryStatus

from .exceptions import TelemetryError

logger = logging.getLogger(__name__)


class TelemetryClient:
    """Client for fetching Azure Monitor telemetry data.

    This client wraps Azure's LogsQueryClient to fetch telemetry from
    Application Insights tables like AppDependencies, AppExceptions,
    and AppAvailabilityResults.

    Example:
        >>> from azure.identity import DefaultAzureCredential
        >>> client = TelemetryClient(
        ...     credential=DefaultAzureCredential(),
        ...     workspace_id="a1a4fc6d-..."
        ... )
        >>> data = client.fetch_all(
        ...     tables=["AppDependencies", "AppExceptions"],
        ...     timespan=timedelta(days=1)
        ... )
    """

    SUPPORTED_TABLES = ["AppDependencies", "AppExceptions", "AppAvailabilityResults"]

    def __init__(self, credential: TokenCredential, workspace_id: str):
        """Initialize telemetry client.

        Args:
            credential: Azure credential for authentication
            workspace_id: Azure Log Analytics workspace ID
        """
        self._credential = credential
        self._workspace_id = workspace_id
        self._client: Optional[LogsQueryClient] = None

    def _get_client(self) -> LogsQueryClient:
        """Get or create LogsQueryClient instance.

        Returns:
            LogsQueryClient instance
        """
        if self._client is None:
            self._client = LogsQueryClient(self._credential)
            logger.debug("Created LogsQueryClient")
        return self._client

    def fetch(
        self,
        table: str,
        timespan: timedelta = timedelta(days=1),
        query: Optional[str] = None,
        limit: int = 1000,
    ) -> List[Dict[str, Any]]:
        """Fetch telemetry from a single table.

        Args:
            table: Table name (AppDependencies, AppExceptions, AppAvailabilityResults)
            timespan: Time range to query
            query: Custom KQL query. If None, uses default "{table} | limit {limit}"
            limit: Maximum rows to return (used in default query)

        Returns:
            List of row dictionaries

        Raises:
            TelemetryError: If table is not supported or query fails
        """
        if table not in self.SUPPORTED_TABLES:
            raise TelemetryError(
                f"Unsupported table: {table}. Supported tables: {self.SUPPORTED_TABLES}"
            )

        if query is None:
            query = f"{table} | limit {limit}"

        logger.info(f"Fetching telemetry from {table} (timespan={timespan})")

        try:
            client = self._get_client()
            response = client.query_workspace(
                workspace_id=self._workspace_id,
                query=query,
                timespan=timespan,
            )

            if response.status == LogsQueryStatus.SUCCESS:
                rows = []
                if response.tables:
                    table_data = response.tables[0]
                    columns = [col for col in table_data.columns]
                    for row in table_data.rows:
                        row_dict = dict(zip(columns, row))
                        rows.append(self._serialize_row(row_dict))

                logger.info(f"Fetched {len(rows)} rows from {table}")
                return rows

            elif response.status == LogsQueryStatus.PARTIAL:
                logger.warning(f"Partial results for {table}: {response.partial_error}")
                rows = []
                if response.partial_data:
                    table_data = response.partial_data[0]
                    columns = [col for col in table_data.columns]
                    for row in table_data.rows:
                        row_dict = dict(zip(columns, row))
                        rows.append(self._serialize_row(row_dict))
                return rows

            else:
                raise TelemetryError(f"Query failed for {table}")

        except TelemetryError:
            raise
        except Exception as e:
            raise TelemetryError(f"Failed to fetch telemetry from {table}: {e}") from e

    def _serialize_row(self, row: Dict[str, Any]) -> Dict[str, Any]:
        """Serialize row values to JSON-compatible types.

        Args:
            row: Row dictionary with potentially non-serializable values

        Returns:
            Row dictionary with JSON-serializable values
        """
        serialized = {}
        for key, value in row.items():
            if hasattr(value, "isoformat"):
                serialized[key] = value.isoformat()
            elif value is None:
                serialized[key] = None
            else:
                serialized[key] = (
                    str(value)
                    if not isinstance(value, (str, int, float, bool, list, dict))
                    else value
                )
        return serialized

    def fetch_all(
        self,
        tables: Optional[List[str]] = None,
        timespan: timedelta = timedelta(days=1),
        custom_queries: Optional[Dict[str, str]] = None,
        limit: int = 1000,
    ) -> Dict[str, List[Dict[str, Any]]]:
        """Fetch telemetry from multiple tables.

        Args:
            tables: List of table names. Defaults to all supported tables.
            timespan: Time range to query
            custom_queries: Optional dict mapping table names to custom KQL queries
            limit: Maximum rows per table (used in default queries)

        Returns:
            Dictionary mapping table names to lists of row dictionaries

        Example:
            >>> data = client.fetch_all(
            ...     tables=["AppDependencies", "AppExceptions"],
            ...     timespan=timedelta(days=2)
            ... )
            >>> print(data.keys())
            dict_keys(['AppDependencies', 'AppExceptions'])
        """
        if tables is None:
            tables = self.SUPPORTED_TABLES.copy()

        custom_queries = custom_queries or {}
        results: Dict[str, List[Dict[str, Any]]] = {}

        for table in tables:
            query = custom_queries.get(table)
            try:
                results[table] = self.fetch(
                    table=table,
                    timespan=timespan,
                    query=query,
                    limit=limit,
                )
            except TelemetryError as e:
                logger.error(f"Failed to fetch {table}: {e}")
                results[table] = []

        return results


__all__ = ["TelemetryClient"]
