# token_logger.py — Logg LLM token-bruk til BigQuery (CG3e)

import os
import logging
import uuid
from datetime import datetime, timezone
from typing import Optional

logger = logging.getLogger(__name__)

PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "propane-will-491900-m5")
DATASET    = os.environ.get("BQ_BILLING_DATASET", "billing_data")
TABLE      = "llm_token_usage"
FULL_TABLE = f"{PROJECT_ID}.{DATASET}.{TABLE}"

# Prismodell (USD per 1M tokens) — oppdater ved modellbytte
MODEL_PRICING = {
    "gemini-2.5-pro":   {"input": 1.25, "output": 10.00},
    "gemini-2.5-flash": {"input": 0.075, "output": 0.30},
}
DEFAULT_PRICING = {"input": 1.25, "output": 10.00}


def _estimate_cost(model_name: str, input_tokens: int, output_tokens: int) -> float:
    pricing = MODEL_PRICING.get(model_name, DEFAULT_PRICING)
    cost = (input_tokens / 1_000_000) * pricing["input"] + \
           (output_tokens / 1_000_000) * pricing["output"]
    return round(cost, 8)


def log_token_usage(
    agent_name: str,
    model_name: str,
    input_tokens: int,
    output_tokens: int,
    request_id: Optional[str] = None,
) -> None:
    """
    Logg ett LLM-kall til BigQuery-tabellen llm_token_usage.
    Feiler stille slik at applikasjonen aldri krasjer pga logging.
    """
    try:
        from google.cloud import bigquery
        client = bigquery.Client(project=PROJECT_ID)

        total_tokens = input_tokens + output_tokens
        estimated_cost = _estimate_cost(model_name, input_tokens, output_tokens)

        row = {
            "timestamp":          datetime.now(timezone.utc).isoformat(),
            "agent_name":         agent_name,
            "model_name":         model_name,
            "input_tokens":       input_tokens,
            "output_tokens":      output_tokens,
            "total_tokens":       total_tokens,
            "estimated_cost_usd": estimated_cost,
            "request_id":         request_id or str(uuid.uuid4()),
        }

        errors = client.insert_rows_json(FULL_TABLE, [row])
        if errors:
            logger.warning(f"[token_logger] BQ insert errors: {errors}")
        else:
            logger.info(
                f"[token_logger] Logged: agent={agent_name} model={model_name} "
                f"in={input_tokens} out={output_tokens} cost=${estimated_cost:.6f}"
            )
    except Exception as e:
        logger.warning(f"[token_logger] Failed to log token usage (non-fatal): {e}")


# BQ table schema — brukes som referanse ved manuell oppretting eller Terraform
BQ_SCHEMA = [
    {"name": "timestamp",          "type": "TIMESTAMP", "mode": "REQUIRED"},
    {"name": "agent_name",         "type": "STRING",    "mode": "REQUIRED"},
    {"name": "model_name",         "type": "STRING",    "mode": "REQUIRED"},
    {"name": "input_tokens",       "type": "INTEGER",   "mode": "REQUIRED"},
    {"name": "output_tokens",      "type": "INTEGER",   "mode": "REQUIRED"},
    {"name": "total_tokens",       "type": "INTEGER",   "mode": "REQUIRED"},
    {"name": "estimated_cost_usd", "type": "FLOAT",     "mode": "REQUIRED"},
    {"name": "request_id",         "type": "STRING",    "mode": "NULLABLE"},
]