import json
import queue
import re
import bisect
import sys
from collections import defaultdict, deque
from PyQt5.QtCore import QObject, pyqtSignal, QTimer, pyqtSlot

EMPTY_LIST = []
EMPTY_DICT = {}


class BaseLog:
    __slots__ = ('log_index', 'system_time', 'time_ms', 'class_name')
    
    def __init__(self, index, time_str, ms, cls_name):
        self.log_index = index
        self.system_time = sys.intern(time_str) if time_str else ""
        self.time_ms = ms
        self.class_name = sys.intern(cls_name) if cls_name else ""

    def get(self, key, default=None):
        return getattr(self, key, default)
    
    def __getitem__(self, key):
        return getattr(self, key)


class MessageLog(BaseLog):
    __slots__ = ('sender', 'receiver', 'performative', 'content', 'sender_action')
    
    def __init__(self, index, time_str, ms, sender, receiver, perf, content, action):
        super().__init__(index, time_str, ms, "Channel")
        self.sender = sys.intern(sender) if sender else ""
        self.receiver = receiver
        self.performative = sys.intern(perf) if perf else ""
        self.content = content[:500] if len(content) > 500 else content
        self.sender_action = action[:100] if action and len(action) > 100 else (action or "")


class AgentLog(BaseLog):
    __slots__ = (
        'agent_name', 'cycle', 'beliefs', 'goals', 'running_intentions',
        'perceptions', 'envs', 'chs', 'events', 'desc', 'last_event', 
        'last_intention', '_cached_action'
    )
    
    def __init__(self, index, time_str, ms, name, data):
        super().__init__(index, time_str, ms, "Agent")
        self.agent_name = sys.intern(name) if name else ""
        self.cycle = data.get('cycle', 0)
        
        self.beliefs = self._compact_list(data.get('beliefs'))
        self.goals = self._compact_list(data.get('goals'))
        self.running_intentions = self._compact_list(data.get('running_intentions'))
        self.perceptions = data.get('perceptions') or EMPTY_DICT
        self.envs = self._compact_list(data.get('envs'))
        self.chs = self._compact_list(data.get('chs'))
        self.events = self._compact_list(data.get('events'))
        self.desc = (data.get('desc', '') or "")[:200]
        self.last_event = data.get('last_event', '') or ""
        self.last_intention = data.get('last_intention', '') or ""
        self._cached_action = None
    
    def _compact_list(self, lst):
        if not lst:
            return EMPTY_LIST
        return tuple(lst) if len(lst) <= 50 else tuple(lst[:50])

    @property
    def action(self):
        if self._cached_action is None:
            self._cached_action = self._extract_action()
        return self._cached_action

    def _extract_action(self):
        desc = self.desc
        if not desc:
            return "Idle"
        if "action" in desc:
            match = re.search(r"(?:doing action|action:)\s*\*?(\w+)", desc)
            if match:
                return f"Running: {match.group(1)}"
        if self.running_intentions:
            first = self.running_intentions[0]
            if "->" in first or "]" in first:
                try:
                    match = re.search(r"(?:->|\]\s*,)\s*([\w\d_]+)\(", first)
                    if match:
                        return f"Plan: {match.group(1)}"
                except:
                    pass
        return "Processing"


class IntentionLog(BaseLog):
    __slots__ = ('agent_name', 'intention', 'trigger', 'cycle')
    
    def __init__(self, index, time_str, ms, agent, intention, trigger, cycle):
        super().__init__(index, time_str, ms, "Intention")
        self.agent_name = sys.intern(agent) if agent else ""
        self.intention = intention[:300] if len(intention) > 300 else intention
        self.trigger = sys.intern(trigger) if trigger else ""
        self.cycle = cycle


