import torch
from model.nano_gpt import AgentGPT, Config
from agent.recursive_reasoning import RecursiveAgenticLoop
import tiktoken
import sys
import io

# UTF-8 Safety
if sys.stdout.encoding != 'utf-8':
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

def test_greedy_search():
    print("--- 🧠 EAM: Greedy Search Execution ---")
    
    config = Config()
    config.n_layer = 10
    config.n_embd = 640
    model = AgentGPT(config)
    model.eval()
    
    enc = tiktoken.get_encoding("cl100k_base")
    class SimpleWrapper:
        def encode(self, s, **k): return torch.tensor([enc.encode(s)])
        def decode(self, i): 
            if torch.is_tensor(i): i = i.tolist()
            if isinstance(i, list) and len(i) > 0 and isinstance(i[0], list): i = i[0]
            return enc.decode(i)

    loop = RecursiveAgenticLoop(model, SimpleWrapper(), demo_mode=False)
    
    print("Executing Deterministic Reasoning Loop (Greedy)...")
    # We use a short max_recursion for the test
    loop.max_recursion = 2
    output = loop.generate_with_reasoning("Test system logic.")
    
    print(f"\nFinal Answer: {output}")
    print("--- Greedy Search Complete ---")

if __name__ == "__main__":
    test_greedy_search()