"""Agent execution loop implementation."""

import asyncio
import time
import uuid
from collections import Counter
from typing import AsyncIterator, Literal

from pydantic import BaseModel

from yier_agents.src.memory import CompactionConfig, InMemorySessionStore, SessionStore
from yier_agents.src.message import Message, LLMResponse
from yier_agents.src.skill import Skill, SkillCatalog, SkillRegistry
from yier_agents.src.tool import Tool, ToolContext, ToolOutput
from yier_agents.src.middleware import ToolMiddleware
from yier_agents.src.llm import LLM
from yier_agents.src.tools.skill_load import SkillLoadTool
from yier_agents.src.events import (
    AgentEventType,
    AgentEndEvent,
    LLMEndEvent,
    ToolCallStartEvent,
    ToolCallEndEvent,
    MessageCompactEvent,
)


class AgentResult(BaseModel):
    """Agent execution result."""
    session_id: str
    messages: list[Message]
    final_message: Message
    finish_reason: Literal["stop", "tool_calls", "length", "max_iterations"]


class Agent:
    """Agent with finish_reason-driven execution loop and session memory."""

    def __init__(
        self,
        llm: LLM,
        tools: list[Tool] | None = None,
        max_iterations: int = 20,
        system_prompt: str | None = None,
        verbose: bool = False,
        enable_memory: bool = True,
        skills: list[Skill] | None = None,
        skill_catalog: SkillCatalog | None = None,
        enable_skill_tool: bool = True,
        session_store: SessionStore | None = None,
        compaction_config: CompactionConfig | None = None,
    ):
        """Initialize Agent.

        Args:
            llm: LLM instance for completion
            tools: List of available tools
            max_iterations: Maximum number of iterations to prevent infinite loops
            system_prompt: Optional system prompt to prepend to conversation
            verbose: Whether to print debug information during execution
            enable_memory: Whether to enable session memory across multiple run() calls
            skills: Optional reusable skill bundles that contribute tools and prompts
            skill_catalog: Optional catalog of lazily loadable file-based skills
            enable_skill_tool: Whether to expose the built-in lazy skill loading tool
            session_store: Optional storage backend for session memory
            compaction_config: Optional settings for long-context compaction
        """
        self.llm = llm
        self.skill_registry = SkillRegistry(skills or [])
        self.skill_catalog = skill_catalog
        self.enable_skill_tool = enable_skill_tool
        self.tools = self._build_tools(tools)
        self.max_iterations = max_iterations
        self.system_prompt = self._build_system_prompt(system_prompt)
        self.verbose = verbose
        self.enable_memory = enable_memory
        self.session_store = session_store or InMemorySessionStore()
        self.compaction_config = compaction_config or CompactionConfig()

        self.middleware = ToolMiddleware(tools=self.tools)

    def _build_tools(self, tools: list[Tool] | None) -> list[Tool]:
        collected_tools = self.skill_registry.collect_tools(tools)

        if self.enable_skill_tool and self.skill_catalog and not self.skill_catalog.is_empty():
            collected_tools.append(SkillLoadTool(self.skill_catalog))

        deduped_tools: dict[str, Tool] = {}
        for tool in collected_tools:
            deduped_tools[tool.name] = tool
        return list(deduped_tools.values())

    def _build_system_prompt(self, base_prompt: str | None) -> str | None:
        prompt_parts = [
            part.strip()
            for part in [self.skill_registry.build_system_prompt(base_prompt)]
            if part and part.strip()
        ]

        if self.enable_skill_tool and self.skill_catalog and not self.skill_catalog.is_empty():
            prompt_parts.append(
                "\n".join(
                    [
                        "Skills provide specialized instructions and workflows for specific tasks.",
                        "Use the skill_load tool to load a skill when a task matches its description.",
                        self.skill_catalog.format(verbose=True),
                    ]
                )
            )

        if not prompt_parts:
            return None

        return "\n\n".join(prompt_parts)

    async def run_stream(
        self,
        user_message: str,
        session_id: str
    ) -> AsyncIterator[AgentEventType]:
        """Run agent with streaming events.

        Args:
            user_message: User's input message
            session_id: Session identifier

        Yields:
            AgentEventType events as execution progresses
        """
        messages = self._load_or_initialize_messages(session_id)

        # Emit start event
        # yield AgentStartEvent(
        #     session_id=session_id,
        #     timestamp=time.time(),
        #     user_message=user_message,
        #     iteration=0
        # )

        # Add user message
        messages.append(Message(role="user", content=user_message))

        # Execution loop
        iterations = 0
        while iterations < self.max_iterations:
            iterations += 1

            original_count = len(messages)
            messages = await self._maybe_compact(messages)
            if len(messages) != original_count:
                yield MessageCompactEvent(
                    session_id=session_id,
                    timestamp=time.time(),
                    iteration=iterations,
                    original_count=original_count,
                    compacted_count=len(messages)
                )

            # Call LLM
            # yield LLMStartEvent(
            #     session_id=session_id,
            #     timestamp=time.time(),
            #     iteration=iterations
            # )

            tool_schemas = [tool.to_openai_schema() for tool in self.tools] if self.tools else None
            response = await self.llm.complete(messages=messages, tools=tool_schemas)

            # Add assistant message to history
            messages.append(response.message)

            # Emit LLM end event
            yield LLMEndEvent(
                session_id=session_id,
                timestamp=time.time(),
                iteration=iterations,
                message=response.message,
                finish_reason=response.finish_reason,
                usage=response.usage
            )

            # Handle finish reason
            if response.finish_reason == "stop":
                # Save session memory
                if self.enable_memory:
                    self.session_store.save(session_id, messages)
                    # yield MemoryEvent(
                    #     session_id=session_id,
                    #     timestamp=time.time(),
                    #     action="save",
                    #     message_count=len(messages)
                    # )

                # Emit end event
                yield AgentEndEvent(
                    session_id=session_id,
                    timestamp=time.time(),
                    finish_reason="stop",
                    total_iterations=iterations
                )
                return

            elif response.finish_reason == "length":
                # Emit compact event
                original_count = len(messages)
                messages = await self._maybe_compact(messages, force=True)
                yield MessageCompactEvent(
                    session_id=session_id,
                    timestamp=time.time(),
                    iteration=iterations,
                    original_count=original_count,
                    compacted_count=len(messages)
                )
                continue

            elif response.finish_reason == "tool_calls":
                # Execute tools
                if response.message.tool_calls:
                    message_id = str(uuid.uuid4())

                    for tc in response.message.tool_calls:
                        # Emit tool start event
                        yield ToolCallStartEvent(
                            session_id=session_id,
                            timestamp=time.time(),
                            iteration=iterations,
                            tool_name=tc.name,
                            tool_call_id=tc.id,
                            arguments=tc.arguments
                        )

                        # Execute tool
                        ctx = ToolContext(
                            session_id=session_id,
                            message_id=message_id,
                            call_id=tc.id
                        )
                        try:
                            result = await self.middleware.execute_tool_call(tc, ctx)

                            # Emit tool end event
                            yield ToolCallEndEvent(
                                session_id=session_id,
                                timestamp=time.time(),
                                iteration=iterations,
                                tool_name=tc.name,
                                tool_call_id=tc.id,
                                result=result.content,
                                is_error=result.is_error
                            )

                            # Add tool result to messages
                            messages.append(Message(
                                role="tool",
                                content=result.content,
                                tool_call_id=result.call_id
                            ))

                        except Exception as e:
                            # Emit error event
                            from yier_agents.src.events import ErrorEvent
                            yield ErrorEvent(
                                session_id=session_id,
                                timestamp=time.time(),
                                iteration=iterations,
                                error_message=str(e),
                                error_type=type(e).__name__
                            )

                continue

        # Max iterations reached
        if self.enable_memory:
            self.session_store.save(session_id, messages)
            # yield MemoryEvent(
            #     session_id=session_id,
            #     timestamp=time.time(),
            #     action="save",
            #     message_count=len(messages)
            # )

        yield AgentEndEvent(
            session_id=session_id,
            timestamp=time.time(),
            finish_reason="max_iterations",
            total_iterations=iterations
        )

    def _log(self, message: str, level: str = "info"):
        """Print debug message if verbose is enabled.

        Args:
            message: Message to print
            level: Log level (info, success, warning, tool)
        """
        if not self.verbose:
            return

        # ANSI color codes
        colors = {
            "info": "\033[36m",      # Cyan
            "success": "\033[32m",   # Green
            "warning": "\033[33m",   # Yellow
            "tool": "\033[35m",      # Magenta
            "reset": "\033[0m"
        }

        # Icons
        icons = {
            "info": "ℹ",
            "success": "✓",
            "warning": "⚠",
            "tool": "⚙"
        }

        color = colors.get(level, colors["info"])
        icon = icons.get(level, icons["info"])
        reset = colors["reset"]

        print(f"{color}{icon} {message}{reset}")

    async def run(self, user_message: str, session_id: str) -> AgentResult:
        """Run agent execution loop.

        Args:
            user_message: User's input message
            session_id: Session identifier

        Returns:
            AgentResult containing conversation history and final state
        """
        messages = self._load_or_initialize_messages(session_id)
        if self.enable_memory:
            self._log(f"Loaded {len(messages)} messages from session memory", "info")
        elif self.system_prompt:
            self._log("Added system prompt", "info")

        self._log("=" * 60, "info")
        self._log(f"Starting Agent execution (session: {session_id[:8]}...)", "info")
        self._log("=" * 60, "info")

        # Add user message
        messages.append(Message(role="user", content=user_message))
        self._log(f"User: {user_message[:100]}{'...' if len(user_message) > 100 else ''}", "info")

        # Execution loop
        iterations = 0
        while iterations < self.max_iterations:
            iterations += 1
            self._log(f"\n{'─' * 60}", "info")
            self._log(f"Iteration {iterations}/{self.max_iterations}", "info")
            self._log(f"{'─' * 60}", "info")

            original_count = len(messages)
            messages = await self._maybe_compact(messages)
            if len(messages) != original_count:
                self._log(
                    f"Compacted message history ({original_count} -> {len(messages)})",
                    "warning",
                )

            # Call LLM
            self._log("Calling LLM...", "info")
            tool_schemas = [tool.to_openai_schema() for tool in self.tools] if self.tools else None
            response = await self.llm.complete(messages=messages, tools=tool_schemas)

            # Add assistant message to history
            messages.append(response.message)

            # Log LLM response
            usage_str = f"{response.usage} tokens" if response.usage else "unknown tokens"
            self._log(f"LLM response received (finish_reason: {response.finish_reason}, usage: {usage_str})", "success")

            if response.message.reasoning_content:
                reasoning_preview = response.message.reasoning_content[:150]
                self._log(f"💭 Reasoning: {reasoning_preview}{'...' if len(response.message.reasoning_content) > 150 else ''}", "info")

            if response.message.content:
                content_preview = response.message.content[:150]
                self._log(f"Content: {content_preview}{'...' if len(response.message.content) > 150 else ''}", "info")

            # Handle finish reason
            if response.finish_reason == "stop":
                # Normal completion
                self._log("Agent completed successfully", "success")
                self._log("=" * 60, "info")

                # Save session memory
                if self.enable_memory:
                    self.session_store.save(session_id, messages)
                    self._log(f"Saved {len(messages)} messages to session memory", "info")

                return AgentResult(
                    session_id=session_id,
                    messages=messages,
                    final_message=response.message,
                    finish_reason="stop"
                )

            elif response.finish_reason == "length":
                # Token limit reached - try to compact and continue
                self._log("Token limit reached, attempting to compact messages", "warning")
                messages = await self._maybe_compact(messages, force=True)
                continue

            elif response.finish_reason == "tool_calls":
                # Execute tools
                if response.message.tool_calls:
                    tool_count = len(response.message.tool_calls)
                    self._log(f"Executing {tool_count} tool call(s)...", "tool")

                    for i, tc in enumerate(response.message.tool_calls, 1):
                        args_str = str(tc.arguments) if isinstance(tc.arguments, dict) else tc.arguments
                        self._log(f"  [{i}/{tool_count}] {tc.name}({args_str[:50]}{'...' if len(args_str) > 50 else ''})", "tool")

                    message_id = str(uuid.uuid4())
                    tool_results = await self._execute_tools(
                        tool_calls=response.message.tool_calls,
                        session_id=session_id,
                        message_id=message_id
                    )

                    # Add tool result messages
                    self._log(f"Received {len(tool_results)} tool result(s)", "success")
                    for i, result in enumerate(tool_results, 1):
                        result_preview = result.content[:100] if result.content else "No content"
                        self._log(f"  [{i}] {result_preview}{'...' if len(result.content) > 100 else ''}", "tool")

                        messages.append(Message(
                            role="tool",
                            content=result.content,
                            tool_call_id=result.call_id
                        ))
                continue

        # Max iterations reached
        self._log(f"Max iterations ({self.max_iterations}) reached", "warning")
        self._log("=" * 60, "info")

        # Save session memory
        if self.enable_memory:
            self.session_store.save(session_id, messages)
            self._log(f"Saved {len(messages)} messages to session memory", "info")

        return AgentResult(
            session_id=session_id,
            messages=messages,
            final_message=messages[-1],
            finish_reason="max_iterations"
        )

    async def _execute_tools(
        self,
        tool_calls: list,
        session_id: str,
        message_id: str
    ) -> list:
        """Execute tool calls concurrently.

        Args:
            tool_calls: List of ToolCall objects
            session_id: Session identifier
            message_id: Message identifier

        Returns:
            List of ToolResult objects
        """
        tasks = []
        for tool_call in tool_calls:
            ctx = ToolContext(
                session_id=session_id,
                message_id=message_id,
                call_id=tool_call.id
            )
            task = self.middleware.execute_tool_call(tool_call, ctx)
            tasks.append(task)

        # Execute all tools concurrently
        results = await asyncio.gather(*tasks)
        return results

    async def _maybe_compact(
        self,
        messages: list[Message],
        force: bool = False,
    ) -> list[Message]:
        """Compact messages when token limit is reached.

        This is an extension point for future message compression strategies.
        Older history is summarized into a single system message while keeping
        the most recent turns intact.

        Args:
            messages: Current message history

        Returns:
            Potentially compacted message history
        """
        if not self.compaction_config.enabled:
            return messages

        if not force and len(messages) < self.compaction_config.trigger_message_count:
            return messages

        prefix_messages, remaining_messages = self._split_prefix_messages(messages)
        preserve_count = self.compaction_config.preserve_recent_messages

        if len(remaining_messages) <= preserve_count:
            return messages

        summary_source = remaining_messages[:-preserve_count]
        recent_messages = remaining_messages[-preserve_count:]

        if not summary_source:
            return messages

        summary_message = await self._build_summary_message(summary_source)
        return prefix_messages + [summary_message] + recent_messages

    def _load_or_initialize_messages(self, session_id: str) -> list[Message]:
        if self.enable_memory:
            stored_messages = self.session_store.load(session_id)
            if stored_messages is not None:
                return stored_messages

        messages: list[Message] = []
        if self.system_prompt:
            messages.append(Message(role="system", content=self.system_prompt))
        return messages

    def _split_prefix_messages(
        self,
        messages: list[Message],
    ) -> tuple[list[Message], list[Message]]:
        prefix_messages: list[Message] = []
        remaining_messages = messages.copy()

        while (
            remaining_messages
            and remaining_messages[0].role == "system"
            and not self._is_summary_message(remaining_messages[0])
        ):
            prefix_messages.append(remaining_messages.pop(0))

        return prefix_messages, remaining_messages

    def _is_summary_message(self, message: Message) -> bool:
        return bool(
            message.role == "system"
            and message.content
            and message.content.startswith("<conversation_summary>")
        )

    async def _build_summary_message(self, messages: list[Message]) -> Message:
        summary_prompt = (
            "Summarize the conversation history for continued agent execution.\n"
            "Keep confirmed facts, user preferences, unfinished work, important tool results, "
            "and open questions. Be concise and do not invent new details."
        )
        transcript = self._format_messages_for_summary(messages)

        response = await self.llm.complete(
            messages=[
                Message(role="system", content=summary_prompt),
                Message(
                    role="user",
                    content=f"Conversation transcript:\n\n{transcript}",
                ),
            ],
            tools=None,
            max_tokens=self.compaction_config.summary_max_tokens,
            temperature=0.0,
        )

        summary_content = response.message.content or self._fallback_summary(messages)
        return Message(
            role="system",
            content=(
                "<conversation_summary>\n"
                f"{summary_content.strip()}\n"
                "</conversation_summary>"
            ),
        )

    def _format_messages_for_summary(self, messages: list[Message]) -> str:
        formatted_messages: list[str] = []

        for index, message in enumerate(messages, 1):
            sections = [f"{index}. role={message.role}"]

            if message.content:
                sections.append(f"content={message.content}")

            if message.reasoning_content:
                sections.append(f"reasoning={message.reasoning_content}")

            if message.tool_calls:
                tool_descriptions = [
                    f"{tool_call.name}({tool_call.arguments})"
                    for tool_call in message.tool_calls
                ]
                sections.append(f"tool_calls={'; '.join(tool_descriptions)}")

            if message.tool_call_id:
                sections.append(f"tool_call_id={message.tool_call_id}")

            formatted_messages.append(" | ".join(sections))

        return "\n".join(formatted_messages)

    def _fallback_summary(self, messages: list[Message]) -> str:
        summary_lines: list[str] = []

        for message in messages[-6:]:
            preview = (message.content or "").strip()
            if not preview:
                continue
            summary_lines.append(f"- {message.role}: {preview[:160]}")

        return "\n".join(summary_lines) or "No summary available."

    def clear_session(self, session_id: str) -> bool:
        """Clear memory for a specific session.

        Args:
            session_id: Session identifier to clear

        Returns:
            True if session was found and cleared, False otherwise
        """
        was_cleared = self.session_store.clear_session(session_id)
        if was_cleared:
            self._log(f"Cleared session memory for {session_id[:8]}...", "info")
        return was_cleared

    def clear_all_sessions(self):
        """Clear all session memory."""
        count = self.session_store.clear_all()
        self._log(f"Cleared {count} session(s) from memory", "info")

    def get_session_messages(self, session_id: str) -> list[Message] | None:
        """Get messages for a specific session.

        Args:
            session_id: Session identifier

        Returns:
            List of messages if session exists, None otherwise
        """
        return self.session_store.get_session_messages(session_id)

    def get_session_count(self) -> int:
        """Get total number of active sessions.

        Returns:
            Number of sessions in memory
        """
        return self.session_store.get_session_count()