class EnvironmentLog(BaseLog):
    __slots__ = ('env_name', 'type', 'content', 'agent_action')
    
    def __init__(self, index, time_str, ms, env, log_type, content, action):
        super().__init__(index, time_str, ms, "Environment")
        self.env_name = sys.intern(env) if env else ""
        self.type = sys.intern(log_type) if log_type else ""
        self.content = content[:300] if len(content) > 300 else content
        self.agent_action = action[:100] if action and len(action) > 100 else (action or "")


class LearningLog(BaseLog):
    __slots__ = ('model_name', 'environment', 'agent', 'policy', 'desc', 'training_complete')
    
    def __init__(self, index, time_str, ms, model_name, environment, agent, policy, desc):
        super().__init__(index, time_str, ms, "Learning")
        self.model_name = sys.intern(model_name) if model_name else ""
        self.environment = sys.intern(environment) if environment else ""
        self.agent = sys.intern(agent) if agent else ""
        self.policy = policy or EMPTY_DICT
        self.desc = desc or ""
        self.training_complete = "Training Complete" in (desc or "")


class LogStore(QObject):
    
    store_updated = pyqtSignal()
    environment_state_updated = pyqtSignal(dict)
    agent_list_updated = pyqtSignal(list)
    environment_history_updated = pyqtSignal(str, object)
    participants_updated = pyqtSignal(list)

    def __init__(self, new_queue):
        super().__init__()
        self.new_queue = new_queue

        self.all_logs_timeline = deque()
        self.timeline_ms_map = deque()

        self.message_logs = deque()
        self.msg_timestamps = deque()
        
        self.intention_logs = deque()
        self.int_timestamps = deque()

        self._logs_by_agent = defaultdict(deque)
        
        self._total_messages_processed = 0
        self._total_intentions_processed = 0
        self._total_logs_processed = 0

        self.environment_states = {}
        self._environment_history = defaultdict(deque)
        self._env_history_timestamps = defaultdict(deque)
        self._environment_states_history = defaultdict(deque)
        self._env_states_timestamps = defaultdict(deque)

        self.known_agents = set()
        self._agent_last_intention_tracker = {}
        self.learning_models = {}
        self.learning_logs = deque()
        self.learning_timestamps = deque()
        self.agent_models = defaultdict(set)

        self.total_duration_ms = 0
        self.current_timeline_ms = 0
        self.current_timeline_index = 0
        self.is_live = True
        self._emitted_agent_list = False

        self.processing_timer = QTimer(self)
        self.processing_timer.setInterval(20)
        self.processing_timer.timeout.connect(self._read_and_process_queue)
        
        self.ui_notify_timer = QTimer(self)
        self.ui_notify_timer.setInterval(500)
        self.ui_notify_timer.timeout.connect(self._notify_ui)
        
        self._has_new_data = False

    def start_polling(self):
        self.processing_timer.start()
        self.ui_notify_timer.start()
    
    def stop_polling(self):
        self.processing_timer.stop()
        self.ui_notify_timer.stop()

    
    def _read_and_process_queue(self):
        try:
            packets_limit = 5000
            count = 0
            while count < packets_limit:
                try:
                    log_bundle = self.new_queue.get_nowait()
                    for json_string in log_bundle:
                        try:
                            data = json.loads(json_string)
                            self._process_log_entry(data)
                            self._has_new_data = True
                        except (json.JSONDecodeError, TypeError):
                            pass
                    count += 1
                except queue.Empty:
                    break
        except Exception as e:
            print(f"LogStore Error: {e}")

    def _time_to_ms(self, time_str):
        if not time_str:
            return 0
        try:
            h, m, s_full = time_str.split(':')
            s, ms = s_full.split('.')
            return (int(h) * 3600000) + (int(m) * 60000) + (int(s) * 1000) + int(ms)
        except:
            return self.total_duration_ms

    def _process_log_entry(self, data):
        sys_time = data.get("system_time", "00:00:00.000")
        ms = self._time_to_ms(sys_time)
        class_name = data.get("class_name")
        
        if ms > self.total_duration_ms:
            self.total_duration_ms = ms
            if self.is_live:
                self.current_timeline_ms = ms

        log_added = False
        log_obj = None

        if class_name == "Agent":
            log_obj, log_added = self._process_agent_log(data, sys_time, ms)
            
        elif class_name == "Channel":
            log_obj, log_added = self._process_channel_log(data, sys_time, ms)
            
        elif class_name == "Environment":
            log_obj, log_added = self._process_environment_log(data, sys_time, ms)
            
        elif class_name == "Learning":
            log_obj, log_added = self._process_learning_log(data, sys_time, ms)
        
        else:
            log_index = self._total_logs_processed
            log_obj = BaseLog(log_index, sys_time, ms, class_name or "Unknown")
            log_added = True

        if log_added and log_obj:
            self.all_logs_timeline.append(log_obj)
            self.timeline_ms_map.append(ms)
            self._total_logs_processed += 1
            
        if self.is_live:
            self.current_timeline_index = len(self.timeline_ms_map) - 1

    def _process_agent_log(self, data, sys_time, ms):
        name = data.get("my_name")
        if not name:
            return None, False
            
        should_add = True
        agent_logs = self._logs_by_agent[name]
        
        if agent_logs:
            prev = agent_logs[-1]
            if (data.get('beliefs') == list(prev.beliefs) and
                data.get('goals') == list(prev.goals) and
                data.get('running_intentions') == list(prev.running_intentions) and
                data.get('last_event') == prev.last_event):
                should_add = False

        if should_add:
            log_index = self._total_logs_processed
            log_obj = AgentLog(log_index, sys_time, ms, name, data)
            agent_logs.append(log_obj)
            
            if name not in self.known_agents:
                self.known_agents.add(name)
                self._emitted_agent_list = False

            curr_intention = log_obj.last_intention
            if curr_intention and curr_intention != 'null':
                prev_intention = self._agent_last_intention_tracker.get(name)
                if curr_intention != prev_intention:
                    trig = (log_obj.last_event or "").replace('gain:', '').replace('lose:', '')
                    int_log = IntentionLog(
                        self._total_intentions_processed, 
                        sys_time, ms, name, curr_intention, trig, log_obj.cycle
                    )
                    self.intention_logs.append(int_log)
                    self.int_timestamps.append(ms)
                    self._total_intentions_processed += 1
                    self._agent_last_intention_tracker[name] = curr_intention

            desc = data.get("desc", "")
            if "Adding model for" in desc:
                match = re.search(r"Adding model for\s+(\S+)", desc)
                if match:
                    self.agent_models[name].add(match.group(1))
            
            return log_obj, True
        
        return None, False

    def _process_channel_log(self, data, sys_time, ms):
        desc = data.get("desc", "")
        
        if "sending" in desc:
            match = re.search(r"(\w+)\s+sending\s+([^:]+):\s*(.*)\s+to\s+([\w\s,\[\]']+)", desc)
            if match:
                sender, perf, content, receiver = match.groups()

                sender_logs = self._logs_by_agent.get(sender)
                action = sender_logs[-1].action if sender_logs else "Unknown"
                
                log_obj = MessageLog(
                    self._total_messages_processed,
                    sys_time, ms, sender, receiver.strip(), 
                    perf.strip(), content.strip(), action
                )
                self.message_logs.append(log_obj)
                self.msg_timestamps.append(ms)
                self._total_messages_processed += 1
                
                return log_obj, True
        
        return None, False

    def _process_environment_log(self, data, sys_time, ms):
        env_name = data.get("my_name")
        if not env_name:
            return None, False
            
        desc = data.get("desc", "")
        log_added = False
        log_obj = None

        if desc.startswith(("Creating", "Deleting", "Changing")):
            action = f"{data.get('action')} ({data.get('agent')})"
            p_new = data.get('new_percept') or data.get('percept(s)')
            content = str(p_new)
            
            if "Changing" in desc:
                content = f"{str(data.get('old_percept'))} -> {content}"

            first_word = desc.split()[0]
            l_type = ("create" if "Creating" in first_word else 
                     "delete" if "Deleting" in first_word else 
                     "change" if "Changing" in first_word else 
                     first_word.lower())

            log_obj = EnvironmentLog(
                self._total_logs_processed, sys_time, ms, 
                env_name, l_type, content, action
            )
            self._environment_history[env_name].append(log_obj)
            self._env_history_timestamps[env_name].append(ms)
            
            if self.is_live:
                self.environment_history_updated.emit(env_name, log_obj)
            log_added = True

        percepts = data.get("percepts")
        if percepts:
            clean = {k: str(v) for k, v in (percepts[0] if isinstance(percepts, list) else percepts).items()}
            state_snapshot = {
                'percepts': clean,
                'connected_agents': data.get('connected_agents', [])
            }

            env_states = self._environment_states_history[env_name]
            should_save = True
            if env_states and env_states[-1] == state_snapshot:
                should_save = False
            
            if should_save:
                self.environment_states[env_name] = state_snapshot
                env_states.append(state_snapshot)
                self._env_states_timestamps[env_name].append(ms)
                
                if self.is_live:
                    self.environment_state_updated.emit(self.environment_states)

        return log_obj, log_added

    def _process_learning_log(self, data, sys_time, ms):
        model_name = data.get("my_name")
        log_obj = LearningLog(
            len(self.learning_logs), sys_time, ms,
            model_name, data.get("environment"),
            data.get("agent"), data.get("policy"), data.get("desc")
        )
        
        self.learning_logs.append(log_obj)
        self.learning_timestamps.append(ms)
        
        if model_name not in self.learning_models:
            self.learning_models[model_name] = {
                'environment': log_obj.environment,
                'policy': {},
                'agents': set(),
                'training_complete': False
            }
        
        if log_obj.policy:
            self.learning_models[model_name]['policy'] = log_obj.policy
        if log_obj.training_complete:
            self.learning_models[model_name]['training_complete'] = True
        
        return log_obj, True
    
    def get_message_count(self):
        return self._total_messages_processed
    
    def get_intention_count(self):
        return self._total_intentions_processed
    
    def get_messages_in_window(self):
        return len(self.message_logs)
    
    def get_intentions_in_window(self):
        return len(self.intention_logs)

    def get_messages_reference(self):
        return self.message_logs

    def get_message_count_limit(self):
        if self.is_live:
            return len(self.message_logs)
        return self._get_slice_index(list(self.msg_timestamps), self.current_timeline_ms)

    def get_all_messages(self):
        if self.is_live:
            return self.message_logs
        cut_idx = self._get_slice_index(list(self.msg_timestamps), self.current_timeline_ms)
        return list(self.message_logs)[:cut_idx]

    def get_all_intentions_history(self):
        if self.is_live:
            return self.intention_logs
        cut_idx = self._get_slice_index(list(self.int_timestamps), self.current_timeline_ms)
        return list(self.intention_logs)[:cut_idx]

    def get_logs_for_agent(self, agent_name):
        logs = self._logs_by_agent.get(agent_name)
        if not logs:
            return []
        if self.is_live:
            return list(logs)
        
        logs_list = list(logs)
        target = self.current_timeline_ms
        lo, hi = 0, len(logs_list)
        while lo < hi:
            mid = (lo + hi) // 2
            if logs_list[mid].time_ms <= target:
                lo = mid + 1
            else:
                hi = mid
        return logs_list[:lo]
    
    def get_agent_log_history(self, agent_name):
        return self.get_logs_for_agent(agent_name)

    def get_latest_agent_state(self, agent_name):
        logs = self._logs_by_agent.get(agent_name)
        return logs[-1] if logs else None

    def get_latest_agent_state_before_index(self, agent_name, log_index):
        logs = self.get_logs_for_agent(agent_name)
        return logs[-1] if logs else None

    def get_environment_states_at_index(self, index):
        if self.is_live:
            return self.environment_states
        
        result = {}
        target_ms = self.current_timeline_ms
        for env_name in self._environment_states_history.keys():
            timestamps = list(self._env_states_timestamps.get(env_name, []))
            states = list(self._environment_states_history.get(env_name, []))
            if not timestamps:
                continue
            idx = bisect.bisect_right(timestamps, target_ms)
            if idx > 0:
                result[env_name] = states[idx - 1]
        return result
    
    def get_latest_environment_states(self):
        return self.environment_states

    def get_environment_change_history(self, env_name):
        hist = self._environment_history.get(env_name)
        if not hist:
            return []
        if self.is_live:
            return list(hist)
        ts_list = list(self._env_history_timestamps.get(env_name, []))
        cut_idx = self._get_slice_index(ts_list, self.current_timeline_ms)
        return list(hist)[:cut_idx]

    def get_learning_models(self):
        return {
            k: {
                'environment': v['environment'],
                'policy': v['policy'],
                'training_complete': v['training_complete']
            }
            for k, v in self.learning_models.items()
        }
    
    def get_learning_for_agent(self, agent_name):
        model_names = self.agent_models.get(agent_name, set())
        return {
            m: {'policy': self.learning_models[m]['policy']}
            for m in model_names if m in self.learning_models
        }
    
    def get_learning_logs(self):
        if self.is_live:
            return list(self.learning_logs)
        cut_idx = self._get_slice_index(
            list(self.learning_timestamps), 
            self.current_timeline_ms
        )
        return list(self.learning_logs)[:cut_idx]
    
    def _get_slice_index(self, timestamp_list, target_ms):
        if not timestamp_list:
            return 0
        return bisect.bisect_right(timestamp_list, target_ms)

    def _notify_ui(self):
        if self._has_new_data and self.is_live:
            self._has_new_data = False
            self.store_updated.emit()
            if not self._emitted_agent_list and self.known_agents:
                self._emit_agent_list()

    def _emit_agent_list(self):
        if not self.known_agents:
            return
        filtered = [a for a in self.known_agents if f"{a}_1" not in self.known_agents]
        self.agent_list_updated.emit(sorted(filtered))
        self._emitted_agent_list = True
    
    @pyqtSlot(int)
    def set_current_timeline_index(self, index):
        timeline_len = len(self.timeline_ms_map)
        if 0 <= index < timeline_len:
            self.current_timeline_ms = self.timeline_ms_map[index]
            self.current_timeline_index = index
        else:
            self.current_timeline_ms = self.total_duration_ms
            self.current_timeline_index = timeline_len - 1
        
        self.is_live = (index >= timeline_len - 50)
        self.store_updated.emit()

    @pyqtSlot()
    def toggle_live_mode(self):
        self.is_live = not self.is_live
        if self.is_live:
            self.current_timeline_ms = self.total_duration_ms
            self.current_timeline_index = len(self.timeline_ms_map) - 1
            self.store_updated.emit()

    def get_total_duration_ms(self):
        return self.total_duration_ms

    def get_index_from_ms(self, ms_value):
        if not self.timeline_ms_map:
            return 0
        timestamps = list(self.timeline_ms_map)
        idx = bisect.bisect_left(timestamps, ms_value)
        return min(idx, len(timestamps) - 1)
    
    def clear_old_data(self):
        import gc
        gc.collect()
        
        print(f"[LogStore] Memory cleanup - Logs: {len(self.all_logs_timeline)}, "
              f"Messages: {len(self.message_logs)}, "
              f"Intentions: {len(self.intention_logs)}")

    def get_memory_stats(self):
        return {
            'total_logs': len(self.all_logs_timeline),
            'total_messages': len(self.message_logs),
            'total_intentions': len(self.intention_logs),
            'total_messages_processed': self._total_messages_processed,
            'total_intentions_processed': self._total_intentions_processed,
            'agents_count': len(self.known_agents)
        }