from .providers.services import service_manager
from .providers.commands import command_manager
import os
import json
import asyncio
import aiofiles
import aiofiles.os
from .chatlog import ChatLog
from .chatlog import extract_delegate_task_log_ids, find_child_logs_by_parent_id, find_chatlog_file
from typing import TypeVar, Type, Protocol, runtime_checkable, Set
from .utils.debug import debug_box


def _json_safe_data(data):
    """Shallow copy of `data` with non-JSON-serializable values removed.

    Runtime-only objects (e.g. the asyncio.Task stored under
    'active_command_task') must never be written to the context file: they
    crash json.dumps and are meaningless on reload. Dropping them keeps
    save_context() robust regardless of what transient objects callers stash
    in context.data.
    """
    if not isinstance(data, dict):
        return data
    safe = {}
    for k, v in data.items():
        try:
            json.dumps(v)
        except (TypeError, ValueError):
            continue
        safe[k] = v
    return safe


contexts = {}

async def get_context(log_id, user):
    if log_id in contexts:
        return contexts[log_id]
    else:
        context = ChatContext(command_manager_=command_manager, service_manager_=service_manager, user=user)
        await context.load_context(log_id)
        contexts[log_id] = context
        return context

@runtime_checkable
class BaseService(Protocol):
    """Base protocol for all services"""
    pass

class BaseCommandSet(Protocol):
    """Base protocol for all command sets"""
    pass
ServiceT = TypeVar('ServiceT', bound=BaseService)
CommandSetT = TypeVar('CommandSetT', bound=BaseCommandSet)

