from app.services.usage_tracker import UsageTracker
from app.services.smart_router import ENGINE_COST_PROFILES


class UsageRecorder:
    def __init__(self, tracker: UsageTracker):
        self.tracker = tracker

    def calculate_cost(self, engine_type: str, input_tokens: int, output_tokens: int) -> float:
        profile = ENGINE_COST_PROFILES.get(engine_type)
        if not profile:
            return 0.0

        input_cost = (input_tokens / 1_000_000) * profile.input_price_per_million
        output_cost = (output_tokens / 1_000_000) * profile.output_price_per_million
        return round(input_cost + output_cost, 6)

    def record(
        self,
        user_id: str,
        brand_id: str | None,
        engine_type: str,
        query: str,
        input_tokens: int,
        output_tokens: int,
        metadata: dict | None = None,
    ) -> None:
        cost = self.calculate_cost(engine_type, input_tokens, output_tokens)
        self.tracker.record(
            user_id=user_id,
            brand_id=brand_id or "",
            engine_type=engine_type,
            query=query,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cost=cost,
            metadata=metadata or {},
        )
