import asyncio
import json
from pathlib import Path
from uuid import uuid4

from .utils.llm import process_messages_for_store
from .utils.log import logger

_ALL_CONTEXTS = object()


# FIX: other memory implimentaion?
class Memory:
    """
    The Memory class is used to store the memory of the agent.

    Args:
        name: The name of the memory.

    Attributes:
        name: The name of the memory.
        id: The ID of the memory.
        extra_data: The extra data of the memory.
    """

    def __init__(self, name: str, file_path: str | None = None, persist_delay: float = 2.0):
        self.name = name
        self.id = str(uuid4())
        self._messages: list[dict] = []
        self.extra_data: dict[str, object] = {}
        
        # Debounced persistence (opt-in: only when file_path is set)
        self._file_path: str | None = file_path
        self._persist_delay = persist_delay
        self._persist_task: asyncio.Task | None = None
        self._dirty: bool = False

    def __getitem__(self, key: int | slice):
        """Get a message or slice of messages from the memory."""
        if isinstance(key, int):
            return self._messages[key]
        elif isinstance(key, slice):
            new_memory = Memory(self.name)
            new_memory._messages = self._messages[key]
            new_memory.extra_data = self.extra_data.copy()
            return new_memory
        else:
            raise ValueError(f"Invalid key: {key}")

    def save(self, file_path: str):
        """
        Save the memory to a file.

        Args:
            file_path: The path to save the memory to.
        """
        with open(file_path, "w", encoding="utf-8") as f:
            json.dump(
                {
                    "id": self.id,
                    "name": self.name,
                    "messages": self._messages,
                    "extra_data": self.extra_data,
                },
                f,
                indent=4,
            )

    @classmethod
    def load(cls, file_path: str):
        """
        Load the memory from a file.

        Args:
            file_path: The path to load the memory from.
        """
        with open(file_path, "r", encoding="utf-8") as f:
            data = json.load(f)
            memory = cls(data["name"])
            memory.id = data["id"]
            memory._messages = data["messages"]
            memory.extra_data = data.get("extra_data", {})
            return memory

    def clear(self):
        """
        Clear the memory.
        """
        self._messages = []
        self.extra_data = {}

    def add_messages(self, messages: list[dict]):
        """
        Add messages to the memory.

        Args:
            messages: The messages to add to the memory.
        """
        messages = process_messages_for_store(messages)
        self._messages.extend(messages)
        self._schedule_persist()  # Trigger debounced auto-persistence

    def mark_dirty(self):
        """Mark memory as dirty to trigger delayed persistence.
        
        Call this after modifying extra_data or other fields that don't
        automatically trigger persistence (unlike add_messages).
        """
        self._schedule_persist()
    
    def _schedule_persist(self):
        """Schedule debounced persistence (non-blocking).
        
        Only schedules if file_path is set (opt-in behavior).
        Uses debounce window to batch multiple writes.
        """
        if not self._file_path:
            return
        
        self._dirty = True
        
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            # No event loop running, skip auto-persist
            # (Will be persisted by MemoryManager.save() later)
            return
        
        # Cancel existing scheduled task (debounce)
        if self._persist_task and not self._persist_task.done():
            self._persist_task.cancel()
        
        async def _delayed_persist():
            await asyncio.sleep(self._persist_delay)
            self._do_persist()
        
        self._persist_task = loop.create_task(_delayed_persist())

    def _do_persist(self):
        """Actually write memory to disk."""
        if self._file_path and self._dirty:
            try:
                self.save(self._file_path)
                self._dirty = False
                logger.debug(f"Memory '{self.name}' auto-persisted to {self._file_path}")
            except Exception as e:
                logger.error(f"Failed to auto-persist memory '{self.name}': {e}")

    async def flush(self):
        """Force immediate persistence, cancel any pending debounce.
        
        Use this for graceful shutdown to ensure all data is saved.
        """
        if self._persist_task and not self._persist_task.done():
            self._persist_task.cancel()
            try:
                await self._persist_task
            except asyncio.CancelledError:
                pass
        self._do_persist()

    def get_messages(self, execution_context_id=_ALL_CONTEXTS, for_llm: bool = True) -> list[dict]:
        """
        Get the messages from the memory.

        Args:
            execution_context_id: Filter by execution context ID.
            for_llm: If True, applies processing for LLM consumption:
                     - Excludes internal system messages
                     - Handles compression truncation
                     - Normalizes compression messages to user role

        Returns:
            The messages from the memory.
        """
        if execution_context_id is _ALL_CONTEXTS:
            messages = list(self._messages)
        elif execution_context_id is None:
            messages = [
                msg
                for msg in self._messages
                if msg.get("execution_context_id") is None
            ]
        else:
            messages = [
                msg
                for msg in self._messages
                if msg.get("execution_context_id") == execution_context_id
            ]

        if not for_llm:
            return messages

        # --- LLM Processing Logic ---
        from copy import deepcopy

        # 1. Skip system messages (managed externally/prepended)
        filtered = [m for m in messages if m.get("role") != "system"]

        # 2. Find last compression message for truncation
        last_compression_idx = -1
        for i, msg in enumerate(filtered):
            if msg.get("role") == "compression":
                last_compression_idx = i

        # 3. Apply truncation
        if last_compression_idx >= 0:
            filtered = filtered[last_compression_idx:]

        # 4. Final processing (convert compression to user, cleanup metadata)
        final_messages = []
        for msg in filtered:
            msg_copy = deepcopy(msg)

            # Metadata is preserved here for cost tracking. 
            # It will be removed by call_llm_provider before sending to API.
            # if "_metadata" in msg_copy:
            #    del msg_copy["_metadata"]

            if msg_copy.get("role") == "compression":
                msg_copy["role"] = "user"
                # Ensure content is string
                if not isinstance(msg_copy.get("content"), str):
                    msg_copy["content"] = str(msg_copy.get("content", ""))

            final_messages.append(msg_copy)

        return final_messages

    def get_user_turns(self) -> list[tuple[int, dict]]:
        """
        Get all user turns from the memory.

        Returns:
            A list of tuples containing the index and the message.
        """
        return [
            (i, msg)
            for i, msg in enumerate(self._messages)
            if msg.get("role") == "user"
        ]

    def revert_to_message(self, index: int):
        """
        Revert the memory to a specific message index.
        The message at the given index will be REMOVED, along with all subsequent messages.
        Effectively rollback to the state BEFORE that message.

        Args:
            index: The index of the message to revert to (inclusive).
        """
        if 0 <= index < len(self._messages):
            self._messages = self._messages[:index]
            self._dirty = True
            # Force immediate save for revert to prevent race conditions or confusion
            if self._file_path:
                self.save(self._file_path)
        else:
            raise ValueError(f"Invalid message index: {index}")

    def cleanup(self):
        """Cleanup the memory after the agent is interrupted."""
        while True:
            if len(self._messages) == 0:
                break

            last_message = self._messages[-1]
            if "role" not in last_message:
                logger.debug(
                    f"Popping message: {last_message}, len={len(self._messages)}"
                )
                self._messages.pop()
                continue
            if last_message["role"] == "user":
                logger.debug(
                    f"Popping user message: {last_message}, len={len(self._messages)}"
                )
                self._messages.pop()
                continue
            if last_message.get("content") is None:
                logger.debug(
                    f"Popping message: {last_message}, len={len(self._messages)}"
                )
                self._messages.pop()
                continue
            if last_message["role"] == "assistant":
                if last_message["tool_calls"]:
                    last_message["tool_calls"] = None
            if last_message["role"] == "assistant":
                break
            if last_message["role"] == "tool":
                break

    def _fix_corrupted_messages(self):
        """Remove corrupted messages (missing role and useless).
        
        Removes messages that only contain partial data (e.g. {'agent_name': ...})
        resulting from failed model calls.
        """
        original_len = len(self._messages)
        self._messages = [
            msg for msg in self._messages 
            if msg.get("role") is not None
        ]
        removed_count = original_len - len(self._messages)
        if removed_count > 0:
            logger.warning(
                f"Removed {removed_count} corrupted messages (missing role) from memory '{self.name}'"
            )
            self._dirty = True

    def _fix_orphaned_tool_calls(self):
        """Add placeholder responses for orphaned tool_calls and fix context IDs.

        1. Inserts [INTERNAL_ERROR] tool responses for any tool_call that lacks
           a corresponding tool message.
        2. Updates existing placeholder messages to ensure they match key metadata
           (execution_context_id, agent_name) of the parent assistant message.
        """
        import time
        from copy import deepcopy

        # Helper to find tool message for a tool_call_id
        tool_msgs_map = {
            msg.get("tool_call_id"): msg
            for msg in self._messages
            if msg.get("role") == "tool" and msg.get("tool_call_id")
        }

        insertions = []  # (index, placeholder_message)

        for i, msg in enumerate(self._messages):
            if msg.get("role") == "assistant" and msg.get("tool_calls"):
                parent_context_id = msg.get("execution_context_id")
                parent_agent_name = msg.get("agent_name")

                # Check each tool call in this assistant message
                for tc in msg["tool_calls"]:
                    tc_id = tc.get("id")
                    if not tc_id:
                        continue

                    existing_tool_msg = tool_msgs_map.get(tc_id)

                    if existing_tool_msg:
                        continue
                    else:
                        # Create Phase: Insert missing tool response
                        placeholder = {
                            "role": "tool",
                            "tool_call_id": tc_id,
                            "tool_name": tc.get("function", {}).get("name", "unknown"),
                            "content": "[INTERNAL_ERROR] Session interrupted - tool execution incomplete",
                            "id": str(uuid4()),
                            "timestamp": time.time(),
                            "_recovered": True,
                        }
                        # Propagate parent metadata
                        if parent_context_id:
                            placeholder["execution_context_id"] = parent_context_id
                        if parent_agent_name:
                            placeholder["agent_name"] = parent_agent_name

                        # Use parent's metadata structure if useful (optional, but good for tracking)
                        if "_metadata" in msg:
                             # Minimal metadata copy if needed
                             pass

                        insertions.append((i + 1, placeholder))
                        # Update map to prevent duplicates if multiple refs exist (unlikely)
                        tool_msgs_map[tc_id] = placeholder

        # Insert in reverse order to maintain correct indices
        for idx, placeholder in reversed(insertions):
            self._messages.insert(idx, placeholder)

        if insertions:
            logger.info(
                f"Fixed memory '{self.name}': {len(insertions)} orphaned tool_call(s) inserted."
            )
            self._schedule_persist()

    def ensure_fixed(self):
        """Ensure memory is fixed (idempotent).

        This method can be called multiple times safely. It will only fix
        corrupted messages and orphaned tool calls once, then persist the
        result to disk.

        Use this method when you need to ensure the memory is in a valid
        state for agent execution or LLM API calls.
        """
        if not getattr(self, '_orphans_fixed', False):
            self._fix_corrupted_messages()
            self._fix_orphaned_tool_calls()
            self._orphans_fixed = True
            # Persist the fixed state immediately
            if self._dirty:
                self._schedule_persist()


