"""SimpleAgent — the most fundamental Agent implementation for Fluxibly.

Manages a single LLM, tools, handoffs, agents-as-tools, and skills.
The caller owns conversation history — forward() receives the full
messages list each call.

forward() flow:
1. _resolve_context()        → extract prompt_params, tools, agents, handoffs, skills
2. _prepare(context_*)       → resolve new references, rebuild all_tools
3. _prepare_prompts()        → process templates + inject skill catalog
4. _agent_forward_loop()     → LLM call + agentic tool-call loop (with handoff/skill/monitoring)
5. finalize_response()       → log output, finalize monitoring, return
"""

from __future__ import annotations

import time
from pathlib import Path
from typing import Any

from fluxibly.agent.base import (
    AgentConfig,
    AgentResponse,
    BaseAgent,
    HandoffConfig,
)
from fluxibly.agent.skills import (
    SkillMetadata,
    build_skill_catalog,
    load_skill_metadata,
)
from fluxibly.agent.utils.context import resolve_context
from fluxibly.agent.utils.data import (
    collect_requirements,
    inject_data_prompt,
    inject_sandbox_files_prompt,
    stage_skill,
)
from fluxibly.agent.utils.execution import finalize_response
from fluxibly.agent.utils.monitoring import (
    AgentMonitoringSetup,
    fail_agent_monitoring,
    llm_forward,
    monitored_tool_call,
    setup_agent_monitoring,
)
from fluxibly.agent.utils.prompts import prepare_prompts
from fluxibly.agent.utils.resources import resolve_entry
from fluxibly.agent.utils.tools import (
    apply_input_filter,
    auto_build_tool_service,
    delegate_to_agent,
    execute_function_tool,
    handle_load_skill,
    inject_skill_catalog,
    log_tool_loop_iteration,
    parse_tool_args,
    rebuild_all_tools,
)
from fluxibly.config.loader import load_tool_config_with_path
from fluxibly.llm.registry import create_llm
from fluxibly.monitoring.collector import MonitoringCollector
from fluxibly.monitoring.config import MonitoringConfig
from fluxibly.schema.tools import ToolCall, ToolResult
from fluxibly.tools.executor import ToolService