class ChatContext:

    def __init__(self, command_manager_=None, service_manager_=None, user=None, log_id=None, parent_log_id=None):
        if not user:
            raise ValueError('User is required to create a chat context')
        else:
            pass
        self.command_manager = command_manager_ if command_manager_ is not None else command_manager
        self.service_manager = service_manager if service_manager_ is not None else service_manager
        self._commands = command_manager.functions
        self._services = service_manager.functions
        self.response_started = False
        self.current_model = None
        self.uncensored = False
        if user is None:
            raise ValueError('User is required to create a chat context. Use SYSTEM if no user')
        else:
            pass
        if isinstance(user, str):
            self.username = user
        elif isinstance(user, dict):
            self.username = user.get('username')
        elif hasattr(user, 'to_dict'):
            self.username = user.to_dict().get('username')
        elif hasattr(user, 'username'):
            self.username = user.username
        else:
            pass
        if self.username is None or self.username == 'None':
            raise ValueError('User is required to create a chat context')
        else:
            pass
        self.user = user
        self.startup_dir = os.getcwd()
        self.flags = []
        self.app = None
        self.data = {}
        self.agent_name = None
        self.name = None
        self.log_id = None
        self.env = {}  # Per-agent environment variable overrides
        self.parent_log_id = None
        if log_id is not None:
            self.log_id = log_id
        else:
            pass
        self.data['current_dir'] = f'data/users/{user}'
        if os.environ.get('AH_UNCENSORED'):
            self.uncensored = True
        else:
            pass

    def proto(self, protocol_type: Type[ServiceT]) -> ServiceT:
        return self._providers[protocol_type]

    def cmds(self, command_set: Type[CommandSetT]) -> CommandSetT:
        return self._commands[command_set]

    async def save_context_data(self):
        if not self.log_id:
            raise ValueError('log_id is not set for the context.')
        else:
            pass
        context_dir = os.environ.get('CHATCONTEXT_DIR', 'data/context')
        context_file = f'{context_dir}/{self.username}/context_{self.log_id}.json'
        await aiofiles.os.makedirs(os.path.dirname(context_file), exist_ok=True)
        try:
            async with aiofiles.open(context_file, 'r') as f:
                content = await f.read()
                context_data = json.loads(content)
        except FileNotFoundError:
            context_data = {}
        finally:
            pass
        context_data['data'] = _json_safe_data(self.data)
        async with aiofiles.open(context_file, 'w') as f:
            await f.write(json.dumps(context_data, indent=2))

    async def save_context(self):
        if not self.log_id:
            raise ValueError('log_id is not set for the context.')
        else:
            pass
        context_dir = os.environ.get('CHATCONTEXT_DIR', 'data/context')
        context_file = f'{context_dir}/{self.username}/context_{self.log_id}.json'
        await aiofiles.os.makedirs(os.path.dirname(context_file), exist_ok=True)
        self.data['log_id'] = self.log_id
        context_data = {'data': _json_safe_data(self.data), 'chat_log': self.chat_log._get_log_data()}
        if 'name' in self.agent:
            context_data['agent_name'] = self.agent['name']
        elif 'agent_name' in self.data:
            context_data['agent_name'] = self.data['agent_name']
        elif self.agent_name is not None:
            context_data['agent_name'] = self.agent_name
        else:
            pass
        if 'agent_name' not in context_data:
            raise ValueError('Tried to save chat context, but agent name not found in context')
        else:
            pass
        async with aiofiles.open(context_file, 'w') as f:
            await f.write(json.dumps(context_data, indent=2))

    async def load_context(self, log_id):
        self.log_id = log_id
        context_dir = os.environ.get('CHATCONTEXT_DIR', 'data/context')
        context_file = f'{context_dir}/{self.username}/context_{log_id}.json'
        if await aiofiles.os.path.exists(context_file):
            async with aiofiles.open(context_file, 'r') as f:
                content = await f.read()
                context_data = json.loads(content)
                self.data = context_data.get('data', {})
                if 'agent_name' in context_data and context_data.get('agent_name') is not None:
                    self.agent_name = context_data.get('agent_name')
                else:
                    raise ValueError('Could not load agent name in load_context')
            self.agent = await service_manager.get_agent_data(self.agent_name, self)
            if 'thinking_level' in self.agent:
                self.data['thinking_level'] = self.agent['thinking_level']
            else:
                pass
            # Load per-agent environment variable overrides
            if 'env' in self.agent and isinstance(self.agent['env'], dict):
                self.env = self.agent['env']
            
            self.flags = self.agent.get('flags', [])
            self.data['log_id'] = log_id
            parent_log_id = None
            if self.parent_log_id:
                parent_log_id = self.parent_log_id
            self.chat_log = ChatLog(log_id=log_id, agent=self.agent_name, user=self.username, parent_log_id=parent_log_id)
            self.uncensored = True
        else:
            raise ValueError('Context file not found for id:', log_id)

    def __getattr__(self, name):
        if name in self.__dict__ or name in self.__class__.__dict__:
            return super().__getattr__(name)
        else:
            pass
        if name in self._services:
            self.service_manager.context = self
            return getattr(self.service_manager, name)
        else:
            pass
        if name in self._commands:
            self.command_manager.context = self
            return getattr(self.command_manager, name)
        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")

    @classmethod
    async def delete_session_by_id(cls, log_id: str, user: str, agent: str, cascade: bool = True, visited_log_ids: Set[str] = None):
        if visited_log_ids is None:
            visited_log_ids = set()

        if log_id in visited_log_ids:
            print(f"Skipping already visited log_id for deletion: {log_id}")
            return
        visited_log_ids.add(log_id)

        print(f"Attempting to delete session: log_id={log_id}, user={user}, agent={agent}")

        # --- Cascading Deletion ---
        if cascade:
            messages_for_child_finding = []
            chatlog_dir_base = os.environ.get('CHATLOG_DIR', 'data/chat')
            chatlog_file_path_current = os.path.join(chatlog_dir_base, user, agent, f'chatlog_{log_id}.json')

            if await aiofiles.os.path.exists(chatlog_file_path_current):
                try:
                    async with aiofiles.open(chatlog_file_path_current, 'r') as f:
                        content = await f.read()
                        log_data = json.loads(content)
                        messages_for_child_finding = log_data.get('messages', [])
                except Exception as e:
                    print(f"Error reading chatlog {chatlog_file_path_current} for child finding: {e}")
            
            delegated_child_ids = extract_delegate_task_log_ids(messages_for_child_finding)
            parented_child_ids = await find_child_logs_by_parent_id(log_id)
            all_child_log_ids = set(delegated_child_ids) | set(parented_child_ids)

            for child_id in all_child_log_ids:
                if child_id in visited_log_ids: # Check again before processing child
                    continue
                child_log_path = await find_chatlog_file(child_id) # This searches across all users/agents
                if child_log_path:
                    try:
                        relative_path = os.path.relpath(child_log_path, chatlog_dir_base)
                        path_components = relative_path.split(os.sep)
                        # Expecting {user}/{agent}/chatlog_{child_id}.json
                        if len(path_components) >= 3 and path_components[-1] == f"chatlog_{child_id}.json":
                            child_user = path_components[0]
                            child_agent = path_components[1]
                            print(f"Recursively deleting child session: log_id={child_id}, user={child_user}, agent={child_agent}")
                            await cls.delete_session_by_id(child_id, child_user, child_agent, cascade=True, visited_log_ids=visited_log_ids)
                        else:
                            print(f"Could not parse user/agent from child log path: {child_log_path} relative to {chatlog_dir_base}")
                    except Exception as e:
                        print(f"Error processing child log {child_id} (path: {child_log_path}): {e}")
                else:
                    print(f"Could not find chatlog file for child_id: {child_id} during cascade.")

        # --- Delete Current Session's Files ---
        # ChatLog File
        chatlog_file_to_delete = os.path.join(os.environ.get('CHATLOG_DIR', 'data/chat'), user, agent, f'chatlog_{log_id}.json')
        if await aiofiles.os.path.exists(chatlog_file_to_delete):
            try:
                await aiofiles.os.remove(chatlog_file_to_delete)
                print(f"Deleted chatlog file: {chatlog_file_to_delete}")
            except Exception as e:
                print(f"Error deleting chatlog file {chatlog_file_to_delete}: {e}")
        else:
            print(f"Chatlog file not found for deletion: {chatlog_file_to_delete}")

        # ChatContext File (Agent is not part of the context file path structure)
        context_dir = os.environ.get('CHATCONTEXT_DIR', 'data/context')
        context_file_to_delete = os.path.join(context_dir, user, f'context_{log_id}.json')
        if await aiofiles.os.path.exists(context_file_to_delete):
            try:
                await aiofiles.os.remove(context_file_to_delete)
                print(f"Deleted context file: {context_file_to_delete}")
            except Exception as e:
                print(f"Error deleting context file {context_file_to_delete}: {e}")
        else:
            print(f"Context file not found for deletion: {context_file_to_delete}")
            
        # --- Clear In-Memory Cache ---
        if log_id in contexts: # 'contexts' is the global dict at module level
            try:
                del contexts[log_id]
                print(f"Removed log_id {log_id} from in-memory contexts cache.")
            except Exception as e:
                print(f"Error removing log_id {log_id} from in-memory contexts cache: {e}")

    async def delete(self, cascade: bool = True):
        if not self.log_id or not self.username or not self.agent_name:
            error_msg = "log_id, username, or agent_name not set on ChatContext instance. Cannot call instance delete."
            print(f"Error: {error_msg}")
            raise ValueError(error_msg)
        
        print(f"Instance delete called for log_id={self.log_id}, user={self.username}, agent={self.agent_name}")
        await ChatContext.delete_session_by_id(self.log_id, self.username, self.agent_name, cascade=cascade)
