# -*- coding: utf-8 -*-
import asyncio
import datetime
import logging
import os
from typing import Optional, Type, List, Sequence, Tuple, Any

from agentscope.agent import ReActAgent
from agentscope.formatter import OpenAIChatFormatter
from agentscope.memory import InMemoryMemory
from agentscope.message import Msg
from agentscope.model import OpenAIChatModel
from agentscope.tool import Toolkit
from pydantic import BaseModel

from .prompt import (
    build_system_prompt_from_working_dir,
    build_bootstrap_guidance,
)
from .skills_manager import (
    ensure_skills_initialized,
    get_working_skills_dir,
    list_available_skills,
)
from .tools import (
    execute_shell_command,
    read_file,
    write_file,
    edit_file,
    send_file_to_user,
    desktop_screenshot,
    browser_use,
    create_memory_search_tool,
)
from .utils import (
    process_file_and_media_blocks_in_message,
    count_message_tokens,
    check_valid_messages,
    is_first_user_interaction,
    prepend_to_message_content,
)
from ..agents.memory import MemoryManager
from ..config import load_config
from ..constant import (
    MEMORY_COMPACT_THRESHOLD,
    MEMORY_COMPACT_KEEP_RECENT,
    WORKING_DIR,
)
from ..providers import get_active_llm_config

logger = logging.getLogger(__name__)


def create_file_block_support_formatter(base_formatter_class):
    """Factory function to add file block support to any Formatter class."""

    class FileBlockSupportFormatter(base_formatter_class):
        @staticmethod
        def convert_tool_result_to_string(
            output: str | List[dict],
        ) -> tuple[str, Sequence[Tuple[str, dict]]]:
            """Extend parent class to support file blocks.

            Uses try-first strategy for compatibility.
            """
            if isinstance(output, str):
                return output, []

            # Try parent class method first
            try:
                return base_formatter_class.convert_tool_result_to_string(
                    output,
                )
            except ValueError as e:
                if "Unsupported block type: file" not in str(e):
                    raise

                # Handle output containing file blocks
                textual_output = []
                multimodal_data = []

                for block in output:
                    if not isinstance(block, dict) or "type" not in block:
                        raise ValueError(
                            f"Invalid block: {block}, "
                            "expected a dict with 'type' key",
                        ) from e

                    if block["type"] == "file":
                        file_path = block.get("path", "") or block.get(
                            "url",
                            "",
                        )
                        file_name = block.get("name", file_path)

                        textual_output.append(
                            f"The returned file '{file_name}' "
                            f"can be found at: {file_path}",
                        )
                        multimodal_data.append((file_path, block))
                    else:
                        # Delegate other block types to parent class
                        (
                            text,
                            data,
                        ) = base_formatter_class.convert_tool_result_to_string(
                            [block],
                        )
                        textual_output.append(text)
                        multimodal_data.extend(data)

                if len(textual_output) == 0:
                    return "", multimodal_data
                elif len(textual_output) == 1:
                    return textual_output[0], multimodal_data
                else:
                    return (
                        "\n".join("- " + _ for _ in textual_output),
                        multimodal_data,
                    )

    FileBlockSupportFormatter.__name__ = (
        f"FileBlockSupport{base_formatter_class.__name__}"
    )
    return FileBlockSupportFormatter


# Create formatter with file block support
ClawdAgentFormatter = create_file_block_support_formatter(
    OpenAIChatFormatter,
)


