# createsonline/agents/base.py
"""
Core Agent and SubAgent classes for CREATESONLINE.

Provides a flexible agent framework that works with:
- Local ML models (TransformerLM from createsonline.ml)
- External LLM APIs (OpenAI, Anthropic via createsonline.ai)
- Custom inference functions

Agents support tools, memory, sub-agents, and ReAct-style reasoning.
"""

import json
import time
import asyncio
import traceback
from typing import Any, Callable, Dict, List, Optional, Union

from .tools import ToolRegistry, ToolResult
from .memory import ConversationMemory, MemoryManager, Message
from .prompts import ReActPrompt


class AgentResponse:
    """Structured response from an agent."""
    
    __slots__ = (
        'content', 'tool_calls', 'thoughts', 'metadata',
        'agent_name', 'duration', 'success', 'error',
    )
    
    def __init__(self, content: str = '', success: bool = True, error: str = None,
                 agent_name: str = '', duration: float = 0.0):
        self.content = content
        self.success = success
        self.error = error
        self.agent_name = agent_name
        self.duration = duration
        self.tool_calls: List[ToolResult] = []
        self.thoughts: List[str] = []
        self.metadata: dict = {}
    
    def to_dict(self) -> dict:
        return {
            'content': self.content,
            'success': self.success,
            'error': self.error,
            'agent_name': self.agent_name,
            'duration': self.duration,
            'tool_calls': [tc.to_dict() for tc in self.tool_calls],
            'thoughts': self.thoughts,
            'metadata': self.metadata,
        }
    
    def __str__(self):
        return self.content
    
    def __repr__(self):
        status = 'OK' if self.success else 'FAIL'
        return f"AgentResponse({status}, agent={self.agent_name!r}, len={len(self.content)})"


