from pydantic import BaseModel, Field
from typing import Dict, List, Callable, Optional
from futuremind.tool.tools.search.log import logger

class SearchSnippet(BaseModel):
    title: str
    url: str
    description: str
    weight: float = 0.0  # TODO
    date: str = ""
    
class LLMUsage(BaseModel):
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0

class TokenTracker:
    def __init__(self, budget: int | None = None):
        self.usages: List[TokenUsage] = []
        self.budget = budget
        self.usage_listeners: List[Callable[[LLMUsage], None]] = []

    def add_usage_listener(self, listener: Callable[[LLMUsage], None]):
        self.usage_listeners.append(listener)

    def track_usage(self, tool: str, usage: LLMUsage):
        self.usages.append(TokenUsage(tool=tool, usage=usage))
        for listener in self.usage_listeners:
            listener(usage)

    def get_total_usage(self) -> LLMUsage:
        total = LLMUsage()
        for usage in self.usages:
            total.prompt_tokens += usage.usage.prompt_tokens
            total.completion_tokens += usage.usage.completion_tokens
            total.total_tokens += usage.usage.total_tokens
        return total

    def get_usage_breakdown(self) -> Dict[str, int]:
        breakdown = defaultdict(int)
        for usage in self.usages:
            breakdown[usage.tool] += usage.usage.total_tokens
        return dict(breakdown)

    def print_summary(self):
        breakdown = self.get_usage_breakdown()

    def reset(self):
        self.usages.clear()
