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

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

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}"

# Gyldige caller_type-verdier
# agent   = OPAX/Jason kjørt som agent via /run eller /run/dag
# browser = direkte kall fra browser-chat (opax.vauco.no)
# cli     = Gemini TUI CLI eller lokal agent.py kjørt manuelt
# cron    = bakgrunnsjobb (billing_agent, anomaly_detector, dag-runner)
# api     = ekstern REST-kall til /run uten sesjon
CallerType   = Literal["agent", "browser", "cli", "cron", "api"]

# Gyldige module_name-verdier (utvides etter hvert som moduler comes online)
# jason/light     = OPAX-agenten, gemini-2.5-flash
# jason/heavy     = OPAX-agenten, gemini-2.5-pro
# costguard       = CostGuard-analyse og anomali-endepunkter
# billing_agent   = ml/billing_agent.py bakgrunns-analyse
# threadstone     = Threadstone-modul (fremtidig)
# cli/gemini-tui  = lokal Gemini TUI (arkivert, men kan aktiveres)
ModuleName = str  # ikke enum — for fremtidssikkerhet

# 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,          # deprecated alias — bruk module_name
    model_name: str,
    input_tokens: int,
    output_tokens: int,
    request_id: Optional[str]  = None,
    module_name: Optional[str] = None,   # CG4: f.eks. 'jason/light', 'costguard'
    caller_type: CallerType    = "agent", # CG4: hvem/hva som kalte
    session_id: Optional[str]  = None,   # CG4: sesjon-ID for gruppering
) -> None:
    """
    Logg ett LLM-kall til BigQuery-tabellen llm_token_usage.
    Feiler stille slik at applikasjonen aldri krasjer pga logging.

    Labeling-konvensjoner:
      module_name  — hvilken komponent (jason/light, costguard, billing_agent, …)
      caller_type  — hvordan kallet ble trigget (agent, browser, cli, cron, api)
      agent_name   — beholdt for bakoverkompatibilitet; settes til module_name om gitt
    """
    try:
    
        from google.cloud import bigquery
        client = bigquery.Client(project=PROJECT_ID)

        effective_module = module_name or agent_name
        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":         effective_module,       # bakoverkompatibelt felt
            "module_name":        effective_module,       # CG4: nytt felt
            "caller_type":        caller_type,            # CG4: nytt felt
            "session_id":         session_id or "",       # CG4: nytt felt
            "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: module={effective_module} "
                f"caller={caller_type} 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, Terraform eller create_bq_table()
BQ_SCHEMA = [
    {"name": "timestamp",          "type": "TIMESTAMP", "mode": "REQUIRED"},
    {"name": "agent_name",         "type": "STRING",    "mode": "REQUIRED"},  # bakoverkompatibelt
    {"name": "module_name",        "type": "STRING",    "mode": "REQUIRED"},  # CG4
    {"name": "caller_type",        "type": "STRING",    "mode": "REQUIRED"},  # CG4: agent|browser|cli|cron|api
    {"name": "session_id",         "type": "STRING",    "mode": "NULLABLE"},  # CG4
    {"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"},
]


def create_bq_table_if_not_exists() -> None:
    """
    Opprett llm_token_usage-tabellen i BigQuery hvis den ikke finnes.
    Trygt å kalle ved app-oppstart (CREATE TABLE IF NOT EXISTS-semantikk).
    """
    try:
        from google.cloud import bigquery
        client = bigquery.Client(project=PROJECT_ID)
        dataset_ref = client.dataset(DATASET)
        table_ref   = dataset_ref.table(TABLE)

        try:
            client.get_table(table_ref)
            logger.info(f"[token_logger] Tabell {FULL_TABLE} finnes allerede.")
            return
        except Exception:
            pass  # tabell finnes ikke — opprett

        schema = [bigquery.SchemaField(f["name"], f["type"], mode=f["mode"]) for f in BQ_SCHEMA]
        table  = bigquery.Table(table_ref, schema=schema)
        client.create_table(table)
        logger.info(f"[token_logger] Opprettet tabell {FULL_TABLE}.")
    except Exception as e:
        logger.warning(f"[token_logger] Kunne ikke opprette BQ-tabell (non-fatal): {e}")