DEFAULT_CHAT_NAME = "New Chat"


class MemoryManager:
    """
    The MemoryManager class is used to manage multiple memories.

    Args:
        path: The path to the memory files.

    Attributes:
        path: The path to the memory files.
        memory_store: The in-RAM store of the memories.
    """

    def __init__(self, path: str | Path):
        self.path = Path(path)
        self.memory_store: dict[str, Memory] = {}
        self.load()

    def new_memory(self, name: str | None = None) -> Memory:
        """
        Create a new memory.

        Args:
            name: The name of the memory.
        """
        if name is None:
            name = DEFAULT_CHAT_NAME
        memory = Memory(name)
        self.memory_store[memory.id] = memory
        # Enable auto-persistence for managed memories
        memory._file_path = str(self.path / f"{memory.id}.json")
        return memory

    def get_memory(self, id: str, auto_fix: bool = False) -> Memory:
        """
        Get a memory by its ID.

        Args:
            id: The ID of the memory.
            auto_fix: If True, automatically fix corrupted messages and orphaned
                     tool calls. This is required for agent execution and LLM API
                     calls. If False (default), return the memory as-is for
                     read-only operations like frontend queries.

        Raises:
            KeyError: If the memory with the given ID does not exist.

        Performance Note:
            Setting auto_fix=False skips the fix operations, which can save
            5-10ms for read-only queries. The fix will be applied automatically
            when the memory is used for agent execution.
        """
        if id not in self.memory_store:
            raise KeyError(f"Chat '{id}' not found. It may have been deleted.")
        memory = self.memory_store[id]

        # Only fix if explicitly requested (for agent execution)
        if auto_fix:
            memory.ensure_fixed()

        return memory

    def delete_memory(self, id: str):
        """
        Delete a memory by its ID.

        Args:
            id: The ID of the memory.
        """
        # Delete from memory store
        del self.memory_store[id]
        
        # Immediately delete the file from disk
        file_path = self.path / f"{id}.json"
        if file_path.exists():
            file_path.unlink()
            logger.debug(f"Deleted memory file: {file_path}")

    def list_memories(self):
        """
        List all the memories.

        Returns:
            The list of the memories.
        """
        return list(self.memory_store.keys())

    def save_one(self, memory_id: str):
        """
        Save a single memory to the file system.
        
        Args:
            memory_id: The ID of the memory to save.
        """
        memory = self.memory_store.get(memory_id)
        if memory:
            memory.save(str(self.path / f"{memory.id}.json"))
        else:
            logger.warning(f"Memory {memory_id} not found in memory store, cannot save")
    
    def save(self):
        """
        Save all the memories to the file system.
        """
        for memory in self.memory_store.values():
            memory.save(str(self.path / f"{memory.id}.json"))
        for file in self.path.glob("*.json"):
            if file.stem not in self.memory_store:
                file.unlink()

    def load(self):
        """
        Load all the memories from the file system.
        """
        if not self.path.exists():
            self.path.mkdir(parents=True)
        for file in self.path.glob("*.json"):
            try:
                memory = Memory.load(str(file))
                # Enable auto-persistence for managed memories
                memory._file_path = str(file)
                logger.debug(f"Loaded memory: {memory.name} from {file}")
                self.memory_store[memory.id] = memory
                if memory.extra_data.get("running"):
                    memory.extra_data["running"] = False
            except Exception as e:
                logger.error(f"Failed to load memory from {file}: {e}")

    def update_memory_name(self, memory_id: str, name: str):
        """
        Update the name of a memory.

        Args:
            memory_id: The ID of the memory.
            name: The new name of the memory.
        """
        # Read-only: updating memory name, no need to fix
        memory = self.get_memory(memory_id)
        memory.name = name
        self.save_one(memory_id)
