from typing import Any
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult


class TokenUsageCallback(BaseCallbackHandler):
    """
    Wyłapuje całkowitą liczbę użytych tokenów po przeprocesowaniu każdego LLM wezwania.
    Inkrementuje licznik bezpośrednio do bazy PostgreSQL dla aktualnego usera.
    """

    def __init__(self, user_id: str):
        self.user_id = user_id

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
        try:
            if response.llm_output and "token_usage" in response.llm_output:
                tokens = response.llm_output["token_usage"].get("total_tokens", 0)
                if tokens > 0:
                    from core.subscription.tracker import increment_tokens

                    increment_tokens(self.user_id, tokens)
        except Exception as e:
            from core.audit_logger import audit_log

            try:
                audit_log("ERROR", f"Błąd inkrementacji tokenów callback: {e}")
            except Exception:
                pass