class Agent:
    """
    An intelligent agent that can reason, use tools, and delegate to sub-agents.
    
    The agent supports multiple inference backends:
    
    1. **Local ML**: Use a TransformerLM model trained with createsonline.ml.
    2. **External LLM**: Use OpenAI/Anthropic via createsonline.ai services.
    3. **Custom**: Provide any async callable as the inference function.
    4. **Rule-based**: Override the `think()` method for deterministic logic.
    
    Usage::
    
        # With external LLM
        from createsonline.ai import ai_config
        agent = Agent(
            name="assistant",
            system_prompt="You are a helpful coding assistant.",
            inference_fn=lambda messages: ai_config.get_service("openai").chat(messages),
        )
        
        # With tools
        tools = ToolRegistry()
        tools.register("search", search_fn, "Search the web")
        agent = Agent(name="searcher", tools=tools)
        
        # Run
        response = await agent.run("Find information about Python asyncio")
    
    Architecture:
        The agent follows a Perceive → Think → Act loop:
        1. Perceive: Add user input to memory
        2. Think: Generate response (via LLM or rules)
        3. Act: Execute tool calls if needed
        4. Observe: Feed tool results back into context
        5. Repeat until final answer or max iterations
    """
    
    def __init__(
        self,
        name: str = 'agent',
        system_prompt: str = '',
        inference_fn: Callable = None,
        tools: ToolRegistry = None,
        memory: Union[ConversationMemory, MemoryManager] = None,
        sub_agents: Dict[str, 'Agent'] = None,
        max_iterations: int = 10,
        verbose: bool = False,
        metadata: dict = None,
    ):
        self.name = name
        self.system_prompt = system_prompt or f"You are {name}, a helpful AI assistant."
        self.inference_fn = inference_fn
        self.tools = tools or ToolRegistry()
        self.max_iterations = max_iterations
        self.verbose = verbose
        self.metadata = metadata or {}
        
        # Memory
        if isinstance(memory, MemoryManager):
            self.memory_manager = memory
            self.memory = memory.conversation
        elif isinstance(memory, ConversationMemory):
            self.memory = memory
            self.memory_manager = None
        else:
            self.memory = ConversationMemory(max_turns=50, max_tokens=8000)
            self.memory_manager = MemoryManager()
        
        # Sub-agents
        self.sub_agents: Dict[str, Agent] = sub_agents or {}
        
        # ReAct prompt builder
        self._react = ReActPrompt(system=self.system_prompt)
        
        # Hooks
        self._pre_hooks: List[Callable] = []
        self._post_hooks: List[Callable] = []
        
        # State
        self._iteration_count = 0
        self._total_runs = 0
        self._is_running = False
        
        # Initialize system message in memory
        self.memory.add_system(self.system_prompt)
    
    # ------------------------------------------------------------------
    # Sub-Agent management
    # ------------------------------------------------------------------
    def add_sub_agent(self, agent: 'Agent', name: str = None):
        """Register a sub-agent that this agent can delegate to."""
        key = name or agent.name
        self.sub_agents[key] = agent
        
        # Also register sub-agent as a tool
        async def _delegate(task: str) -> str:
            resp = await agent.run(task)
            return resp.content
        
        self.tools.register(
            f"delegate_{key}",
            _delegate,
            f"Delegate a task to the {key} sub-agent: {agent.system_prompt[:100]}",
            parameters={'task': {'type': 'string', 'description': 'The task to delegate'}},
            required_params=['task'],
            tags=['sub_agent'],
        )
    
    def remove_sub_agent(self, name: str):
        self.sub_agents.pop(name, None)
        self.tools.unregister(f"delegate_{name}")
    
    # ------------------------------------------------------------------
    # Hooks
    # ------------------------------------------------------------------
    def before_run(self, fn: Callable):
        """Register a pre-run hook."""
        self._pre_hooks.append(fn)
        return fn
    
    def after_run(self, fn: Callable):
        """Register a post-run hook."""
        self._post_hooks.append(fn)
        return fn
    
    # ------------------------------------------------------------------
    # Core inference
    # ------------------------------------------------------------------
    async def _infer(self, messages: List[dict]) -> str:
        """
        Run inference to generate a response.
        
        Override this method to use a custom inference backend.
        Falls back to built-in rule-based responses if no inference_fn provided.
        """
        if self.inference_fn is not None:
            result = self.inference_fn(messages)
            if asyncio.iscoroutine(result):
                result = await result
            if isinstance(result, dict):
                return result.get('content', result.get('text', str(result)))
            return str(result)
        
        # Built-in: simple rule-based fallback
        return self._rule_based_response(messages)
    
    def _rule_based_response(self, messages: List[dict]) -> str:
        """
        Simple rule-based response when no LLM is available.
        
        This processes the last user message and:
        1. Checks if a tool call is needed (looks for keywords)
        2. Returns a direct response otherwise
        """
        last_user_msg = ''
        for m in reversed(messages):
            if m.get('role') == 'user':
                last_user_msg = m.get('content', '')
                break
        
        if not last_user_msg:
            return "Final Answer: I'm ready to help. What would you like me to do?"
        
        # Check if any tool keywords match
        msg_lower = last_user_msg.lower()
        for tool in self.tools.list_tools():
            # Simple keyword matching
            tool_words = tool.name.replace('_', ' ').lower().split()
            if any(w in msg_lower for w in tool_words if len(w) > 3):
                params = {}
                # Try to extract the main parameter
                for pname in tool.required_params:
                    params[pname] = last_user_msg
                    break
                return (
                    f"Thought: I should use the {tool.name} tool\n"
                    f"Action: {tool.name}\n"
                    f"Action Input: {json.dumps(params)}"
                )
        
        return f"Final Answer: {last_user_msg}"
    
    # ------------------------------------------------------------------
    # Main run loop
    # ------------------------------------------------------------------
    async def run(self, user_input: str, context: dict = None) -> AgentResponse:
        """
        Run the agent on a user input.
        
        This implements the full Perceive → Think → Act → Observe loop.
        
        Args:
            user_input: The user's message/query.
            context: Optional additional context dict.
        
        Returns:
            AgentResponse with the final answer and metadata.
        """
        t0 = time.time()
        self._is_running = True
        self._total_runs += 1
        
        response = AgentResponse(agent_name=self.name)
        
        try:
            # Pre-hooks
            for hook in self._pre_hooks:
                result = hook(user_input, context)
                if asyncio.iscoroutine(result):
                    result = await result
                if isinstance(result, str):
                    user_input = result
            
            # Add user message to memory
            meta = {'context': context} if context else None
            self.memory.add_user(user_input, meta)
            if self.memory_manager:
                self.memory_manager.add_message('user', user_input, meta)
            
            # RAG: Retrieve relevant context if memory manager is available
            rag_context = ''
            if self.memory_manager:
                rag_context = self.memory_manager.build_rag_context(user_input)
            
            # Build system prompt with tools and RAG
            system = self.system_prompt
            if rag_context:
                system += '\n\n' + rag_context
            if len(self.tools) > 0:
                system = self._react.build_system_prompt(
                    self.tools.to_prompt_description()
                )
                if rag_context:
                    system += '\n\n' + rag_context
            
            # ReAct loop
            for iteration in range(self.max_iterations):
                self._iteration_count = iteration + 1
                
                # Build messages for inference
                messages = self.memory.get_messages()
                if messages and messages[0].get('role') == 'system':
                    messages[0]['content'] = system
                
                # Infer
                raw_response = await self._infer(messages)
                
                if self.verbose:
                    print(f"[{self.name}] Iteration {iteration + 1}: {raw_response[:200]}")
                
                # Parse ReAct response
                parsed = ReActPrompt.parse_response(raw_response)
                
                if parsed['thought']:
                    response.thoughts.append(parsed['thought'])
                
                # Check for final answer
                if parsed['final_answer'] is not None:
                    response.content = parsed['final_answer']
                    break
                
                # Execute tool if requested
                if parsed['action']:
                    tool_name = parsed['action']
                    tool_input = parsed['action_input'] or {}
                    
                    if self.verbose:
                        print(f"[{self.name}] Tool call: {tool_name}({tool_input})")
                    
                    if isinstance(tool_input, str):
                        tool_input = {'input': tool_input}
                    
                    tool_result = await self.tools.invoke(tool_name, **tool_input)
                    response.tool_calls.append(tool_result)
                    
                    # Feed observation back into memory
                    observation = (
                        str(tool_result.output) if tool_result.success
                        else f"Error: {tool_result.error}"
                    )
                    
                    react_text = (
                        f"Thought: {parsed['thought'] or 'Processing...'}\n"
                        f"Action: {tool_name}\n"
                        f"Action Input: {json.dumps(tool_input)}\n"
                        f"Observation: {observation}"
                    )
                    self.memory.add_assistant(react_text)
                    
                    continue
                
                # No action and no final answer — treat raw response as final
                if not parsed['action'] and not parsed['final_answer']:
                    response.content = raw_response
                    break
            
            else:
                # Max iterations reached
                response.content = (
                    response.thoughts[-1] if response.thoughts
                    else "I was unable to complete the task within the iteration limit."
                )
            
            # Add final response to memory
            self.memory.add_assistant(response.content)
            if self.memory_manager:
                self.memory_manager.add_message('assistant', response.content)
            
            # Post-hooks
            for hook in self._post_hooks:
                result = hook(response)
                if asyncio.iscoroutine(result):
                    await result
        
        except Exception as exc:
            response.success = False
            response.error = f"{type(exc).__name__}: {exc}"
            response.metadata['traceback'] = traceback.format_exc()
            if self.verbose:
                print(f"[{self.name}] Error: {response.error}")
        
        finally:
            response.duration = time.time() - t0
            self._is_running = False
        
        return response
    
    def run_sync(self, user_input: str, context: dict = None) -> AgentResponse:
        """Synchronous wrapper for run()."""
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None
        
        if loop and loop.is_running():
            import concurrent.futures
            with concurrent.futures.ThreadPoolExecutor() as pool:
                future = pool.submit(asyncio.run, self.run(user_input, context))
                return future.result()
        else:
            return asyncio.run(self.run(user_input, context))
    
    # ------------------------------------------------------------------
    # Convenience methods
    # ------------------------------------------------------------------
    async def chat(self, message: str) -> str:
        """Simple chat interface — returns just the text response."""
        resp = await self.run(message)
        return resp.content
    
    def reset(self):
        """Reset the agent's memory and state."""
        self.memory.clear()
        self.memory.add_system(self.system_prompt)
        if self.memory_manager:
            self.memory_manager.clear()
        self._iteration_count = 0
    
    def save_state(self, directory: str):
        """Save agent state to a directory."""
        import os
        os.makedirs(directory, exist_ok=True)
        
        # Save memory
        if self.memory_manager:
            self.memory_manager.save(os.path.join(directory, 'memory'))
        else:
            self.memory.save(os.path.join(directory, 'conversation.json'))
        
        # Save agent config
        config = {
            'name': self.name,
            'system_prompt': self.system_prompt,
            'max_iterations': self.max_iterations,
            'verbose': self.verbose,
            'metadata': self.metadata,
            'total_runs': self._total_runs,
            'sub_agents': list(self.sub_agents.keys()),
        }
        with open(os.path.join(directory, 'config.json'), 'w') as f:
            json.dump(config, f, indent=2, default=str)
    
    def load_state(self, directory: str):
        """Load agent state from a directory."""
        import os
        
        # Load memory
        mem_dir = os.path.join(directory, 'memory')
        if os.path.exists(mem_dir) and self.memory_manager:
            self.memory_manager.load(mem_dir)
        
        conv_path = os.path.join(directory, 'conversation.json')
        if os.path.exists(conv_path):
            self.memory.load(conv_path)
        
        # Load config
        config_path = os.path.join(directory, 'config.json')
        if os.path.exists(config_path):
            with open(config_path, 'r') as f:
                config = json.load(f)
            self._total_runs = config.get('total_runs', 0)
            self.metadata.update(config.get('metadata', {}))
    
    @property
    def is_running(self) -> bool:
        return self._is_running
    
    @property
    def stats(self) -> dict:
        return {
            'name': self.name,
            'total_runs': self._total_runs,
            'memory_turns': self.memory.turn_count,
            'tools_count': len(self.tools),
            'sub_agents': list(self.sub_agents.keys()),
            'is_running': self._is_running,
        }
    
    def __repr__(self):
        return (
            f"Agent({self.name!r}, tools={len(self.tools)}, "
            f"sub_agents={len(self.sub_agents)})"
        )


