import sys
import os
sys.path.append(os.getcwd())

from orionagent import Agent, Manager, Gemini
from orionagent.memory.session import Session
from orionagent.agents.hitl import HitlConfig
from unittest.mock import patch

# Mock model to capture the prompt
class MockModel:
    def __init__(self):
        self.last_prompt = None
        self.verbose = True
        self.debug = True
        
    def generate(self, prompt, **kwargs):
        self.last_prompt = prompt
        return "Mock response"

def test_memory_prompt_format():
    mock_llm = MockModel()
    agent = Agent(
        name="Orion-Test",
        model=mock_llm,
        memory="session"
    )
    
    # Setup session with one turn
    session = Session("default_user", "Orion-Test", "test_session_99")
    session.messages = [
        {"role": "user", "content": "hi"},
        {"role": "assistant", "content": "Hello! How can I help?"}
    ]
    agent._session_manager.save(session)
    
    # Ask something
    agent.ask("tell me a joke", stream=False)
    
    prompt = mock_llm.last_prompt
    print("\n--- CAPTURED PROMPT ---\n")
    print(prompt)
    print("\n-----------------------\n")
    
    # Assertions for the focus on the current task
    assert "--- CONVERSATION HISTORY ---" in prompt
    assert "[USER]: hi" in prompt
    assert "==== ACTIVE TASK ====" in prompt
    
    print("\n✅ Verification successful: Memory is clearly separated from the active task.")

def test_hitl_enhanced():
    """Verify that HITL now uses LLM for risk assessment and defaults plan_review to True."""
    # User's test API key
    os.environ["GEMINI_API_KEY"] = "AIzaSyDM2USTaKiMWaYT86D1d6kMwEPgy7MQFPE"
    
    # Initialize with token count to verify "lightweight" claim
    llm = Gemini(model_name="gemini-2.0-flash", token_count=True)
    
    # Create Manager with Medium HITL (triggers LLM risk assessment)
    h_cfg = HitlConfig(permission_level="medium", use_llm=True)
    
    specialist = Agent(name="Specialist", role="System Admin", use_default_tools=True)
    manager = Manager(agents=[specialist], model=llm, hitl=h_cfg, strategy="direct")

    print("\n--- [TEST 1] Risky Task (Should Trigger HITL) ---")
    risky_task = "Delete the server logs and format the disk."
    
    # Mock input to 'y' to simulate approval
    with patch('builtins.input', return_value='y') as mocked_input:
        manager.ask(risky_task, stream=False)
        if mocked_input.called:
            print("✅ HITL correctly triggered for risky task via LLM reasoning.")
        else:
            raise AssertionError("HITL failed to trigger for risky task!")

    print("\n--- [TEST 2] Safe Task (Should NOT Trigger HITL) ---")
    safe_task = "What is 2 + 2?"
    
    # This should NOT call input() because the LLM will say NO to risk
    with patch('builtins.input', side_effect=AssertionError("HITL triggered for safe task!")) as mocked_input:
        res = manager.ask(safe_task, stream=False)
        print(f"Result: {res}")
        print("✅ HITL correctly skipped for safe task via LLM reasoning.")

    llm.print_session_tokens()

if __name__ == "__main__":
    test_memory_prompt_format()
    test_hitl_enhanced()
