#!/usr/bin/env python3
"""
TEST: TITAN MEMORY INTEGRATION
==============================
Verifies the 'Titan Memory' (Context Caching) capability.
"""

import sys
import time
from pathlib import Path

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from core.gemini_executor import GeminiExecutor, TitanMemoryManager

def test_titan_memory():
    print("=== TESTING TITAN MEMORY INTEGRATION ===\n")
    
    # 1. Initialize Executor
    print("[1] Initializing GeminiExecutor...")
    executor = GeminiExecutor()
    if not executor.api_key:
        print("FAIL: No API Key found.")
        return
    print("PASS: Executor initialized.")
    
    # 2. Initialize Titan Manager
    print("\n[2] Initializing TitanMemoryManager...")
    try:
        titan = TitanMemoryManager(executor.api_key)
        print("PASS: Manager initialized.")
    except Exception as e:
        print(f"FAIL: {e}")
        return

    # 3. Create dummy file for cache
    print("\n[3] Preparing test data...")
    test_file = Path("test_titan_data.txt")
    # Pad to exceed 4096 tokens (minimum for caching)
    padding = "This is padding text to ensure we meet the minimum token count requirement for Titan Memory caching. " * 500
    test_file.write_text(f"This is a secret code: OPERATION_AWAKENING_ALPHA. If asked about the secret code, reply with this exact string.\n\n{padding}")
    print(f"PASS: Created {test_file} (Padded for Caching)")
    
    try:
        # 4. Create Memory
        print("\n[4] Creating Titan Memory Block (Context Cache)...")
        print("    (This makes a real API call to Google - may take 10-20s)")
        memory = titan.create_memory(
            files=[test_file],
            display_name=f"test_titan_{int(time.time())}",
            ttl_minutes=5 # Short TTL for test
        )
        
        if not memory:
            print("FAIL: Check logs/console for API error.")
            return
            
        print(f"PASS: Memory Created! Name: {memory.name}")
        print(f"      Tokens: {memory.token_count}")
        
        # 5. Use Memory
        print("\n[5] Executing with Titan Memory...")
        
        # Without memory (should fail to know the secret if not in prompt)
        print("    Query A (No Memory): 'What is the secret code in the file?'")
        res_no_mem = executor.execute("What is the secret code in the file? If you don't know, say 'UNKNOWN'.")
        print(f"    Result: {res_no_mem.response.strip()}")
        
        # With memory
        print(f"    Query B (WITH Titan Memory): 'What is the secret code in the file?'")
        res_mem = executor.execute(
            "What is the secret code in the file?",
            cached_content_name=memory.name
        )
        print(f"    Result: {res_mem.response.strip()}")
        
        if "OPERATION_AWAKENING_ALPHA" in res_mem.response:
            print("PASS: Titan Memory successfully recalled the data!")
        else:
            print("FAIL: Did not recall secret code.")
            
    except Exception as e:
        print(f"ERROR: {e}")
    finally:
        if test_file.exists():
            test_file.unlink()

if __name__ == "__main__":
    test_titan_memory()