"""
Base agent class for AgentRouter SDK
"""

import base64
import asyncio
import json
import threading
import uuid
from abc import ABC, abstractmethod
from datetime import datetime
from typing import Any, Dict, List, Optional, Set

from agentrouter.api import PlanAPIClient, ToolCallAPIClient, UtilityAPIClient, get_utility_client
from agentrouter.exceptions import (
    AgentRouterError,
    ConfigurationError,
    ExecutionError,
    ValidationError,
)
from agentrouter.tools import ToolRegistry, ToolExecutor
from agentrouter.types import (
    AgentConfig,
    AgentStatus,
    ExecutionContext,
    Message,
    ToolDefinition,
    WorkerTask,
    APIResponse,
)
from agentrouter.visualization import (
    ExecutionTracer,
    ExecutionVisualizer,
    PipelineInspector,
)
from agentrouter.utils.logging import log_agent, get_logger


logger = get_logger(__name__)


class BaseAgent(ABC):
    """
    Abstract base class for all agents
    """
    
    def __init__(
        self,
        name: str,
        api_key: str,
        model: str = "usf-mini",
        trace: bool = False,
        **config_kwargs
    ):
        """
        Initialize base agent
        
        Args:
            name: Agent name
            api_key: API key for authentication
            model: Model to use
            trace: Enable execution tracing for visualization (default: False)
            **config_kwargs: Additional configuration options
        """
        # Create unique instance ID
        self.instance_id = str(uuid.uuid4())
        
        # Initialize configuration
        self.config = AgentConfig(
            name=name,
            api_key=api_key,
            model=model,
            **config_kwargs
        )
        
        # Initialize provider base URL if provider is specified
        self._initialize_provider_base_url()
        
        # Initialize API clients with configuration from AgentConfig
        base_url = getattr(self.config, 'base_url', None)
        self.plan_api = PlanAPIClient(
            api_key=api_key,
            base_url=base_url,
            timeout=self.config.api_timeout,
            max_retries=self.config.max_retries,
            retry_delay=self.config.retry_delay,
            retry_multiplier=self.config.retry_multiplier,
            retry_max_wait=self.config.retry_max_wait
        )
        self.tool_call_api = ToolCallAPIClient(
            api_key=api_key,
            base_url=base_url,
            timeout=self.config.api_timeout,
            max_retries=self.config.max_retries,
            retry_delay=self.config.retry_delay,
            retry_multiplier=self.config.retry_multiplier,
            retry_max_wait=self.config.retry_max_wait
        )
        
        # Initialize tool registry and executor
        self.tool_registry = ToolRegistry(name=f"{name}_tools")
        self.tool_executor = ToolExecutor(self.tool_registry)
        
        # Initialize worker agents
        self._workers: Dict[str, 'BaseAgent'] = {}
        self._worker_aliases: Dict[str, str] = {}
        
        # Initialize plugins
        self._plugins: List[Any] = []
        
        # Initialize state
        self.status = AgentStatus.IDLE
        self._execution_context: Optional[ExecutionContext] = None
        self._parent_agent: Optional['BaseAgent'] = None
        
        # Statistics
        self._execution_count = 0
        self._total_tokens_used = 0
        
        # Initialize tracer (zero overhead when disabled)
        self._tracer: Optional[ExecutionTracer] = None
        if trace:
            self._tracer = ExecutionTracer(enabled=True)
            log_agent(name, "Tracing enabled", "debug")
        
        # Simple initialization log
        log_agent(name, f"Initialized ({self.__class__.__name__})", "info")
    
    def _initialize_provider_base_url(self) -> None:
        """Resolve and cache the provider base URL via get-base-url API"""
        model = getattr(self.config, "model", None)
        provider = getattr(self.config, "provider", None)

        if not model or model == "inherited":
            logger.debug(
                "Skipping base URL resolution for agent '%s' until model is provided",
                getattr(self.config, 'name', 'unknown')
            )
            return

        try:
            async def _resolve():
                # Create a fresh UtilityAPIClient bound to the current event loop
                utility_client = UtilityAPIClient()
                try:
                    return await utility_client.get_base_url(
                        provider=provider,
                        model=model
                    )
                finally:
                    # Ensure HTTP resources are released
                    await utility_client.close()

            try:
                loop = asyncio.get_event_loop()
            except RuntimeError:
                loop = None

            if loop and loop.is_running():
                result_container: Dict[str, Any] = {}
                error_container: Dict[str, Exception] = {}

                def _run_in_thread() -> None:
                    try:
                        result_container["value"] = asyncio.run(_resolve())
                    except Exception as thread_exc:  # pragma: no cover - propagate below
                        error_container["error"] = thread_exc

                resolver_thread = threading.Thread(target=_run_in_thread, daemon=True)
                resolver_thread.start()
                resolver_thread.join()

                if "error" in error_container:
                    raise error_container["error"]

                result = result_container.get("value")
            else:
                created_loop = False
                if loop is None:
                    loop = asyncio.new_event_loop()
                    asyncio.set_event_loop(loop)
                    created_loop = True
                try:
                    result = loop.run_until_complete(_resolve())
                finally:
                    if created_loop:
                        loop.close()
                        asyncio.set_event_loop(None)

            if result and result.get("status") == 1:
                data = result.get("data", {})
                resolved_base_url = data.get("baseUrl")
                is_our_model = data.get("isOurModel", False)

                if not resolved_base_url:
                    raise ConfigurationError(
                        "get-base-url response did not include a baseUrl value",
                        parameter="base_url",
                        details={"provider": provider, "model": model, "response": data}
                    )

                self.config.is_our_model = is_our_model

                if is_our_model:
                    # For AgentRouter-managed models we continue to send traffic to the
                    # resolved API base (typically https://api.us.inc/...)
                    self.config.base_url = resolved_base_url
                    self.config.llm_api_base_url = None
                    logger.debug(
                        f"Resolved AgentRouter base URL: {resolved_base_url}"
                    )
                else:
                    # External providers (OpenRouter, OpenAI, etc.) should only use the
                    # resolved base when calling the provider directly for final responses.
                    self.config.llm_api_base_url = resolved_base_url
                    logger.debug(
                        f"Resolved external LLM base URL for provider '{provider or 'default'}': {resolved_base_url}"
                    )
            else:
                message = (result or {}).get("message", "Unknown error") if result else "No response"
                raise ConfigurationError(
                    f"Failed to get base URL: {message}",
                    parameter="provider",
                    details={"provider": provider, "model": model}
                )

        except ConfigurationError:
            raise
        except Exception as e:
            raise ConfigurationError(
                f"Failed to initialize provider base URL: {str(e)}",
                parameter="provider",
                details={"provider": provider, "model": model}
            ) from e
    
    @property
    def name(self) -> str:
        """Get agent name"""
        return self.config.name
    
    @property
    def is_manager(self) -> bool:
        """Check if agent is a manager"""
        return self.__class__.__name__ == "ManagerAgent"
    
    @property
    def is_worker(self) -> bool:
        """Check if agent is a worker"""
        return self.__class__.__name__ == "WorkerAgent"
    
    def register_tool(self, tool_func: Any, alias: Optional[str] = None) -> None:
        """
        Register a tool with the agent
        
        Args:
            tool_func: Tool function decorated with @tool
            alias: Optional alias for the tool
        """
        self.tool_registry.register(tool_func, alias=alias)
        logger.debug(f"Tool registered: {tool_func.__name__ if hasattr(tool_func, '__name__') else 'tool'}")
    
    def unregister_tool(self, name: str) -> bool:
        """
        Unregister a tool
        
        Args:
            name: Tool name or alias
            
        Returns:
            True if tool was unregistered
        """
        return self.tool_registry.unregister(name)
    
    def list_tools(self) -> List[str]:
        """
        List all registered tools
        
        Returns:
            List of tool names
        """
        return self.tool_registry.list_tools()
    
    def attach_worker(
        self,
        worker: 'BaseAgent',
        alias: Optional[str] = None
    ) -> None:
        """
        Attach a worker agent with strict name uniqueness validation
        
        Args:
            worker: Worker agent to attach
            alias: Optional alias for the worker
            
        Raises:
            ValidationError: If worker name is duplicate or invalid
        """
        if not worker.is_worker:
            raise ValidationError(
                f"Can only attach WorkerAgent instances, got {type(worker)}"
            )
        
        worker_name = worker.name
        
        # STRICT NAME UNIQUENESS CHECK (fail-fast approach)
        if worker_name in self._workers:
            raise ValidationError(
                f"Worker '{worker_name}' already attached to '{self.name}'.\n"
                f"Each worker must have a unique name.\n"
                f"Existing workers: {', '.join(self._workers.keys())}"
            )
        
        # Check alias uniqueness if provided
        if alias:
            if alias in self._worker_aliases:
                raise ValidationError(
                    f"Alias '{alias}' already in use for worker '{self._worker_aliases[alias]}'.\n"
                    f"Each alias must be unique.\n"
                    f"Existing aliases: {', '.join(self._worker_aliases.keys())}"
                )
            if alias in self._workers:
                raise ValidationError(
                    f"Alias '{alias}' conflicts with existing worker name.\n"
                    f"Aliases cannot match worker names."
                )
        
        # Attach worker
        self._workers[worker_name] = worker
        worker._parent_agent = self
        
        # Set alias if provided
        if alias:
            self._worker_aliases[alias] = worker_name
        
        logger.debug(f"Worker '{worker_name}' attached" + (f" with alias '{alias}'" if alias else ""))
    
    def detach_worker(self, name: str) -> bool:
        """
        Detach a worker agent
        
        Args:
            name: Worker name or alias
            
        Returns:
            True if worker was detached
        """
        # Resolve alias
        worker_name = self._worker_aliases.get(name, name)
        
        if worker_name in self._workers:
            worker = self._workers[worker_name]
            worker._parent_agent = None
            del self._workers[worker_name]
            
            # Remove aliases
            aliases_to_remove = [
                alias for alias, target in self._worker_aliases.items()
                if target == worker_name
            ]
            for alias in aliases_to_remove:
                del self._worker_aliases[alias]
            
            logger.debug(f"Worker '{worker_name}' detached")
            return True
        
        return False
    
    def get_worker(self, name: str) -> Optional['BaseAgent']:
        """
        Get an attached worker
        
        Args:
            name: Worker name or alias
            
        Returns:
            Worker agent if found
        """
        worker_name = self._worker_aliases.get(name, name)
        return self._workers.get(worker_name)
    
    def list_workers(self) -> List[str]:
        """
        List all attached workers
        
        Returns:
            List of worker names
        """
        return list(self._workers.keys())
    
    def add_plugin(self, plugin: Any) -> None:
        """
        Add a plugin to the agent
        
        Args:
            plugin: Plugin instance
        """
        self._plugins.append(plugin)
        logger.debug(f"Plugin added to agent '{self.name}'")
    
    def remove_plugin(self, plugin: Any) -> bool:
        """
        Remove a plugin from the agent
        
        Args:
            plugin: Plugin instance
            
        Returns:
            True if plugin was removed
        """
        if plugin in self._plugins:
            self._plugins.remove(plugin)
            logger.debug(f"Plugin removed from agent '{self.name}'")
            return True
        return False
    
    async def _apply_plugin_hook(
        self,
        hook_name: str,
        context: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Apply plugin hooks
        
        Args:
            hook_name: Name of the hook
            context: Context to pass to plugins
            
        Returns:
            Modified context
        """
        for plugin in self._plugins:
            if hasattr(plugin, hook_name):
                hook_method = getattr(plugin, hook_name)
                if callable(hook_method):
                    try:
                        context = await hook_method(context) if asyncio.iscoroutinefunction(hook_method) else hook_method(context)
                    except Exception as e:
                        logger.error(
                            f"Plugin hook '{hook_name}' failed: {str(e)}",
                            exc_info=True
                        )
        return context
    
    _SYSTEM_SECTION_SEPARATOR = "\n\n---\n\n"

    def _prepare_system_sections(
        self,
        additional_sections: Optional[List[str]] = None
    ) -> List[str]:
        """Prepare ordered system sections for the agent configuration."""
        sections: List[str] = []

        if self.config.introduction:
            sections.append(f"### Introduction: {self.config.introduction}")

        if self.config.backstory:
            sections.append(f"### Backstory: {self.config.backstory}")

        if self.config.goal:
            sections.append(f"### Goal: {self.config.goal}")

        if self.config.knowledge_cutoff:
            sections.append(f"### Knowledge cutoff: {self.config.knowledge_cutoff}")

        if self.config.current_time:
            sections.append(f"### Current time: {self.config.current_time}")

        if self.config.instruction:
            sections.append(f"### Instructions: {self.config.instruction}")

        if additional_sections:
            sections.extend([section.strip() for section in additional_sections if section and section.strip()])

        return sections

    def _prepare_system_message(
        self,
        additional_sections: Optional[List[str]] = None
    ) -> Optional[Message]:
        """
        Prepare system message with agent configuration

        Args:
            additional_sections: Extra sections to include after the agent configuration

        Returns:
            System message if configuration exists
        """
        parts = self._prepare_system_sections(additional_sections)

        if parts:
            return Message(
                role="system",
                content=self._SYSTEM_SECTION_SEPARATOR.join(parts)
            )

        return None

    def _update_or_add_system_message(
        self,
        messages: List[Message],
        additional_sections: Optional[List[str]] = None
    ) -> List[Message]:
        """
        Update existing system message or add new one
        
        Args:
            messages: Current messages (can be Message objects or dicts)
            
        Returns:
            Updated messages
        """
        separator = self._SYSTEM_SECTION_SEPARATOR
        system_msg = self._prepare_system_message(additional_sections)

        # Always work on a shallow copy to avoid mutating caller state
        working_messages = list(messages) if messages else []

        # Gather existing system sections and strip system messages from the list
        existing_sections: List[str] = []
        retained_messages: List[Message] = []
        for msg in working_messages:
            if isinstance(msg, dict):
                msg_role = msg.get("role")
                msg_content = msg.get("content")
            else:
                msg_role = getattr(msg, "role", None)
                # Pydantic Message stores enum values directly
                if hasattr(msg_role, "value"):
                    msg_role = msg_role.value
                msg_content = getattr(msg, "content", None)

            if msg_role == "system":
                if msg_content:
                    content_str = msg_content if isinstance(msg_content, str) else str(msg_content)
                    parts = [part.strip() for part in content_str.split(separator) if part and part.strip()]
                    existing_sections.extend(parts)
                continue  # Drop system messages; we'll rebuild a single one below

            retained_messages.append(msg)

        # If we have neither configured sections nor existing sections, return as-is
        target_sections = []
        if system_msg:
            target_sections.extend([
                part.strip()
                for part in system_msg.content.split(separator)
                if part and part.strip()
            ])

        combined_sections: List[str] = []
        seen: set[str] = set()

        for section in target_sections + existing_sections:
            key = section.strip()
            if key and key not in seen:
                combined_sections.append(key)
                seen.add(key)

        if combined_sections:
            new_system_msg = Message(
                role="system",
                content=separator.join(combined_sections)
            )
            return [new_system_msg] + retained_messages

        # No system content to add - return retained messages without modification
        return retained_messages
    
    def get_statistics(self) -> Dict[str, Any]:
        """
        Get agent statistics
        
        Returns:
            Statistics dictionary
        """
        return {
            "name": self.name,
            "instance_id": self.instance_id,
            "type": self.__class__.__name__,
            "status": self.status.value,
            "execution_count": self._execution_count,
            "total_tokens_used": self._total_tokens_used,
            "tools_registered": len(self.tool_registry),
            "workers_attached": len(self._workers),
            "plugins_installed": len(self._plugins),
            "tool_execution_stats": self.tool_executor.get_statistics()
        }
    
    def reset_statistics(self) -> None:
        """Reset agent statistics"""
        self._execution_count = 0
        self._total_tokens_used = 0
        self.tool_executor.clear_history()
        logger.debug(f"Statistics reset for agent '{self.name}'")
    
    @abstractmethod
    async def execute(
        self,
        messages: List[Message],
        **kwargs
    ) -> Dict[str, Any]:
        """
        Execute the agent
        
        Args:
            messages: Input messages
            **kwargs: Additional execution parameters
            
        Returns:
            Execution result in OpenAI format
        """
        pass
    
    @abstractmethod
    def create_worker(
        self,
        name: str,
        **config_kwargs
    ) -> 'BaseAgent':
        """
        Create a new worker agent
        
        Args:
            name: Worker name
            **config_kwargs: Worker configuration
            
        Returns:
            Created worker agent
        """
        pass
    
    def enable_tracing(self) -> ExecutionTracer:
        """
        Enable execution tracing for this agent instance.
        
        Returns:
            ExecutionTracer instance
        """
        if not self._tracer:
            self._tracer = ExecutionTracer(enabled=True)
            log_agent(self.name, "Tracing enabled", "debug")
        return self._tracer
    
    def disable_tracing(self) -> None:
        """Disable execution tracing for this agent"""
        if self._tracer:
            self._tracer.enabled = False
            log_agent(self.name, "Tracing disabled", "debug")
    
    def visualize(
        self,
        format: str = "mermaid",
        output: Optional[str] = None,
        **kwargs
    ) -> Any:
        """
        Visualize the last execution (only if tracing was enabled).

        Args:
            format: Output format ('mermaid', 'html', 'json', 'png', 'custom')
            output: Optional file path to save the primary visualization output.
            **kwargs: Additional options including:
                - return_image (bool): When True and format='mermaid', return a
                  dict containing the Mermaid text, PNG bytes, and base64 helpers.
                - image_output (str): Path to write the generated PNG when
                  format='mermaid'.
                - mermaid_renderer (Callable): Custom callable to convert
                  Mermaid text to PNG bytes.
                - mermaid_render_url (str): Override URL for the default renderer.
                - mermaid_render_timeout (float): Timeout for the default renderer.
                - include_styles (bool): Passed through to JSON export.

        Returns:
            For 'mermaid', 'html', and 'json' formats returns a string by default.
            For 'png' returns PNG bytes.
            When return_image=True returns a dict with both Mermaid text and PNG data.
        """
        if not self._tracer or not self._tracer.enabled:
            return "No execution data available. Enable tracing with trace=True or enable_tracing()"

        viz = ExecutionVisualizer(self._tracer)

        mermaid_renderer = kwargs.get("mermaid_renderer")
        mermaid_render_url = kwargs.get("mermaid_render_url")
        mermaid_render_timeout = kwargs.get("mermaid_render_timeout", 30.0)
        image_output_path = kwargs.get("image_output")
        return_image = kwargs.get("return_image", False)

        save_payload: Optional[Any] = None
        save_binary = False

        if format == "mermaid":
            mermaid_diagram = viz.generate_mermaid()
            save_payload = mermaid_diagram

            png_bytes: Optional[bytes] = None
            if return_image or image_output_path:
                png_bytes = viz.generate_mermaid_png(
                    mermaid_diagram=mermaid_diagram,
                    renderer=mermaid_renderer,
                    kroki_url=mermaid_render_url,
                    timeout=mermaid_render_timeout,
                )

                if image_output_path:
                    with open(image_output_path, "wb") as image_file:
                        image_file.write(png_bytes)
                    logger.debug("Mermaid PNG visualization saved to %s", image_output_path)

            if return_image:
                if png_bytes is None:
                    png_bytes = viz.generate_mermaid_png(
                        mermaid_diagram=mermaid_diagram,
                        renderer=mermaid_renderer,
                        kroki_url=mermaid_render_url,
                        timeout=mermaid_render_timeout,
                    )

                encoded = base64.b64encode(png_bytes).decode("ascii") if png_bytes else None
                result: Any = {
                    "mermaid": mermaid_diagram,
                    "png_bytes": png_bytes,
                    "png_base64": encoded,
                    "png_data_uri": f"data:image/png;base64,{encoded}" if encoded else None,
                }
            else:
                result = mermaid_diagram

        elif format == "html":
            result = viz.generate_html()
            save_payload = result

        elif format == "json":
            json_data = viz.generate_json(kwargs.get("include_styles", True))
            result = json.dumps(json_data)
            save_payload = result

        elif format == "png":
            png_bytes = viz.generate_mermaid_png(
                renderer=mermaid_renderer,
                kroki_url=mermaid_render_url,
                timeout=mermaid_render_timeout,
            )
            result = png_bytes
            save_payload = png_bytes
            save_binary = True

        elif format == "custom" and "formatter" in kwargs:
            result = viz.export_custom(kwargs["formatter"])
            save_payload = result
            save_binary = isinstance(result, (bytes, bytearray))

        else:
            raise ValueError(f"Unsupported format: {format}")

        if output and save_payload is not None:
            mode = "wb" if save_binary or isinstance(save_payload, (bytes, bytearray)) else "w"
            with open(output, mode) as f:
                f.write(save_payload)
            logger.debug("Visualization saved to %s", output)

        return result
    
    def inspect(self) -> PipelineInspector:
        """
        Create a pipeline inspector for this agent.
        
        Returns:
            PipelineInspector instance for analyzing agent structure
        """
        return PipelineInspector(self)
    
    def visualize_pipeline(
        self,
        format: str = "console",
        output: Optional[str] = None
    ) -> str:
        """
        Visualize the static pipeline structure (before execution).
        
        Args:
            format: Visualization format ('console', 'mermaid', 'json')
            output: Optional file path to save output
        
        Returns:
            Pipeline visualization as string
        """
        inspector = self.inspect()
        return inspector.visualize(format, output)
    
    def get_tracer_stats(self) -> Dict[str, Any]:
        """
        Get execution tracer statistics.
        
        Returns:
            Tracer statistics if enabled, empty dict otherwise
        """
        if self._tracer and self._tracer.enabled:
            return self._tracer.get_statistics()
        return {"enabled": False}
    
    def __repr__(self) -> str:
        """String representation"""
        tracing = " [TRACING]" if self._tracer and self._tracer.enabled else ""
        return (
            f"{self.__class__.__name__}("
            f"name='{self.name}', "
            f"model='{self.config.model}', "
            f"tools={len(self.tool_registry)}, "
            f"workers={len(self._workers)}){tracing}"
        )
