"""
Context management for PicoGent
"""

import os
from typing import Dict, Any, List, Optional


class Context:
    """Context manager for agent workspace and environment"""
    
    def __init__(self, workspace: str = "."):
        self.workspace = os.path.abspath(workspace)
        self.context_data: Dict[str, Any] = {}
    
    def set_workspace(self, workspace: str):
        """Set the workspace directory"""
        self.workspace = os.path.abspath(workspace)
        if not os.path.exists(self.workspace):
            os.makedirs(self.workspace, exist_ok=True)
    
    def get_workspace(self) -> str:
        """Get the current workspace directory"""
        return self.workspace
    
    def resolve_path(self, path: str) -> str:
        """Resolve a path relative to workspace"""
        if os.path.isabs(path):
            return path
        return os.path.join(self.workspace, path)
    
    def set_context(self, key: str, value: Any):
        """Set context data"""
        self.context_data[key] = value
    
    def get_context(self, key: str, default: Any = None) -> Any:
        """Get context data"""
        return self.context_data.get(key, default)
    
    def clear_context(self):
        """Clear all context data"""
        self.context_data.clear()
    
    def get_system_prompt(self) -> str:
        """Generate system prompt with context"""
        prompt_parts = [
            "You are PicoGent, a helpful coding assistant.",
            f"Your workspace is: {self.workspace}",
            "",
            "Available tools:",
            "- read: Read file contents with optional offset/limit",
            "- write: Write content to files (creates directories as needed)", 
            "- edit: Replace text in files (old_string -> new_string)",
            "- bash: Execute shell commands (30s timeout)",
            "",
            "Guidelines:",
            "1. Always use relative paths when possible",
            "2. Be precise with file operations",
            "3. Ask for confirmation before destructive operations",
            "4. Provide helpful error messages",
            "5. Keep responses concise and actionable"
        ]
        
        # Add custom context if available
        custom_context = self.get_context("system_additions")
        if custom_context:
            prompt_parts.extend(["", "Additional context:", custom_context])
        
        return "\n".join(prompt_parts)