class SimpleAgent(BaseAgent):
    """The most fundamental Agent implementation.

    Manages:
    - A single LLM
    - A ToolService for tool execution
    - Tools, handoffs (transfer_to_xxx), agents-as-tools (agent_xxx), and skills
    - System prompt and user prompt templates
    """

    def __init__(self, config: AgentConfig) -> None:
        super().__init__(config)

        # Primary LLM
        self.llm = create_llm(config.llm)

        # Tool service (can be set externally or built from config)
        self.tool_service: ToolService | None = None

        # Monitoring collector (loaded from env — never from agent config)
        self._collector: MonitoringCollector | None = None
        _mon = MonitoringConfig.from_env()
        if _mon.enabled:
            self._collector = MonitoringCollector.get_or_create(_mon)

        # Resolution maps (populated by _prepare)
        self.tool_map: dict[str, dict[str, Any]] = {}
        self.handoff_map: dict[str, dict[str, Any]] = {}
        self.agent_tool_map: dict[str, dict[str, Any]] = {}
        self.skill_map: dict[str, SkillMetadata] = {}

        # Flat tool list for LLM (rebuilt by _prepare)
        self.all_tools: list[dict[str, Any]] = []
        # Tool details for monitoring: [{name, description, parameters}, ...]
        self._tool_details: list[dict[str, Any]] = []

        # Tool handler auto-discovery (populated by _prepare, used by _auto_build_tool_service)
        self._tool_yaml_paths: dict[str, Path] = {}
        self._tool_service_built: bool = False

        # Tracks which requirements.txt files have been installed (avoid re-installing)
        self._installed_requirements: set[str] = set()

        # Resolve config-declared references
        self._prepare(
            tools=config.tools,
            handoffs=config.handoffs,
            agents=config.agents,
            skills=config.skills,
        )

    # ══════════════════════════════════════════════════════════════
    # _prepare() — Shared Resolution Method
    # ══════════════════════════════════════════════════════════════

    def _prepare(
        self,
        tools: list[str | dict[str, Any]] | None = None,
        handoffs: list[str | HandoffConfig | dict[str, Any]] | None = None,
        agents: list[str | dict[str, Any]] | None = None,
        skills: list[str | dict[str, Any]] | None = None,
    ) -> None:
        """Resolve all references and rebuild the flat tool list.

        Each entry can be a ``str`` (name/path) or ``dict`` (resource bundle).
        Deduplication by name — entries already in the maps are skipped.

        At the end, ``self.all_tools`` is rebuilt from all maps.
        """
        # Lazy import to avoid circular dependency (registry imports SimpleAgent)
        from fluxibly.agent.registry import create_agent

        # Resolve tools
        for entry in tools or []:
            name, resolved_path, self.sandbox_dir = resolve_entry(
                entry,
                "tool",
                self.sandbox_dir,
            )
            if name in self.tool_map:
                continue
            yaml_path: Path | None = None
            if resolved_path:
                # Load from YAML path (with path tracking for handler discovery)
                tool_dict, yaml_path = load_tool_config_with_path(
                    resolved_path
                )
            elif (
                isinstance(entry, dict)
                and "type" in entry
                and entry.get("type") == "function"
            ):
                # Already a full tool dict (e.g. from runtime context)
                tool_dict = entry
                fn_name = entry.get("function", {}).get("name", "")
                name = fn_name or name
            elif isinstance(entry, dict) and "function" in entry:
                tool_dict = entry
                fn_name = entry.get("function", {}).get("name", "")
                name = fn_name or name
            else:
                # str name/path — resolve via loader
                tool_ref = (
                    entry
                    if isinstance(entry, str)
                    else resolved_path or str(entry)
                )
                tool_dict, yaml_path = load_tool_config_with_path(tool_ref)
                fn_name = tool_dict.get("function", {}).get("name", "")
                name = fn_name or name
            self.tool_map[name] = tool_dict
            if yaml_path is not None:
                self._tool_yaml_paths[name] = yaml_path

        # Resolve handoffs
        for entry in handoffs or []:
            hc: HandoffConfig | None = None
            agent_ref: str | None = None

            if isinstance(entry, HandoffConfig):
                hc = entry
                agent_ref = hc.agent
            elif (
                isinstance(entry, dict)
                and "agent" in entry
                and "type" not in entry
            ):
                # HandoffConfig-style dict from YAML
                hc = HandoffConfig(**entry)
                agent_ref = hc.agent
            elif isinstance(entry, dict) and entry.get("type") in ("agent",):
                # Resource bundle
                name, resolved_path, self.sandbox_dir = resolve_entry(
                    entry,
                    "agent",
                    self.sandbox_dir,
                )
                agent_ref = resolved_path or name
            elif isinstance(entry, str):
                agent_ref = entry
            else:
                agent_ref = str(entry)

            if not agent_ref:
                continue

            # Create the sub-agent
            sub_agent = create_agent(agent_ref)
            sub_agent.session_id = self.session_id
            sub_agent.sandbox_dir = self.sandbox_dir
            sub_agent._sandbox_session = self._sandbox_session

            agent_name = sub_agent.config.name
            tool_name = (
                hc.tool_name if hc else None
            ) or f"transfer_to_{agent_name}"

            if tool_name in self.handoff_map:
                continue

            tool_desc = (
                (hc.tool_description if hc else None)
                or sub_agent.config.description
                or f"Transfer to {agent_name}"
            )
            input_filter = hc.input_filter if hc else None

            # Build transfer_to_xxx tool dict
            tool_dict = {
                "type": "function",
                "function": {
                    "name": tool_name,
                    "description": tool_desc,
                    "parameters": {
                        "type": "object",
                        "properties": {},
                    },
                },
            }

            self.handoff_map[tool_name] = {
                "agent": sub_agent,
                "input_filter": input_filter,
                "tool_dict": tool_dict,
            }

        # Resolve agents-as-tools
        for entry in agents or []:
            # Accept BaseAgent instances directly (e.g. from runtime context)
            if isinstance(entry, BaseAgent):
                sub_agent = entry
            else:
                name, resolved_path, self.sandbox_dir = resolve_entry(
                    entry,
                    "agent",
                    self.sandbox_dir,
                )
                agent_ref = resolved_path or (
                    entry if isinstance(entry, str) else name
                )
                sub_agent = create_agent(agent_ref)

            sub_agent.session_id = self.session_id
            sub_agent.sandbox_dir = self.sandbox_dir
            sub_agent._sandbox_session = self._sandbox_session

            agent_name = sub_agent.config.name
            tool_name = f"agent_{agent_name}"

            if tool_name in self.agent_tool_map:
                continue

            tool_desc = (
                sub_agent.config.description or f"Delegate to {agent_name}"
            )
            tool_dict = sub_agent.as_tool(
                tool_name=tool_name, tool_description=tool_desc
            )

            self.agent_tool_map[tool_name] = {
                "agent": sub_agent,
                "tool_dict": tool_dict,
            }

        # Resolve skills
        for entry in skills or []:
            name, resolved_path, self.sandbox_dir = resolve_entry(
                entry,
                "skill",
                self.sandbox_dir,
            )
            if name in self.skill_map:
                continue
            skill_ref = resolved_path or (
                entry if isinstance(entry, str) else name
            )
            metadata = load_skill_metadata(skill_ref)
            self.skill_map[metadata.name] = metadata

        # Rebuild flat tool list
        self.all_tools, self._tool_names, self._tool_details = (
            rebuild_all_tools(
                self.tool_map,
                self.handoff_map,
                self.agent_tool_map,
                self.skill_map,
            )
        )

    # ══════════════════════════════════════════════════════════════
    # forward()
    # ══════════════════════════════════════════════════════════════

    async def forward(
        self,
        messages: list[dict[str, Any]],
        context: dict[str, Any] | None = None,
    ) -> AgentResponse:
        """Execute simple agent logic.

        Handles handoffs (transfer_to_xxx), agent delegation (agent_xxx),
        skill loading (load_skill), and regular tool calls.

        Args:
            messages: Full conversation history from the caller.
            context: Optional runtime context dict. Supported keys:
                - prompt_params: dict — template params
                - tools: list — additional tools (str or dict)
                - agents: list — additional agents (str or dict)
                - handoffs: list — additional handoffs (str, HandoffConfig, or dict)
                - skills: list — additional skills (str or dict)
                - _monitoring: dict — (internal) propagated by parent agents

        Returns:
            AgentResponse with content, metadata, and optional handoff_to.
        """
        context = context or {}

        # Lazy auto-build ToolService from handler .py files (first forward() only)
        if not self._tool_service_built:
            if self.tool_service is None:
                ts, session = await auto_build_tool_service(
                    self._tool_yaml_paths,
                    self.tool_map,
                    self.config,
                    self._sandbox_session,
                    self.sandbox_dir,
                    self.logger,
                )
                if ts is not None:
                    self.tool_service = ts
                if session is not None:
                    self._sandbox_session = session
            self._tool_service_built = True

        # Set up monitoring if enabled
        mon: AgentMonitoringSetup | None = None
        if self._collector:
            mon = await setup_agent_monitoring(
                self._collector,
                messages,
                context,
                self.__class__.__name__,
                self.config.name,
            )

        try:
            # Step 1: Resolve context
            system_params, user_params, runtime_tools, runtime_agents = (
                resolve_context(context, logger=self.logger)
            )
            runtime_handoffs = context.get("handoffs", [])
            runtime_skills = context.get("skills", [])

            # Step 2: _prepare with context-injected resources
            if (
                runtime_tools
                or runtime_agents
                or runtime_handoffs
                or runtime_skills
            ):
                self._prepare(
                    tools=runtime_tools,
                    handoffs=runtime_handoffs,
                    agents=runtime_agents,
                    skills=runtime_skills,
                )

            # Step 2.1: Register direct function tools from context
            direct_fns = context.pop("direct_functions", [])
            if direct_fns:
                if self.tool_service is None:
                    self.tool_service = ToolService()
                    self._tool_service_built = True
                for entry in direct_fns:
                    schema = entry["schema"]
                    handler = entry["handler"]
                    fn_name = schema["function"]["name"]
                    # Always overwrite handler — closures carry request-scoped data
                    self.tool_service.function_registry[fn_name] = handler
                    if fn_name not in [
                        t.get("function", {}).get("name")
                        for t in self.tool_service.tools
                    ]:
                        self.tool_service.tools.append(schema)
                    self.tool_map[fn_name] = schema
                self.all_tools, self._tool_names, self._tool_details = (
                    rebuild_all_tools(
                        self.tool_map,
                        self.handoff_map,
                        self.agent_tool_map,
                        self.skill_map,
                    )
                )

            # Step 2.5: Stage skill directories into sandbox
            if self.skill_map and self.sandbox_dir:
                for skill in self.skill_map.values():
                    if skill.sandbox_path is None:
                        skill.sandbox_path = stage_skill(
                            skill.path,
                            self.sandbox_dir,
                        )

            # Step 2.6: Auto-install requirements.txt from skills and tools
            if self.sandbox_dir and self._sandbox_session is not None:
                skill_dirs = [
                    s.sandbox_path or s.path for s in self.skill_map.values()
                ]
                tool_yamls = list(self._tool_yaml_paths.values())
                req_paths = collect_requirements(skill_dirs, tool_yamls)
                for req_path in req_paths:
                    if req_path not in self._installed_requirements:
                        await self._sandbox_session.install_requirements(
                            req_path
                        )
                        self._installed_requirements.add(req_path)

            # Step 3: Prepare prompts (+ skill catalog injection)
            working_messages = prepare_prompts(
                messages,
                self.config,
                system_params,
                user_params,
                logger=self.logger,
            )

            # Inject skill catalog into system prompt if skills are configured
            if self.skill_map:
                catalog = build_skill_catalog(self.skill_map)
                inject_skill_catalog(working_messages, catalog)

            # Inject data file info into system prompt if data was staged
            staged_data = context.get("_staged_data")
            if staged_data:
                working_messages = inject_data_prompt(
                    working_messages, staged_data
                )

            # Inject sandbox file tree so the LLM knows what scripts/data are available
            if self.sandbox_dir:
                working_messages = inject_sandbox_files_prompt(
                    working_messages,
                    self.sandbox_dir,
                )

            # Step 4: LLM call + tool-call loop with handoff/skill awareness
            llm_response, iteration, llm_call_count, tool_interactions = (
                await self._agent_forward_loop(
                    working_messages=working_messages,
                    context=context,
                    mon=mon,
                )
            )

            # Step 5: Finalize
            return await finalize_response(
                llm_response=llm_response,
                config=self.config,
                context=context,
                messages=messages,
                iteration=iteration,
                llm_call_count=llm_call_count,
                mon=mon,
                logger=self.logger,
                tool_names=self._tool_details or None,
                tool_interactions=tool_interactions or None,
            )

        except Exception as e:
            if mon:
                await fail_agent_monitoring(mon, e)
            raise

    # ══════════════════════════════════════════════════════════════
    # _agent_forward_loop — LLM call + tool-call loop
    # ══════════════════════════════════════════════════════════════

    async def _agent_forward_loop(
        self,
        working_messages: list[dict[str, Any]],
        context: dict[str, Any],
        mon: AgentMonitoringSetup | None,
    ) -> tuple[Any, int, int, list[dict[str, Any]]]:
        """Execute the initial LLM call and the agentic tool-call loop.

        Handles:
        - ``transfer_to_xxx`` → return immediately with handoff_to
        - ``agent_xxx`` → delegate to sub-agent, return result as tool output
        - ``load_skill`` → load skill body, return as tool output
        - Regular tools → execute via ToolService

        Returns:
            Tuple of (final_llm_response, iteration_count, llm_call_count,
            tool_interactions).
        """
        config = self.config
        max_tool_iterations = config.max_tool_iterations
        llm_config = config.llm

        # Initial LLM call
        self.llm.set_messages(working_messages)
        self.llm.set_tools(self.all_tools)
        self.logger.log_input(
            {
                "message_count": len(working_messages),
                "tool_count": len(self.all_tools),
            }
        )
        llm_response = await llm_forward(
            self.llm,
            llm_config,
            mon,
            working_messages,
            tool_names=self._tool_details or None,
        )

        iteration = 0
        llm_call_count = 1

        # Track which messages have skill bodies to evict
        skill_body_indices: list[tuple[int, str]] = (
            []
        )  # (message_index, skill_name)

        # Accumulate all tool interactions for metadata
        all_tool_interactions: list[dict[str, Any]] = []

        while (
            llm_response.output.tool_calls and iteration < max_tool_iterations
        ):
            iteration += 1
            log_tool_loop_iteration(
                self.logger,
                iteration,
                max_tool_iterations,
                llm_response.output.tool_calls,
            )

            # Add assistant message with tool calls
            working_messages.append(
                {
                    "role": "assistant",
                    "content": llm_response.output.output_text,
                    "tool_calls": [
                        tc.model_dump()
                        for tc in llm_response.output.tool_calls
                    ],
                }
            )

            # Process tool calls — check for handoffs, skills, agents, regular tools
            results: list[ToolResult] = []
            handoff_response: AgentResponse | None = None

            for seq, tool_call in enumerate(llm_response.output.tool_calls):
                tool_name = tool_call.name
                tc_start = time.monotonic()

                # --- Handoff: transfer_to_xxx ---
                if tool_name in self.handoff_map:
                    handoff_entry = self.handoff_map[tool_name]
                    target_agent = handoff_entry["agent"]
                    input_filter_name = handoff_entry["input_filter"]

                    # Propagate latest sandbox state
                    target_agent.sandbox_dir = self.sandbox_dir
                    target_agent._sandbox_session = self._sandbox_session

                    # Apply input filter to messages for the receiving agent
                    handoff_messages = apply_input_filter(
                        working_messages,
                        input_filter_name,
                    )

                    self.logger.log_event(
                        "handoff",
                        {
                            "from": self.config.name,
                            "to": target_agent.config.name,
                            "tool": tool_name,
                            "input_filter": input_filter_name,
                        },
                    )

                    # Return with handoff_to set — Runner will swap agents
                    handoff_response = AgentResponse(
                        content=llm_response,
                        metadata={
                            "handoff_to_agent": target_agent.config.name,
                            "handoff_messages": handoff_messages,
                        },
                        handoff_to=target_agent,
                    )
                    break

                # --- Skill: load_skill ---
                elif tool_name == "load_skill":
                    result = handle_load_skill(tool_call, self.skill_map)
                    results.append(result)

                    skill_name_arg = parse_tool_args(tool_call.arguments).get(
                        "name", ""
                    )
                    self.logger.log_event(
                        "skill.load",
                        {
                            "skill": skill_name_arg,
                            "is_error": result.is_error,
                        },
                    )

                    # Record monitoring for skill load
                    if mon:
                        duration_ms = int((time.monotonic() - tc_start) * 1000)
                        await mon.collector.record_tool_call(
                            span_ctx=mon.span_ctx,
                            tool_call=tool_call,
                            result=result,
                            duration_ms=duration_ms,
                            iteration=iteration,
                            sequence_in_batch=seq,
                        )

                    # Track for eviction after next LLM call
                    msg_idx = len(working_messages)  # Will be added next
                    skill_body_indices.append(
                        (msg_idx + len(results) - 1, skill_name_arg)
                    )

                # --- Agent-as-tool: agent_xxx ---
                elif tool_name in self.agent_tool_map:
                    agent_entry = self.agent_tool_map[tool_name]
                    sub_agent = agent_entry["agent"]

                    # Propagate latest sandbox state (may have been
                    # created lazily after _prepare())
                    sub_agent.sandbox_dir = self.sandbox_dir
                    sub_agent._sandbox_session = self._sandbox_session

                    self.logger.log_event(
                        "agent.delegate",
                        {
                            "agent": sub_agent.config.name,
                            "tool": tool_name,
                        },
                    )

                    result = await delegate_to_agent(
                        sub_agent,
                        tool_call,
                        context,
                        self.logger,
                    )
                    results.append(result)

                    # Record monitoring for agent delegation
                    if mon:
                        duration_ms = int((time.monotonic() - tc_start) * 1000)
                        await mon.collector.record_tool_call(
                            span_ctx=mon.span_ctx,
                            tool_call=tool_call,
                            result=result,
                            duration_ms=duration_ms,
                            iteration=iteration,
                            sequence_in_batch=seq,
                        )

                # --- Regular tool ---
                else:

                    async def _exec(tc: ToolCall) -> ToolResult:
                        return await execute_function_tool(
                            tc, self.tool_service
                        )

                    if mon:
                        result = await monitored_tool_call(
                            mon.collector,
                            tool_call,
                            _exec,
                            mon.span_ctx,
                            iteration,
                            seq,
                        )
                    else:
                        result = await _exec(tool_call)
                    results.append(result)

            # If handoff occurred, return immediately
            if handoff_response is not None:
                return (
                    llm_response,
                    iteration,
                    llm_call_count,
                    all_tool_interactions,
                )

            # Evict previous skill bodies (from earlier iterations)
            # TODO: temporarily disabled for debugging — skill bodies kept in history
            # evict_skill_bodies(working_messages, skill_body_indices)
            # skill_body_indices.clear()

            # Add tool results to messages
            for tool_call, result in zip(
                llm_response.output.tool_calls, results
            ):
                working_messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": str(result.content),
                    }
                )

                # Accumulate tool interaction record
                tc_args = tool_call.arguments
                if isinstance(tc_args, str):
                    tc_args = parse_tool_args(tc_args)
                all_tool_interactions.append(
                    {
                        "tool_name": tool_call.name,
                        "tool_call_id": tool_call.id,
                        "arguments": tc_args,
                        "result": result.content,
                        "is_error": result.is_error,
                        "iteration": iteration,
                    }
                )

                # Track skill body messages for next eviction
                if tool_call.name == "load_skill" and not result.is_error:
                    skill_name_arg = parse_tool_args(tool_call.arguments).get(
                        "name", ""
                    )
                    skill_body_indices.append(
                        (len(working_messages) - 1, skill_name_arg)
                    )

            # Re-call LLM with tool results
            self.llm.set_messages(working_messages)
            llm_response = await llm_forward(
                self.llm,
                llm_config,
                mon,
                working_messages,
                tool_names=self._tool_details or None,
            )
            llm_call_count += 1

        # Final eviction of any remaining skill bodies
        # TODO: temporarily disabled for debugging — skill bodies kept in history
        # evict_skill_bodies(working_messages, skill_body_indices)

        return llm_response, iteration, llm_call_count, all_tool_interactions