class AgentToolParams(BaseModel):
    """Parameters for AgentTool."""
    prompt: str
    session_id: str | None = None


class AgentTool(Tool[AgentToolParams]):
    """Tool that wraps an Agent for sub-agent delegation.

    This allows a main agent to delegate tasks to a sub-agent,
    executing it as a tool call.
    """

    def __init__(
        self,
        name: str,
        description: str,
        agent: Agent,
        inherit_context: bool = True
    ):
        """Initialize AgentTool.

        Args:
            name: Tool name
            description: Tool description
            agent: Sub-agent to execute
            inherit_context: Whether to inherit context from parent (currently unused)
        """
        self.agent = agent
        self.inherit_context = inherit_context

        # Define the execute function
        async def execute_fn(params: AgentToolParams, ctx: ToolContext) -> ToolOutput:
            return await self._execute(params, ctx)

        # Initialize parent Tool class
        super().__init__(
            name=name,
            description=description,
            parameters=AgentToolParams,
            execute=execute_fn
        )

    async def _execute(self, params: AgentToolParams, ctx: ToolContext) -> ToolOutput:
        """Execute the sub-agent.

        Args:
            params: AgentToolParams containing prompt and optional session_id
            ctx: Tool execution context

        Returns:
            ToolOutput containing sub-agent result and metadata
        """
        # Determine sub-session ID
        if params.session_id:
            sub_session_id = params.session_id
        else:
            sub_session_id = f"{ctx.session_id}:sub:{uuid.uuid4()}"

        # Run sub-agent
        result = await self.agent.run(params.prompt, sub_session_id)

        # Collect tool summary
        tool_summary = self._summarize_tools(result.messages)

        # Construct output content
        content = result.final_message.content or ""

        # Add metadata section
        metadata_lines = [
            "",
            "<subagent_metadata>",
            f"Session ID: {result.session_id}",
            f"Finish Reason: {result.finish_reason}",
            f"Message Count: {len(result.messages)}",
        ]

        if tool_summary:
            metadata_lines.append(f"Tools Used: {tool_summary}")

        metadata_lines.append("</subagent_metadata>")

        content += "\n".join(metadata_lines)

        return ToolOutput(
            content=content,
            metadata={
                "sub_session_id": result.session_id,
                "finish_reason": result.finish_reason,
                "message_count": len(result.messages)
            }
        )

    def _summarize_tools(self, messages: list[Message]) -> str:
        """Summarize tool calls from messages.

        Args:
            messages: List of messages from agent execution

        Returns:
            Tool summary string in format "tool_name(count), ..."
        """
        tool_counter = Counter()

        for message in messages:
            if message.role == "assistant" and message.tool_calls:
                for tool_call in message.tool_calls:
                    tool_counter[tool_call.name] += 1

        if not tool_counter:
            return ""

        # Format as "tool_name(count), tool_name2(count2)"
        summary_parts = [f"{name}({count})" for name, count in sorted(tool_counter.items())]
        return ", ".join(summary_parts)
