"""
Finault LiteLLM Callback — Seal every LLM call automatically.

Usage:
    import litellm
    from finault.litellm_callback import FinaultCallback

    litellm.success_callback = [FinaultCallback()]

    # Or in litellm config:
    litellm_settings:
      success_callback: ["finault"]

Every successful LLM call is automatically:
- Costed (model + tokens → USD)
- Sealed (SHA-256 hash chain)
- Attributed (customer, agent, chain)
- Scanned (Dark Debt patterns)
"""

import hashlib
import hmac
import json
import os
import time
from datetime import datetime, timezone
from typing import Any, Dict, Optional

import httpx


class FinaultCallback:
    """LiteLLM callback handler that seals every call via Finault."""

    def __init__(
        self,
        finault_key: Optional[str] = None,
        gateway_url: Optional[str] = None,
        chains: Optional[list] = None,
        revenue_usd: Optional[float] = None,
        customer_id: Optional[str] = None,
        environment: str = "production",
        async_mode: bool = True,
    ):
        self.finault_key = finault_key or os.environ.get("FINAULT_API_KEY", "")
        self.gateway_url = (
            gateway_url
            or os.environ.get("FINAULT_GATEWAY_URL", "https://api.finault.ai")
        ).rstrip("/")
        self.chains = chains or []
        self.default_revenue = revenue_usd
        self.customer_id = customer_id
        self.environment = environment
        self.async_mode = async_mode
        self._client = httpx.Client(timeout=5.0)
        self._async_client = None

    def log_success_event(self, kwargs, response_obj, start_time, end_time):
        """Called by LiteLLM after every successful completion."""
        try:
            seal_data = self._build_seal_data(kwargs, response_obj, start_time, end_time)
            if self.async_mode:
                import threading
                threading.Thread(
                    target=self._send_seal, args=(seal_data,), daemon=True
                ).start()
            else:
                self._send_seal(seal_data)
        except Exception as e:
            # Never block the LLM response
            pass

    def log_failure_event(self, kwargs, response_obj, start_time, end_time):
        """Called by LiteLLM after failed completions."""
        try:
            seal_data = self._build_seal_data(
                kwargs, response_obj, start_time, end_time, failed=True
            )
            self._send_seal(seal_data)
        except Exception:
            pass

    def _build_seal_data(self, kwargs, response_obj, start_time, end_time, failed=False):
        model = kwargs.get("model", "unknown")
        messages = kwargs.get("messages", [])

        # Extract usage from response
        usage = {}
        if hasattr(response_obj, "usage") and response_obj.usage:
            usage = {
                "prompt_tokens": getattr(response_obj.usage, "prompt_tokens", 0),
                "completion_tokens": getattr(response_obj.usage, "completion_tokens", 0),
                "total_tokens": getattr(response_obj.usage, "total_tokens", 0),
            }

        # Extract litellm metadata
        metadata = kwargs.get("litellm_params", {}).get("metadata", {})
        customer_id = metadata.get("customer_id") or self.customer_id
        revenue = metadata.get("revenue_usd") or self.default_revenue
        extra_chains = metadata.get("finault_chains", [])

        latency_ms = int((end_time - start_time).total_seconds() * 1000)

        return {
            "who": {
                "organizationId": metadata.get("org_id", ""),
                "customerId": customer_id,
                "agentId": metadata.get("agent_id", "litellm"),
                "environment": metadata.get("environment", self.environment),
                "source": "litellm_callback",
            },
            "what": {
                "model": model,
                "provider": kwargs.get("custom_llm_provider", "unknown"),
                "inputTokens": usage.get("prompt_tokens", 0),
                "outputTokens": usage.get("completion_tokens", 0),
                "totalTokens": usage.get("total_tokens", 0),
                "latencyMs": latency_ms,
                "status": "error" if failed else "success",
            },
            "worth": {
                "revenueUsd": revenue,
            },
            "chains": list(set(self.chains + extra_chains)),
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }

    def _send_seal(self, seal_data):
        try:
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.finault_key}",
            }
            if seal_data.get("chains"):
                headers["X-Finault-Chains"] = ",".join(seal_data["chains"])

            self._client.post(
                f"{self.gateway_url}/v1/seal/ingest",
                json=seal_data,
                headers=headers,
            )
        except Exception:
            pass  # Never block


# For litellm config-based registration:
# litellm recognizes callback classes by name
finault_callback = FinaultCallback
