import unittest
import tempfile
import os
import shutil
from datetime import datetime, timedelta
from src.log_registry import LogRegistry
from src.log_pruner import LogPruner

class TestLogPruningHeuristic(unittest.TestCase):
 def setUp(self) -> None:
  self.temp_dir = tempfile.TemporaryDirectory()
  self.registry_path = os.path.join(self.temp_dir.name, "registry.toml")
  self.logs_dir = os.path.join(self.temp_dir.name, "logs")
  os.makedirs(self.logs_dir, exist_ok=True)
  self.registry = LogRegistry(self.registry_path)
  self.pruner = LogPruner(self.registry, self.logs_dir)

 def tearDown(self) -> None:
  self.temp_dir.cleanup()

 def _create_session(self, session_id, start_time, message_count=None, size_kb=None, whitelisted=False):
  path = os.path.join(self.logs_dir, session_id)
  os.makedirs(path, exist_ok=True)
  # Create a dummy file
  with open(os.path.join(path, "comms.log"), "w") as f:
   f.write("test content")
  
  self.registry.register_session(session_id, path, start_time)
  if message_count is not None or size_kb is not None or whitelisted:
   self.registry.update_session_metadata(
    session_id,
    message_count=message_count if message_count is not None else 10,
    errors=0,
    size_kb=size_kb if size_kb is not None else 10,
    whitelisted=whitelisted,
    reason="Test"
   )

 def test_get_old_non_whitelisted_sessions_includes_empty_sessions(self) -> None:
  now = datetime.now()
  cutoff_time = now - timedelta(days=7)
  
  # 1. Old, not whitelisted (should be included)
  self._create_session("old_nw", now - timedelta(days=10))
  
  # 2. Recent, not whitelisted, but empty (message_count=0) (SHOULD be included based on new heuristic)
  self._create_session("recent_empty_msgs", now - timedelta(days=1), message_count=0)
  
  # 3. Recent, not whitelisted, but empty (size_kb=0) (SHOULD be included based on new heuristic)
  self._create_session("recent_empty_size", now - timedelta(days=1), size_kb=0)
  
  # 4. Recent, not whitelisted, NOT empty (should NOT be included)
  self._create_session("recent_not_empty", now - timedelta(days=1), message_count=5, size_kb=5)
  
  # 5. Old, whitelisted (should NOT be included)
  self._create_session("old_w", now - timedelta(days=10), whitelisted=True)

  sessions = self.registry.get_old_non_whitelisted_sessions(cutoff_time)
  session_ids = {s['session_id'] for s in sessions}
  
  self.assertIn("old_nw", session_ids)
  self.assertIn("recent_empty_msgs", session_ids)
  self.assertIn("recent_empty_size", session_ids)
  self.assertNotIn("recent_not_empty", session_ids)
  self.assertNotIn("old_w", session_ids)

 def test_prune_removes_empty_sessions_regardless_of_age(self) -> None:
  now = datetime.now()
  
  # Create a session that is recent but empty
  session_id = "recent_empty"
  session_path = os.path.join(self.logs_dir, session_id)
  os.makedirs(session_path, exist_ok=True)
  # Actual file size 0
  with open(os.path.join(session_path, "comms.log"), "w") as f:
   pass
   
  self.registry.register_session(session_id, session_path, now - timedelta(hours=1))
  self.registry.update_session_metadata(session_id, message_count=0, errors=0, size_kb=0, whitelisted=False, reason="Empty")
  
  self.assertTrue(os.path.exists(session_path))
  
  # Prune with max_age_days=30 (so 1 hour old is NOT "old" by age)
  self.pruner.prune(max_age_days=30, min_size_kb=1)
  
  self.assertFalse(os.path.exists(session_path))
  self.assertNotIn(session_id, self.registry.data)

if __name__ == '__main__':
 unittest.main()