PRICING = { "input": 0.10, "output": 0.40 }

class CostTracker:
    def __init__(self):
        self.input_tokens = 0
        self.output_tokens = 0

    def add(self, usage_metadata):
        if usage_metadata:
            self.input_tokens += getattr(usage_metadata, 'prompt_token_count', 0)
            self.output_tokens += getattr(usage_metadata, 'candidates_token_count', 0)

    def add_manual(self, inp, out):
        self.input_tokens += inp
        self.output_tokens += out

    def get_cost(self):
        return (self.input_tokens / 1e6 * PRICING["input"]) + (self.output_tokens / 1e6 * PRICING["output"])

    def get_summary(self): # ✅ הפונקציה שהייתה חסרה
        return {
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "cost_usd": self.get_cost()
        }

    def print_report(self):
        print(f"💰 COST: ${self.get_cost():.6f}")

# Static Logger
import os
import json
import time
import tempfile

# V260.9: Smarter logging path
def get_log_file_path():
    base_dir = os.path.join(os.getcwd(), "logs")
    # Try creating/accessing local logs dir
    try:
        if not os.path.exists(base_dir):
            os.makedirs(base_dir)
        # Test write permission
        test_file = os.path.join(base_dir, ".test_write")
        with open(test_file, "w") as f: f.write("ok")
        os.remove(test_file)
        return os.path.join(base_dir, "usage.jsonl")
    except (PermissionError, OSError):
        # Fallback to temp
        return os.path.join(tempfile.gettempdir(), "buddy_math_usage.jsonl")

LOG_FILE = get_log_file_path()
LOG_DIR = os.path.dirname(LOG_FILE)

import contextvars

# Stores the total tokens used in the current solve request
current_request_tokens = contextvars.ContextVar('current_request_tokens', default=0)

def log_api_usage(usage_metadata, source="unknown"):
    """Log API usage to proper location and add to current ContextVar"""
    if not usage_metadata:
        return

    try:
        # LOG_DIR check is now redundant but safe
        if not os.path.exists(LOG_DIR):
            os.makedirs(LOG_DIR, exist_ok=True)
            
        entry = {
            "timestamp": time.time(),
            "source": source,
            "input_tokens": getattr(usage_metadata, 'prompt_token_count', 0),
            "output_tokens": getattr(usage_metadata, 'candidates_token_count', 0),
            "total_tokens": getattr(usage_metadata, 'total_token_count', 0)
        }
        
        with open(LOG_FILE, "a", encoding="utf-8") as f:
            f.write(json.dumps(entry) + "\n")
            
        # Update the ContextVar for the current request
        try:
            current_request_tokens.set(current_request_tokens.get() + entry["total_tokens"])
        except Exception:
            pass

        # Also print to stdout for immediate visibility in console logs
        print(f"💰 [USAGE] {source}: In={entry['input_tokens']}, Out={entry['output_tokens']}")
        
    except Exception as e:
        # Fallback to pure stdout if file logging completely fails
        print(f"💰 [USAGE-STDOUT only] {source}: In={getattr(usage_metadata, 'prompt_token_count', 0)}, Out={getattr(usage_metadata, 'candidates_token_count', 0)}")
        print(f"⚠️ [USAGE LOG ERROR] Failed to log to file {LOG_FILE}: {e}")