"""
Mirix SDK - Simple Python interface for memory-enhanced AI agents.

All I/O methods are async. Use Mirix.create() to construct an instance.
"""

import logging
import os
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

from mirix.local_client.local_client import LocalClient
from mirix.schemas.agent import AgentType, CreateMetaAgent
from mirix.schemas.embedding_config import EmbeddingConfig
from mirix.schemas.llm_config import LLMConfig
from mirix.schemas.memory import Memory

logger = logging.getLogger(__name__)


def load_config(config_path: str) -> Dict[str, Any]:
    """
    Load configuration from a YAML file.

    Args:
        config_path: Path to the YAML configuration file

    Returns:
        Dict containing the configuration

    Example:
        >>> config = load_config("mirix/configs/mirix.yaml")
        >>> client = MirixClient(org_id="demo-org")
        >>> client.initialize_meta_agent(config=config)
    """
    import yaml

    config_file = Path(config_path)
    if not config_file.exists():
        raise FileNotFoundError(f"Config file not found: {config_path}")

    with open(config_file, "r") as f:
        config = yaml.safe_load(f)

    return config


class Mirix:
    """
    Simple SDK interface for Mirix memory agent.

    Example:
        from mirix import Mirix

        memory_agent = Mirix(api_key="your-api-key")
        memory_agent.add("The moon now has a president")
        memories = memory_agent.visualize_memories()
    """

    def __init__(
        self,
        api_key: str,
        model_provider: str = "google_ai",
        model: Optional[str] = None,
        config_path: Optional[str] = None,
        load_from: Optional[str] = None,
        _client: Optional[LocalClient] = None,
        _meta_agent=None,
        **kwargs,
    ):
        """
        Initialize Mirix memory agent. Prefer Mirix.create() for async setup.

        Args:
            api_key: API key for LLM provider (required)
            model_provider: LLM provider name (default: "google_ai")
            model: Model to use (optional). If None, uses default model.
            config_path: Path to custom config file (optional)
            load_from: Path to backup directory to restore from (optional)
            _client: Internal: pre-created LocalClient (used by create()).
            _meta_agent: Internal: pre-resolved meta agent (used by create()).
        """
        if not api_key and _client is None:
            raise ValueError("api_key is required to initialize Mirix")

        if api_key:
            # Set API key environment variable based on provider
            if model_provider.lower() in ["google", "google_ai", "gemini"]:
                os.environ["GEMINI_API_KEY"] = api_key
            elif model_provider.lower() in ["openai", "gpt"]:
                os.environ["OPENAI_API_KEY"] = api_key
            elif model_provider.lower() in ["anthropic", "claude"]:
                os.environ["ANTHROPIC_API_KEY"] = api_key
            else:
                os.environ[f"{model_provider.upper()}_API_KEY"] = api_key
            self._reload_model_settings()

        self._config_path = config_path
        self._model = model
        self._model_provider = model_provider
        self._load_from = load_from
        self._kwargs = kwargs
        self._client = _client
        self._meta_agent = _meta_agent

    @classmethod
    async def create(
        cls,
        api_key: str,
        model_provider: str = "google_ai",
        model: Optional[str] = None,
        config_path: Optional[str] = None,
        load_from: Optional[str] = None,
        **kwargs,
    ) -> "Mirix":
        """
        Create and initialize a Mirix memory agent (async). Use this instead of __init__.

        Example:
            memory_agent = await Mirix.create(api_key="your-api-key")
            await memory_agent.add("The moon now has a president")
        """
        if not api_key:
            raise ValueError("api_key is required")
        if model_provider.lower() in ["google", "google_ai", "gemini"]:
            os.environ["GEMINI_API_KEY"] = api_key
        elif model_provider.lower() in ["openai", "gpt"]:
            os.environ["OPENAI_API_KEY"] = api_key
        elif model_provider.lower() in ["anthropic", "claude"]:
            os.environ["ANTHROPIC_API_KEY"] = api_key
        else:
            os.environ[f"{model_provider.upper()}_API_KEY"] = api_key
        import mirix.settings
        from mirix.settings import ModelSettings

        new_settings = ModelSettings()
        for field_name in ModelSettings.model_fields:
            setattr(
                mirix.settings.model_settings,
                field_name,
                getattr(new_settings, field_name),
            )

        system_prompts_folder = None
        llm_config = LLMConfig.default_config(model or "gemini-2.0-flash")
        embedding_config = EmbeddingConfig.default_config("text-embedding-004")
        if config_path:
            config_path = Path(config_path)
            if config_path.exists():
                import yaml

                with open(config_path, "r") as f:
                    config_data = yaml.safe_load(f)
                    system_prompts_folder = config_data.get("system_prompts_folder")
                    if "llm_config" in config_data:
                        llm_config = LLMConfig(**config_data["llm_config"])
                    if "embedding_config" in config_data:
                        embedding_config = EmbeddingConfig(**config_data["embedding_config"])

        client = await LocalClient.create()
        client.set_default_llm_config(llm_config)
        client.set_default_embedding_config(embedding_config)

        agents = await client.list_agents()
        existing_meta_agent = None
        for agent in agents:
            if agent.agent_type == AgentType.meta_memory_agent:
                existing_meta_agent = agent
                break
        if existing_meta_agent:
            meta_agent = existing_meta_agent
        else:
            create_request = CreateMetaAgent(
                system_prompts_folder=system_prompts_folder,
                llm_config=llm_config,
                embedding_config=embedding_config,
            )
            meta_agent = await client.create_meta_agent(request=create_request)

        inst = cls(
            api_key=api_key,
            model_provider=model_provider,
            model=model,
            config_path=config_path,
            load_from=load_from,
            _client=client,
            _meta_agent=meta_agent,
            **kwargs,
        )
        return inst

    def _require_meta_agent(self):
        """Raise if meta agent is not available (e.g. client has no write_scope)."""
        if not self._meta_agent:
            raise RuntimeError(
                "Meta agent is not available. This usually means the client has no write_scope configured. "
                "A write_scope is required to create and use memory agents."
            )

    async def add(self, content: str, **kwargs) -> Dict[str, Any]:
        """
        Add information to memory.

        Args:
            content: Information to memorize
            **kwargs: Additional options (images, metadata, etc.)

        Returns:
            Response from the memory system

        Example:
            await memory_agent.add("John likes pizza")
        """
        self._require_meta_agent()
        response = await self._client.send_message(agent_id=self._meta_agent.id, role="user", message=content, **kwargs)
        if hasattr(response, "messages") and response.messages:
            for msg in reversed(response.messages):
                if msg.role == "assistant":
                    return {"response": msg.text, "success": True}
        return {"response": str(response), "success": True}

    async def list_users(self) -> List[Any]:
        """List all users in the system."""
        return await self._client.server.user_manager.list_users()

    async def construct_system_message(self, message: str, user_id: str) -> str:
        """Construct a system message from a message."""
        self._require_meta_agent()
        return await self._client.construct_system_message(
            agent_id=self._meta_agent.id, message=message, user_id=user_id
        )

    async def extract_memory_for_system_prompt(self, message: str, user_id: str) -> str:
        """Extract memory for system prompt from a message."""
        self._require_meta_agent()
        return await self._client.extract_memory_for_system_prompt(
            agent_id=self._meta_agent.id, message=message, user_id=user_id
        )

    async def get_user_by_name(self, user_name: str):
        """Get a user by their name. Returns User if found, None otherwise."""
        users = await self.list_users()
        for user in users:
            if user.name == user_name:
                return user
        return None

    async def clear_conversation_history(self, user_id: Optional[str] = None) -> Dict[str, Any]:
        """
        Clear conversation history while preserving memories.

        This removes persisted conversation message rows while preserving
        memory tables and agent configuration.

        Args:
            user_id: User ID to clear messages for. If None, clears all retained
                    conversation rows for the meta agent. If provided, only clears
                    messages for that specific user.

        Returns:
            Dict containing success status, message, and count of deleted messages

        Example:
            # Clear all conversation history
            result = memory_agent.clear_conversation_history()

            # Clear history for specific user
            result = memory_agent.clear_conversation_history(user_id="user_123")

            if result['success']:
                logger.debug("Cleared %s messages", result['messages_deleted'])
            else:
                logger.debug("Failed to clear: %s", result['error'])
        """
        self._require_meta_agent()
        try:
            if user_id is None:
                messages_count = 0  # count not available without per-user query
                await self._client.server.agent_manager.reset_messages(
                    agent_id=self._meta_agent.id,
                    actor=self._client.client,
                    user_id=None,
                )
                return {
                    "success": True,
                    "message": "Successfully cleared retained conversation history.",
                    "messages_deleted": messages_count,
                }
            else:
                target_user = await self._client.server.user_manager.get_user_by_id(user_id)
                if not target_user:
                    return {
                        "success": False,
                        "error": f"User with ID '{user_id}' not found",
                        "messages_deleted": 0,
                    }
                current_messages = await self._client.server.message_manager.get_messages_for_agent_user(
                    agent_id=self._meta_agent.id,
                    user_id=target_user.id,
                    actor=self._client.client,
                    limit=10000,
                )
                user_messages_count = len(current_messages)
                await self._client.server.agent_manager.reset_messages(
                    agent_id=self._meta_agent.id,
                    actor=self._client.client,
                    user_id=target_user.id,
                )
                return {
                    "success": True,
                    "message": f"Successfully cleared conversation history for {target_user.name}.",
                    "messages_deleted": user_messages_count,
                }
        except Exception as e:
            return {"success": False, "error": str(e), "messages_deleted": 0}

    def _reload_model_settings(self):
        """
        Force reload of model_settings to pick up new environment variables.

        This is necessary because Pydantic BaseSettings loads environment variables
        at class instantiation time, which happens at import. Since the SDK sets
        environment variables after import, we need to manually update the singleton.
        """
        from mirix.settings import ModelSettings

        # Create a new instance with current environment variables
        new_settings = ModelSettings()

        # Update the global singleton instance with new values
        import mirix.settings

        for field_name in ModelSettings.model_fields:
            setattr(
                mirix.settings.model_settings,
                field_name,
                getattr(new_settings, field_name),
            )

    async def create_user(
        self,
        user_name: str,
        timezone: str = "UTC",
        organization_id: Optional[str] = None,
    ) -> Any:
        """Create a new user in the system."""
        from mirix.schemas.user import UserCreate
        from mirix.services.organization_manager import OrganizationManager

        if organization_id is None:
            organization_id = OrganizationManager.DEFAULT_ORG_ID
        return await self._client.server.user_manager.create_user(
            pydantic_user=UserCreate(
                name=user_name,
                timezone=timezone,
                organization_id=organization_id,
            )
        )

    async def __call__(self, message: str) -> Dict[str, Any]:
        """Use the agent as a callable for adding memories: await memory_agent(message)."""
        return await self.add(message)

    async def insert_tool(
        self,
        name: str,
        source_code: str,
        description: str,
        args_info: Optional[Dict[str, str]] = None,
        returns_info: Optional[str] = None,
        tags: Optional[List[str]] = None,
        apply_to_agents: Union[List[str], str] = "all",
    ) -> Dict[str, Any]:
        """
        Insert a custom tool into the system.

        Args:
            name: The name of the tool function
            source_code: The Python source code for the tool function (without docstring)
            description: Description of what the tool does
            args_info: Optional dict mapping argument names to their descriptions
            returns_info: Optional description of what the function returns
            tags: Optional list of tags for categorization (defaults to ["user_defined"])
            apply_to_all_agents: Whether to add this tool to all existing agents (default: True)

        Returns:
            Dict containing success status, tool data, and any error messages

        Example:
            result = memory_agent.insert_tool(
                name="calculate_sum",
                source_code="def calculate_sum(a: int, b: int) -> int:\n    return a + b",
                description="Calculate the sum of two numbers",
                args_info={"a": "First number", "b": "Second number"},
                returns_info="The sum of a and b",
                tags=["math", "utility"]
            )
        """
        from mirix.schemas.enums import ToolType
        from mirix.schemas.tool import Tool as PydanticTool
        from mirix.services.tool_manager import ToolManager

        # Initialize tool manager
        tool_manager = ToolManager()

        # Check if tool name already exists
        existing_tool = await tool_manager.get_tool_by_name(tool_name=name, actor=self._client.client)

        if existing_tool:
            created_tool = existing_tool

        else:
            # Set default tags if not provided
            if tags is None:
                tags = ["user_defined"]

            # Construct complete source code with docstring
            complete_source_code = self._build_complete_source_code(source_code, description, args_info, returns_info)

            # Generate JSON schema from the complete source code
            from mirix.functions.functions import derive_openai_json_schema

            json_schema = derive_openai_json_schema(source_code=complete_source_code, name=name)

            # Create the tool object
            pydantic_tool = PydanticTool(
                name=name,
                source_code=complete_source_code,
                source_type="python",
                tool_type=ToolType.USER_DEFINED,
                tags=tags,
                description=description,
                json_schema=json_schema,
            )

            # Use the tool manager's create_or_update_tool method
            created_tool = await tool_manager.create_or_update_tool(
                pydantic_tool=pydantic_tool, actor=self._client.client
            )

        # Apply tool to all existing agents if requested
        if apply_to_agents:
            # Get all existing agents
            all_agents = await self._client.server.agent_manager.list_agents(
                actor=self._client.client,
                limit=1000,
            )

            if apply_to_agents != "all":
                all_agents = [agent for agent in all_agents if agent.name in apply_to_agents]

            # Add the tool to each agent
            for agent in all_agents:
                # Get current agent tools
                existing_tools = agent.tools
                existing_tool_ids = [tool.id for tool in existing_tools]

                # Add the new tool if not already present
                if created_tool.id not in existing_tool_ids:
                    new_tool_ids = existing_tool_ids + [created_tool.id]

                    # Update the agent with the new tool
                    from mirix.schemas.agent import UpdateAgent

                    await self._client.server.agent_manager.update_agent(
                        agent_id=agent.id,
                        agent_update=UpdateAgent(tool_ids=new_tool_ids),
                        actor=self._client.client,
                    )

        return {
            "success": True,
            "message": f"Tool '{name}' inserted successfully"
            + (" and applied to all existing agents" if apply_to_agents else ""),
            "tool": {
                "id": created_tool.id,
                "name": created_tool.name,
                "description": created_tool.description,
                "tags": created_tool.tags,
                "tool_type": created_tool.tool_type.value if created_tool.tool_type else None,
            },
        }

    def _build_complete_source_code(
        self,
        source_code: str,
        description: str,
        args_info: Optional[Dict[str, str]] = None,
        returns_info: Optional[str] = None,
    ) -> str:
        """
        Build complete source code with proper docstring from user-provided components.

        Args:
            source_code: The bare function code without docstring
            description: Function description
            args_info: Optional dict mapping argument names to descriptions
            returns_info: Optional return value description

        Returns:
            Complete source code with properly formatted docstring
        """
        import re

        # Find the function definition line
        func_match = re.search(r"(def\s+\w+\([^)]*\)\s*(?:->\s*[^:]+)?:)", source_code)
        if not func_match:
            raise ValueError("Invalid function definition in source_code")

        func_def = func_match.group(1)
        func_body = source_code[func_match.end() :].lstrip("\n")

        # Build docstring
        docstring_lines = ['    """', f"    {description}"]

        if args_info:
            docstring_lines.extend(["", "    Args:"])
            for arg_name, arg_desc in args_info.items():
                docstring_lines.append(f"        {arg_name}: {arg_desc}")

        if returns_info:
            docstring_lines.extend(["", "    Returns:", f"        {returns_info}"])

        docstring_lines.append('    """')

        # Combine everything
        complete_code = func_def + "\n" + "\n".join(docstring_lines) + "\n" + func_body

        return complete_code

    async def visualize_memories(self, user_id: Optional[str] = None) -> Dict[str, Any]:
        """
        Visualize all memories for a specific user.

        Args:
            user_id: User ID to get memories for. If None, uses current active user.

        Returns:
            Dict containing all memory types organized by category

        Example:
            memories = memory_agent.visualize_memories(user_id="user_123")
            logger.debug("Episodic memories: %s", len(memories['episodic']))
            logger.debug("Semantic memories: %s", len(memories['semantic']))
        """
        self._require_meta_agent()
        try:
            # Find the target user
            if user_id:
                target_user = await self._client.server.user_manager.get_user_by_id(user_id)
                if not target_user:
                    return {
                        "success": False,
                        "error": f"User with ID '{user_id}' not found",
                    }
            else:
                users = await self._client.server.user_manager.list_users()
                active_user = next((user for user in users if user.status == "active"), None)
                target_user = active_user if active_user else (users[0] if users else None)

            if not target_user:
                return {"success": False, "error": "No user found"}

            memories = {}

            # Get episodic memory
            try:
                episodic_manager = self._client.server.episodic_memory_manager
                episodic_agent = None
                child_agents = await self._client.list_agents(parent_id=self._meta_agent.id)
                for agent in child_agents:
                    if agent.agent_type == AgentType.episodic_memory_agent:
                        episodic_agent = agent
                        break

                if episodic_agent:
                    events = await episodic_manager.list_episodic_memory(
                        agent_state=episodic_agent,
                        actor=target_user,
                        limit=50,
                        timezone_str=target_user.timezone,
                    )
                else:
                    events = []

                memories["episodic"] = []
                for event in events:
                    memories["episodic"].append(
                        {
                            "timestamp": event.occurred_at.isoformat() if event.occurred_at else None,
                            "summary": event.summary,
                            "details": event.details,
                            "event_type": event.event_type,
                        }
                    )
            except Exception:
                memories["episodic"] = []

            # Get semantic memory
            try:
                semantic_manager = self._client.server.semantic_memory_manager
                # Find semantic memory agent from meta agent's children
                semantic_agent = None
                for agent in child_agents:
                    if agent.agent_type == AgentType.semantic_memory_agent:
                        semantic_agent = agent
                        break

                if semantic_agent:
                    semantic_items = await semantic_manager.list_semantic_items(
                        agent_state=semantic_agent,
                        actor=target_user,
                        limit=50,
                        timezone_str=target_user.timezone,
                    )
                else:
                    semantic_items = []

                memories["semantic"] = []
                for item in semantic_items:
                    memories["semantic"].append(
                        {
                            "title": item.name,
                            "type": "semantic",
                            "summary": item.summary,
                            "details": item.details,
                        }
                    )
            except Exception:
                memories["semantic"] = []

            # Get procedural memory
            try:
                procedural_manager = self._client.server.procedural_memory_manager
                # Find procedural memory agent from meta agent's children
                procedural_agent = None
                for agent in child_agents:
                    if agent.agent_type == AgentType.procedural_memory_agent:
                        procedural_agent = agent
                        break

                if procedural_agent:
                    procedural_items = await procedural_manager.list_procedures(
                        agent_state=procedural_agent,
                        actor=target_user,
                        limit=50,
                        timezone_str=target_user.timezone,
                    )
                else:
                    procedural_items = []

                memories["procedural"] = []
                for item in procedural_items:
                    import json

                    # Parse steps if it's a JSON string
                    steps = item.steps
                    if isinstance(steps, str):
                        try:
                            steps = json.loads(steps)
                            # Extract just the instruction text for simpler display
                            if isinstance(steps, list) and steps and isinstance(steps[0], dict):
                                steps = [step.get("instruction", str(step)) for step in steps]
                        except (json.JSONDecodeError, KeyError, TypeError):
                            # If parsing fails, keep as string and split by common delimiters
                            if isinstance(steps, str):
                                steps = [s.strip() for s in steps.replace("\n", "|").split("|") if s.strip()]
                            else:
                                steps = []

                    memories["procedural"].append(
                        {
                            "title": item.entry_type,
                            "type": "procedural",
                            "summary": item.summary,
                            "steps": steps if isinstance(steps, list) else [],
                        }
                    )
            except Exception:
                memories["procedural"] = []

            # Get resource memory
            try:
                resource_manager = self._client.server.resource_memory_manager
                # Find resource memory agent from meta agent's children
                resource_agent = None
                for agent in child_agents:
                    if agent.agent_type == AgentType.resource_memory_agent:
                        resource_agent = agent
                        break

                if resource_agent:
                    resources = await resource_manager.list_resources(
                        agent_state=resource_agent,
                        actor=target_user,
                        limit=50,
                        timezone_str=target_user.timezone,
                    )
                else:
                    resources = []

                memories["resources"] = []
                for resource in resources:
                    memories["resources"].append(
                        {
                            "filename": resource.title,
                            "type": resource.resource_type,
                            "summary": resource.summary
                            or (resource.content[:200] + "..." if len(resource.content) > 200 else resource.content),
                            "last_accessed": resource.updated_at.isoformat() if resource.updated_at else None,
                        }
                    )
            except Exception:
                memories["resources"] = []

            # Get core memory
            try:
                blocks = await self._client.server.block_manager.get_blocks(user=target_user)
                core_blocks = []
                for block in blocks:
                    b = await self._client.server.block_manager.get_block_by_id(block.id, user=target_user)
                    if b:
                        core_blocks.append(b)
                core_memory = Memory(blocks=core_blocks)

                memories["core"] = []
                total_characters = 0

                for block in core_memory.blocks:
                    if block.value and block.value.strip() and block.label.lower() != "persona":
                        block_chars = len(block.value)
                        total_characters += block_chars

                        memories["core"].append(
                            {
                                "aspect": block.label,
                                "understanding": block.value,
                                "character_count": block_chars,
                                "total_characters": total_characters,
                                "max_characters": block.limit,
                                "last_updated": None,
                            }
                        )
            except Exception:
                memories["core"] = []

            # Get credentials memory
            try:
                knowledge_vault_manager = self._client.server.knowledge_vault_manager
                # Find knowledge vault agent from meta agent's children
                knowledge_vault_memory_agent = None
                for agent in child_agents:
                    if agent.agent_type == AgentType.knowledge_vault_memory_agent:
                        knowledge_vault_memory_agent = agent
                        break

                if knowledge_vault_memory_agent:
                    vault_items = await knowledge_vault_manager.list_knowledge(
                        actor=target_user,
                        agent_state=knowledge_vault_memory_agent,
                        limit=50,
                        timezone_str=target_user.timezone,
                    )
                else:
                    vault_items = []

                memories["credentials"] = []
                for item in vault_items:
                    memories["credentials"].append(
                        {
                            "caption": item.caption,
                            "entry_type": item.entry_type,
                            "source": item.source,
                            "sensitivity": item.sensitivity,
                            "content": "••••••••••••" if item.sensitivity == "high" else item.secret_value,
                        }
                    )
            except Exception:
                memories["credentials"] = []

            return {
                "success": True,
                "user_id": target_user.id,
                "user_name": target_user.name,
                "memories": memories,
                "summary": {
                    "episodic_count": len(memories.get("episodic", [])),
                    "semantic_count": len(memories.get("semantic", [])),
                    "procedural_count": len(memories.get("procedural", [])),
                    "resources_count": len(memories.get("resources", [])),
                    "core_count": len(memories.get("core", [])),
                    "credentials_count": len(memories.get("credentials", [])),
                },
            }

        except Exception as e:
            return {"success": False, "error": str(e)}

    async def update_core_memory(self, label: str, text: str, user_id: Optional[str] = None) -> Dict[str, Any]:
        """
        Update a specific core memory block with new text.

        Core memory blocks are persistent memory aspects that help the agent
        understand and remember key information about users, preferences, and context.

        Args:
            label: The label/name of the core memory block to update (e.g., "persona", "user_preferences")
            text: The new text content for the memory block
            user_id: User ID to update core memory for. If None, uses the current active user.

        Returns:
            Dict containing success status and message

        Example:
            # Update core memory for current user
            result = memory_agent.update_core_memory(
                label="user_preferences",
                text="User prefers concise responses and technical details"
            )

            # Update core memory for specific user
            result = memory_agent.update_core_memory(
                label="user_preferences",
                text="Alice prefers detailed explanations",
                user_id="user_123"
            )

            if result['success']:
                logger.debug("Core memory updated successfully")
            else:
                logger.debug("Update failed: %s", result['message'])
        """
        self._require_meta_agent()
        try:
            # If user_id is provided, get the specific user
            if user_id:
                target_user = await self._client.server.user_manager.get_user_by_id(user_id)
                if not target_user:
                    return {
                        "success": False,
                        "message": f"User with ID '{user_id}' not found",
                    }
            else:
                target_user = self._client.user

            await self._client.update_in_context_memory(agent_id=self._meta_agent.id, section=label, value=text)

            return {
                "success": True,
                "message": f"Core memory block '{label}' updated successfully"
                + (f" for user {user_id}" if user_id else ""),
            }
        except Exception as e:
            return {
                "success": False,
                "message": f"Error updating core memory: {str(e)}",
            }
