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

class ContextSimulation(BaseSimulation):
 def run(self) -> None:
  print("\n--- Running Context & Chat Simulation ---")
  # 1. Skip Discussion Creation, use 'main'
  print("[Sim] Using existing 'main' discussion")
  self.sim.switch_discussion("main")
  time.sleep(1.5)
  # 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
  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', [])
  if not entries:
   print("[Sim] !!! WARNING: entries list is EMPTY. Waiting another 2 seconds for eventual consistency...")
   time.sleep(2)
   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', [])
  print(f"[DEBUG] Entries after truncation: {entries}")
  chat_entries = [e for e in entries if e.get('role') in ('User', 'AI')]
  assert len(chat_entries) == 2, f"Expected exactly 2 chat entries after truncation, found {len(chat_entries)}"

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