class ClawdAgent(ReActAgent):
    def __init__(
        self,
        env_context: Optional[str] = None,
        enable_memory_manager: bool = False,
        mcp_clients: Optional[List[Any]] = None,
        memory_manager: MemoryManager | None = None,
    ):
        """Initialize ClawdAgent.

        Args:
            env_context: Optional environment context
            enable_memory_manager: Whether to enable memory manager
        """
        toolkit = Toolkit()
        self._mcp_clients = mcp_clients or []
        toolkit.register_tool_function(execute_shell_command)
        toolkit.register_tool_function(read_file)
        toolkit.register_tool_function(write_file)
        toolkit.register_tool_function(edit_file)
        toolkit.register_tool_function(browser_use)
        # toolkit.register_tool_function(append_file)
        toolkit.register_tool_function(desktop_screenshot)
        toolkit.register_tool_function(send_file_to_user)

        # Check skills initialization
        ensure_skills_initialized()

        working_skills_dir = get_working_skills_dir()
        available_skills = list_available_skills()

        for skill_name in available_skills:
            skill_dir = working_skills_dir / skill_name
            if skill_dir.exists():
                try:
                    toolkit.register_agent_skill(str(skill_dir))
                    logger.info("Registered skill: %s", skill_name)
                except Exception as e:
                    logger.error(
                        "Failed to register skill '%s': %s",
                        skill_name,
                        e,
                    )

        sys_prompt = build_system_prompt_from_working_dir()
        if env_context is not None:
            sys_prompt = env_context + "\n\n" + sys_prompt

        # Resolve model / api_key / base_url from the active LLM slot
        llm_cfg = get_active_llm_config()
        if llm_cfg and llm_cfg.api_key:
            model_name = llm_cfg.model or "qwen3-max"
            api_key = llm_cfg.api_key
            base_url = llm_cfg.base_url
        else:
            logger.warning(
                "No active LLM configured — "
                "falling back to DASHSCOPE_API_KEY env var",
            )
            model_name = "qwen3-max"
            api_key = os.getenv("DASHSCOPE_API_KEY", "")
            base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"

        super().__init__(
            name="Friday",
            model=OpenAIChatModel(
                model_name,
                api_key=api_key,
                stream=True,
                client_kwargs={"base_url": base_url},
            ),
            sys_prompt=sys_prompt,
            toolkit=toolkit,
            memory=InMemoryMemory(),
            formatter=ClawdAgentFormatter(),
        )
        self.memory_manager = memory_manager

        # Register memory_search tool if memory_manager is available
        if self.memory_manager is not None:
            memory_search_tool = create_memory_search_tool(self.memory_manager)
            self.toolkit.register_tool_function(memory_search_tool)
            logger.info("Registered memory_search tool")

        self.register_instance_hook(
            hook_type="pre_reasoning",
            hook_name="bootstrap_hook",
            hook=ClawdAgent._pre_reasoning_bootstrap_hook,
        )
        logger.info("Registered bootstrap hook")

        if enable_memory_manager and self.memory_manager is not None:
            self.register_instance_hook(
                hook_type="pre_reasoning",
                hook_name="memory_compact_hook",
                hook=ClawdAgent._pre_reasoning_compact_hook,
            )
            logger.info("Registered memory compaction hook")

        self.summary_task: Optional[asyncio.Task] = None
        self._bootstrap_checked = False

    async def register_mcp_clients(self) -> None:
        """Register MCP clients on this agent's toolkit after construction."""
        for client in self._mcp_clients:
            await self.toolkit.register_mcp_client(client)

    async def _pre_reasoning_bootstrap_hook(  # pylint: disable=unused-argument
        self,
        kwargs: dict[str, Any],
    ) -> dict[str, Any] | None:
        """Check and load BOOTSTRAP.md on first user interaction."""
        if self._bootstrap_checked:
            return None

        self._bootstrap_checked = True

        try:
            bootstrap_path = WORKING_DIR / "BOOTSTRAP.md"
            if not bootstrap_path.exists():
                return None

            messages = await self.memory.get_memory()
            if not is_first_user_interaction(messages):
                return None

            config = load_config()
            language = config.agents.language
            bootstrap_content = bootstrap_path.read_text(encoding="utf-8")
            bootstrap_guidance = build_bootstrap_guidance(
                bootstrap_content,
                language,
            )

            logger.info(
                "Found BOOTSTRAP.md [%s], prepending guidance",
                language,
            )

            system_prompt_count = sum(
                1 for msg in messages if msg.role == "system"
            )
            for msg in messages[system_prompt_count:]:
                if msg.role == "user":
                    prepend_to_message_content(msg, bootstrap_guidance)
                    break

            logger.info("Bootstrap guidance prepended to first user message")

        except Exception as e:
            logger.error(
                "Failed to process bootstrap: %s",
                e,
                exc_info=True,
            )

        return None

    async def _pre_reasoning_compact_hook(  # pylint: disable=unused-argument
        self,
        kwargs: dict[str, Any],
    ) -> dict[str, Any] | None:
        """Pre-reasoning hook to check and compact memory if needed.

        This hook is called before each reasoning step. It extracts system
        prompt messages (consecutive system messages at the start) and recent
        messages, then counts tokens for the middle compactable messages only.
        If the token count exceeds the threshold, it triggers compaction.

        Memory structure:
            [System Prompt (preserved)] + [Compactable (counted)] +
            [Recent (preserved)]

        Args:
            kwargs: Input arguments to the _reasoning method (not modified)

        Returns:
            None
        """
        # Only compact if memory manager is enabled
        if self.memory_manager is None:
            return None

        try:
            messages = await self.memory.get_memory()

            # Extract system prompt (consecutive system messages at start)
            system_prompt_messages = []
            for msg in messages:
                if msg.role == "system":
                    system_prompt_messages.append(msg)
                else:
                    break

            # Get remaining messages after system prompt
            remaining_messages = messages[len(system_prompt_messages) :]

            # Skip if not enough messages to compact
            if len(remaining_messages) <= MEMORY_COMPACT_KEEP_RECENT:
                return None

            # ensure the messages_to_keep is valid
            keep_length = MEMORY_COMPACT_KEEP_RECENT
            while keep_length > 0 and not check_valid_messages(
                remaining_messages[-keep_length:],
            ):
                keep_length -= 1

            # Split into compactable and recent messages
            if keep_length > 0:
                messages_to_compact = remaining_messages[:-keep_length]
                messages_to_keep = remaining_messages[-keep_length:]
            else:
                messages_to_compact = remaining_messages
                messages_to_keep = []

            # Count tokens for compactable messages only
            prompt = await self.formatter.format(msgs=messages_to_compact)
            estimated_tokens = await count_message_tokens(prompt)

            # Check if the compactable part exceeds threshold
            if estimated_tokens > MEMORY_COMPACT_THRESHOLD:
                logger.info(
                    "Memory compaction triggered: estimated %d tokens "
                    "(threshold: %d), system_prompt_msgs: %d, "
                    "compactable_msgs: %d, keep_recent_msgs: %d",
                    estimated_tokens,
                    MEMORY_COMPACT_THRESHOLD,
                    len(system_prompt_messages),
                    len(messages_to_compact),
                    len(messages_to_keep),
                )

                current_date = datetime.datetime.now().strftime("%Y-%m-%d")

                # Wait for previous summary task if it exists and is still
                # running
                if (
                    self.summary_task is not None
                    and not self.summary_task.done()
                ):
                    logger.info(
                        "Waiting for previous summary task to complete",
                    )
                    try:
                        await self.summary_task
                    except Exception as e:
                        logger.warning(
                            "Previous summary task failed: %s",
                            e,
                            exc_info=True,
                        )

                self.summary_task = asyncio.create_task(
                    self.memory_manager.summary_memory(
                        messages=messages_to_compact,
                        date=current_date,
                    ),
                )
                compact_content: str = (
                    await self.memory_manager.compact_memory(
                        messages_to_summarize=messages_to_compact,
                    )
                )

                self.memory.update_compressed_summary(compact_content)
                updated_count = await self.memory.update_messages_mark(
                    new_mark="compacted",
                    msg_ids=[msg.id for msg in messages_to_compact],
                )
                logger.info(
                    "Marked %d messages as compacted",
                    updated_count,
                )

        except Exception as e:
            # todo: handle the exception
            logger.error(
                "Failed to compact memory in pre_reasoning hook: %s",
                e,
                exc_info=True,
            )

        return None

    async def reply(
        self,
        msg: Msg | list[Msg] | None = None,
        structured_model: Type[BaseModel] | None = None,
    ) -> Msg:
        """Override reply to process file and media blocks."""
        if msg is not None:
            await process_file_and_media_blocks_in_message(msg)
        return await super().reply(msg=msg, structured_model=structured_model)