class SubAgent(Agent):
    """
    A specialized agent designed for delegation from a parent agent.
    
    SubAgents have:
    - A parent reference for reporting back
    - A specialized focus/role
    - Shared or independent memory
    - Automatic delegation protocol
    
    Usage::
    
        # Create specialized sub-agents
        researcher = SubAgent(
            name="researcher",
            system_prompt="You research topics thoroughly and return factual summaries.",
            tools=research_tools,
        )
        
        coder = SubAgent(
            name="coder",
            system_prompt="You write clean, tested Python code.",
            tools=coding_tools,
        )
        
        # Attach to parent agent
        main_agent = Agent(name="orchestrator")
        main_agent.add_sub_agent(researcher)
        main_agent.add_sub_agent(coder)
    """
    
    def __init__(self, name: str = 'sub_agent', system_prompt: str = '',
                 specialization: str = '', parent: Agent = None, **kwargs):
        super().__init__(name=name, system_prompt=system_prompt, **kwargs)
        self.specialization = specialization
        self.parent = parent
        self._task_history: List[Dict[str, Any]] = []
    
    async def run(self, user_input: str, context: dict = None) -> AgentResponse:
        """Run with task history tracking."""
        task_record = {
            'input': user_input,
            'context': context,
            'start_time': time.time(),
        }
        
        response = await super().run(user_input, context)
        
        task_record['output'] = response.content
        task_record['success'] = response.success
        task_record['duration'] = response.duration
        self._task_history.append(task_record)
        
        return response
    
    async def report_to_parent(self, result: str):
        """Report a result back to the parent agent."""
        if self.parent:
            self.parent.memory.add_tool(
                result, tool_name=f"sub_agent_{self.name}"
            )
    
    @property
    def task_history(self) -> List[dict]:
        return list(self._task_history)
    
    @property
    def success_rate(self) -> float:
        if not self._task_history:
            return 0.0
        successes = sum(1 for t in self._task_history if t.get('success'))
        return successes / len(self._task_history)
    
    def __repr__(self):
        return (
            f"SubAgent({self.name!r}, spec={self.specialization!r}, "
            f"tasks={len(self._task_history)})"
        )
