import asyncio
import os
import time
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple, cast

from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.integrations.datadog.datadog_handler import (
    get_datadog_env,
    get_datadog_hostname,
    get_datadog_pod_name,
    get_datadog_service,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.custom_httpx.http_handler import (
    get_async_httpx_client,
    httpxSpecialProvider,
)
from litellm.types.integrations.datadog_cost_management import (
    DatadogFOCUSCostEntry,
)
from litellm.types.utils import StandardLoggingPayload

# Reserved tag keys whose values come from trusted sources (infra env, LiteLLM
# core payload fields, or proxy-controlled auth metadata). User-supplied
# request_tags / metadata cannot overwrite these, even when the key is
# allowlisted via cost_tag_keys, because that would let an authenticated caller
# spoof cost attribution (e.g. request_tags=["team:victim-team"]).
_RESERVED_TAG_KEYS: frozenset = frozenset(
    {
        "env",
        "service",
        "host",
        "pod_name",
        "provider",
        "model",
        "model_id",
        "team",
        "user",
        "model_group",
    }
)


class DatadogCostManagementLogger(CustomBatchLogger):
    def __init__(self, cost_tag_keys: Optional[List[str]] = None, **kwargs):
        self.cost_tag_keys: List[str] = list(cost_tag_keys) if cost_tag_keys else []
        self.dd_api_key = os.getenv("DD_API_KEY")
        self.dd_app_key = os.getenv("DD_APP_KEY")
        self.dd_site = os.getenv("DD_SITE", "datadoghq.com")

        if not self.dd_api_key or not self.dd_app_key:
            verbose_logger.warning(
                "Datadog Cost Management: DD_API_KEY and DD_APP_KEY are required. Integration will not work."
            )

        self.upload_url = f"https://api.{self.dd_site}/api/v2/cost/custom_costs"

        self.async_client = get_async_httpx_client(
            llm_provider=httpxSpecialProvider.LoggingCallback
        )

        # Initialize lock and start periodic flush task
        self.flush_lock = asyncio.Lock()
        asyncio.create_task(self.periodic_flush())

        # Check if flush_lock is already in kwargs to avoid double passing (unlikely but safe)
        if "flush_lock" not in kwargs:
            kwargs["flush_lock"] = self.flush_lock

        super().__init__(**kwargs)

    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
        try:
            standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
                "standard_logging_object", None
            )

            if standard_logging_object is None:
                return

            # Only log if there is a cost associated
            if standard_logging_object.get("response_cost", 0) > 0:
                self.log_queue.append(standard_logging_object)

                if len(self.log_queue) >= self.batch_size:
                    await self.async_send_batch()

        except Exception as e:
            verbose_logger.exception(
                f"Datadog Cost Management: Error in async_log_success_event: {str(e)}"
            )

    async def async_send_batch(self):
        if not self.log_queue:
            return

        batch_to_send = self.log_queue[:]
        self.log_queue = []

        try:
            aggregated_entries = self._aggregate_costs(batch_to_send)
            if not aggregated_entries:
                verbose_logger.debug(
                    "Datadog Cost Management: batch produced no aggregable entries; "
                    "dropping %d log(s) from queue.",
                    len(batch_to_send),
                )
                return
            await self._upload_to_datadog(aggregated_entries)
        except Exception as e:
            self.log_queue = batch_to_send + self.log_queue
            verbose_logger.exception(
                f"Datadog Cost Management: Error in async_send_batch: {str(e)}"
            )

    def _aggregate_costs(
        self, logs: List[StandardLoggingPayload]
    ) -> List[DatadogFOCUSCostEntry]:
        """
        Aggregates costs by Provider, Model, and Date.
        Returns a list of DatadogFOCUSCostEntry.
        """
        aggregator: Dict[
            Tuple[str, str, str, Tuple[Tuple[str, str], ...]], DatadogFOCUSCostEntry
        ] = {}

        for log in logs:
            try:
                # Extract keys for aggregation
                provider = log.get("custom_llm_provider") or "unknown"
                model = log.get("model") or "unknown"
                cost = log.get("response_cost", 0)

                if cost == 0:
                    continue

                # Get date strings (FOCUS format requires specific keys, but for aggregation we group by Day)
                # UTC date
                # We interpret "ChargePeriod" as the day of the request.
                ts = log.get("startTime") or time.time()
                dt = datetime.fromtimestamp(ts)
                date_str = dt.strftime("%Y-%m-%d")

                # ChargePeriodStart and End
                # If we want daily granularity, end date is usually same day or next day?
                # Datadog Custom Costs usually expects periods.
                # "ChargePeriodStart": "2023-01-01", "ChargePeriodEnd": "2023-12-31" in example.
                # If we send daily, we can say Start=Date, End=Date.

                # Grouping Key: Provider + Model + Date + Tags?
                # For simplicity, let's aggregate by Provider + Model + Date first.
                # If we handle tags, we need to include them in the key.

                tags = self._extract_tags(log)
                tags_key = tuple(sorted(tags.items())) if tags else ()

                key = (provider, model, date_str, tags_key)

                if key not in aggregator:
                    aggregator[key] = {
                        "ProviderName": provider,
                        "ChargeDescription": f"LLM Usage for {model}",
                        "ChargePeriodStart": date_str,
                        "ChargePeriodEnd": date_str,
                        "BilledCost": 0.0,
                        "BillingCurrency": "USD",
                        "Tags": tags if tags else None,
                    }

                aggregator[key]["BilledCost"] += cost

            except Exception as e:
                verbose_logger.warning(
                    f"Error processing log for cost aggregation: {e}"
                )
                continue

        return list(aggregator.values())

    def _extract_tags(self, log: StandardLoggingPayload) -> Dict[str, str]:
        tags: Dict[str, str] = {
            "env": get_datadog_env(),
            "service": get_datadog_service(),
            "host": get_datadog_hostname(),
            "pod_name": get_datadog_pod_name(),
        }

        # Always-on canonical FOCUS dimensions from top-level payload fields.
        # Non-sensitive and required for Datadog Custom Costs per-model attribution.
        self._add_tag(tags, "provider", log.get("custom_llm_provider"))
        self._add_tag(tags, "model", log.get("model"))
        self._add_tag(tags, "model_id", log.get("model_id"))

        # cast because StandardLoggingMetadata is a TypedDict; we iterate it
        # as a generic mapping below.
        metadata: Dict[str, Any] = cast(Dict[str, Any], log.get("metadata") or {})

        # Backwards-compat: team/user/model_group preserved regardless of allowlist.
        if metadata.get("user_api_key_alias"):
            tags["user"] = str(metadata["user_api_key_alias"])
        team_tag = (
            metadata.get("user_api_key_team_alias")
            or metadata.get("team_alias")
            or metadata.get("user_api_key_team_id")
            or metadata.get("team_id")
        )
        if team_tag:
            tags["team"] = str(team_tag)
        if metadata.get("model_group"):
            tags["model_group"] = str(metadata["model_group"])

        # Allowlist-gated: request_tags (split on `:`) and arbitrary metadata.*.
        # Reserved keys are hard-blocked here regardless of allowlist membership —
        # see _RESERVED_TAG_KEYS for the rationale.
        if self.cost_tag_keys:
            allow = set(self.cost_tag_keys)
            for rt in log.get("request_tags") or []:
                if not isinstance(rt, str) or ":" not in rt:
                    continue
                k, _, v = rt.partition(":")
                if k in allow and v:
                    self._set_custom_tag(tags, k, v)
            for k, v in metadata.items():
                if k in allow and v is not None and not isinstance(v, (dict, list)):
                    self._set_custom_tag(tags, k, str(v))
            for nested_key in ("spend_logs_metadata", "requester_metadata"):
                nested = metadata.get(nested_key)
                if isinstance(nested, dict):
                    for k, v in nested.items():
                        if (
                            k in allow
                            and v is not None
                            and not isinstance(v, (dict, list))
                        ):
                            self._set_custom_tag(tags, k, str(v))

        return tags

    @staticmethod
    def _set_custom_tag(tags: Dict[str, str], key: str, value: str) -> None:
        if key in _RESERVED_TAG_KEYS:
            verbose_logger.debug(
                "Datadog Cost Management: dropping user-supplied tag %r=%r — "
                "key is reserved for trusted cost attribution.",
                key,
                value,
            )
            return
        tags[key] = value

    @staticmethod
    def _add_tag(tags: Dict[str, str], key: str, value: Any) -> None:
        if value:
            tags[key] = str(value)

    async def _upload_to_datadog(self, payload: List[Dict]):
        if not self.dd_api_key or not self.dd_app_key:
            return

        headers = {
            "Content-Type": "application/json",
            "DD-API-KEY": self.dd_api_key,
            "DD-APPLICATION-KEY": self.dd_app_key,
        }

        # The API endpoint expects a list of objects directly in the body (file content behavior)
        data_json = safe_dumps(payload)

        response = await self.async_client.put(
            self.upload_url, content=data_json, headers=headers
        )

        response.raise_for_status()

        verbose_logger.debug(
            f"Datadog Cost Management: Uploaded {len(payload)} cost entries. Status: {response.status_code}"
        )
