import sys
import os
import hashlib
from unittest.mock import patch, MagicMock
from types import SimpleNamespace
from pathlib import Path

# Ensure project root is in path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from src import ai_client

def test_token_usage_tracking() -> None:
 """
    Verify that ai_client.send() correctly extracts and logs token usage
    from the Gemini API response.
    """
 ai_client.reset_session()
 
 # Mock the google-genai Client and chats.create
 with patch("src.ai_client._ensure_gemini_client"), \
   patch("src.ai_client._gemini_client") as mock_client:
  
  mock_chat = MagicMock()
  mock_client.chats.create.return_value = mock_chat
  
  # Create a mock response with usage metadata
  mock_usage = SimpleNamespace(
   prompt_token_count=100,
   candidates_token_count=50,
   total_token_count=150,
   cached_content_token_count=20
  )
  
  mock_candidate = SimpleNamespace(
   content=SimpleNamespace(parts=[SimpleNamespace(text="Mock Response", function_call=None)]),
   finish_reason="STOP"
  )
  
  mock_response = SimpleNamespace(
   candidates=[mock_candidate],
   usage_metadata=mock_usage
  )
  
  mock_chat.send_message.return_value = mock_response
  
  # Set provider to gemini
  ai_client.set_provider("gemini", "gemini-2.5-flash-lite")
  
  # Send a message
  ai_client.send("Context", "Hello")
  
  # Verify usage was logged in the comms log
  comms = ai_client.get_comms_log()
  response_entries = [e for e in comms if e.get("direction") == "IN" and e["kind"] == "response"]
  assert len(response_entries) > 0
  usage = response_entries[0]["payload"]["usage"]
  assert usage["input_tokens"] == 100
  assert usage["output_tokens"] == 50
  assert usage["cache_read_input_tokens"] == 20