from datetime import datetime
from typing import Dict, List, Optional, Any, Union
from uuid import uuid4
import json


class Context:
    def __init__(self, context_id: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, max_history_length: int = 100):
        self.context_id = context_id or str(uuid4())
        self.metadata = metadata or {}
        self.max_history_length = max_history_length
        self.history: List[Dict[str, Any]] = []
        self.state: Dict[str, Any] = {}
        self.created_at = datetime.utcnow()
        self.updated_at = self.created_at
        
    def add_interaction(self, input_data: Dict[str, Any], output_data: Dict[str, Any]) -> None:
        interaction = {
            "timestamp": datetime.utcnow().isoformat(),
            "input": input_data,
            "output": output_data
        }
        
        self.history.append(interaction)
        
        if len(self.history) > self.max_history_length:
            self.history = self.history[-self.max_history_length:]
            
        self.updated_at = datetime.utcnow()
        
    def update_state(self, new_state: Dict[str, Any], merge: bool = True) -> None:
        if merge:
            self.state.update(new_state)
        else:
            self.state = new_state
            
        self.updated_at = datetime.utcnow()
        
    def get_history(self, limit: Optional[int] = None) -> List[Dict[str, Any]]:
        if limit is None or limit >= len(self.history):
            return self.history
        return self.history[-limit:]
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "context_id": self.context_id,
            "metadata": self.metadata,
            "state": self.state,
            "history": self.history,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
        }
    
    def to_json(self) -> str:
        return json.dumps(self.to_dict())
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Context':
        context = cls(context_id=data.get("context_id"))
        context.metadata = data.get("metadata", {})
        context.state = data.get("state", {})
        context.history = data.get("history", [])
        
        if "created_at" in data:
            context.created_at = datetime.fromisoformat(data["created_at"])
        if "updated_at" in data:
            context.updated_at = datetime.fromisoformat(data["updated_at"])
            
        return context
    
    @classmethod
    def from_json(cls, json_str: str) -> 'Context':
        data = json.loads(json_str)
        return cls.from_dict(data)


class ContextManager:
    def __init__(self):
        self.contexts: Dict[str, Context] = {}
        
    def create_context(self, metadata: Optional[Dict[str, Any]] = None, context_id: Optional[str] = None) -> Context:
        context = Context(context_id=context_id, metadata=metadata)
        self.contexts[context.context_id] = context
        return context
    
    def get_context(self, context_id: str) -> Optional[Context]:
        return self.contexts.get(context_id)
    
    def update_context(self, context_id: str, updates: Dict[str, Any]) -> Optional[Context]:
        context = self.get_context(context_id)
        if not context:
            return None
        
        if "metadata" in updates:
            context.metadata.update(updates["metadata"])
            
        if "state" in updates:
            context.update_state(updates["state"])
            
        return context
    
    def delete_context(self, context_id: str) -> bool:
        if context_id in self.contexts:
            del self.contexts[context_id]
            return True
        return False
    
    def list_contexts(self) -> List[str]:
        return list(self.contexts.keys())