import torch
import os
from agent.inference import InferenceEngine
from model.config import ModelConfig
from model.transformer import GPT
from model.draft_model import DraftModel

def test_inference_and_reasoning():
    print("--- Testing Agentic Intelligence Suite ---")
    
    # 1. Initialize Large Model with MoE
    config = ModelConfig(n_experts=8, top_k=2)
    model = GPT(config)
    print(f"Model initialized with {config.n_experts} experts (Top-{config.top_k} routing).")
    
    # 2. Test SparseMoE forward pass
    dummy_input = torch.randint(0, config.vocab_size, (1, 10))
    logits, _ = model(dummy_input)
    print(f"Forward pass successful. Logits shape: {logits.shape}")
    
    # 3. Test Speculative Decoding Hook
    engine = InferenceEngine()
    engine.model = model # Use dummy model for testing architecture
    engine.tokenizer = engine.tokenizer if hasattr(engine, 'tokenizer') else None
    
    # Initialize Draft Model
    engine.draft_model = DraftModel(config, draft_layers=1, draft_embd=64)
    print("Draft model initialized for speculative decoding.")
    
    # 4. Mock Agentic Loop (Manually triggering the chat method structure)
    print("\n--- Testing Agentic Loop (Simulation) ---")
    prompt = "What is the capital of France?"
    
    # We simulate a tool call here to see if the engine handles it
    test_text = "[SYSTEM] Assistant [USER] Calculate 5+5 [THOUGHT] I need to use the calculator tool. [TOOL_CALL] calculator [TOOL_ARG] 5+5"
    
    from agent.tool_executor import parse_and_execute_tools
    result = parse_and_execute_tools(test_text)
    print(f"Tool Parsing Test: {test_text[:40]}... -> Result: {result}")
    
    if result == "[TOOL_RESULT] 10":
        print("SUCCESS: Tool execution logic verified.")
    else:
        print(f"FAILURE: Expected [TOOL_RESULT] 10, got {result}")

    # 5. Testing Self-Critique Detection
    critique_prompt = "[THOUGHT] I might have made an error in the logic..."
    if "error" in critique_prompt.lower():
        print("SUCCESS: Self-Critique trigger detected correctly.")
    
    print("\nTest Suite Completed.")

if __name__ == "__main__":
    test_inference_and_reasoning()