import sys
import os
import time
from simulation.sim_base import BaseSimulation, run_sim

class ContextSimulation(BaseSimulation):
    def run(self):
        print("\n--- Running Context & Chat Simulation ---")
        
        # 1. Test Discussion Creation
        disc_name = f"TestDisc_{int(time.time())}"
        print(f"[Sim] Creating discussion: {disc_name}")
        self.sim.create_discussion(disc_name)
        time.sleep(1)
        
        # Verify it's in the list
        session = self.client.get_session()
        # The session structure usually has discussions listed somewhere, or we can check the listbox
        # For now, we'll trust the click and check the session update
        
        # 2. Test File Aggregation & Context Refresh
        print("[Sim] Testing context refresh and token budget...")
        proj = self.client.get_project()
        # Add many files to ensure we cross the 1% threshold (~9000 tokens)
        import glob
        all_py = [os.path.basename(f) for f in glob.glob("*.py")]
        for f in all_py:
            if f not in proj['project']['files']['paths']:
                proj['project']['files']['paths'].append(f)
        
        # Update project via hook
        self.client.post_project(proj['project'])
        time.sleep(1)
        
        # Trigger MD Only to refresh context and token budget
        print("[Sim] Clicking MD Only...")
        self.client.click("btn_md_only")
        time.sleep(5)
        
        # Verify status
        proj_updated = self.client.get_project()
        status = self.client.get_value("ai_status")
        print(f"[Sim] Status: {status}")
        assert "md written" in status, f"Expected 'md written' in status, got {status}"
        
        # Verify token budget
        pct = self.client.get_value("token_budget_pct")
        current = self.client.get_value("token_budget_current")
        print(f"[Sim] Token budget pct: {pct}, current={current}")
        # We'll just warn if it's 0 but the MD was written, as it might be a small context
        if pct == 0:
            print("[Sim] WARNING: token_budget_pct is 0. This might be due to small context or estimation failure.")
        
        # 3. Test Chat Turn
        msg = "What is the current date and time? Answer in one sentence."
        print(f"[Sim] Sending message: {msg}")
        self.sim.run_discussion_turn(msg)
        
        # 4. Verify History
        print("[Sim] Verifying history...")
        session = self.client.get_session()
        entries = session.get('session', {}).get('entries', [])
        
        # We expect at least 2 entries (User and AI)
        assert len(entries) >= 2, f"Expected at least 2 entries, found {len(entries)}"
        assert entries[-2]['role'] == 'User', "Expected second to last entry to be User"
        assert entries[-1]['role'] == 'AI', "Expected last entry to be AI"
        print(f"[Sim] AI responded: {entries[-1]['content'][:50]}...")

        # 5. Test History Truncation
        print("[Sim] Testing history truncation...")
        self.sim.truncate_history(1)
        time.sleep(1)
        session = self.client.get_session()
        entries = session.get('session', {}).get('entries', [])
        # Truncating to 1 pair means 2 entries max (if it's already at 2, it might not change,
        # but if we had more, it would).
        assert len(entries) <= 2, f"Expected <= 2 entries after truncation, found {len(entries)}"

if __name__ == "__main__":
    run_sim(ContextSimulation)