"""Runner — top-level execution loop for Fluxibly agents.

Manages:
- Agent resolution (str/AgentConfig/BaseAgent → ready-to-run agent)
- Handoff chains (detects handoff_to, swaps active agent, repeats)
- Session sandbox lifecycle (create, share across agents, cleanup)
"""

from __future__ import annotations

import asyncio
import os
import shutil
import uuid
from pathlib import Path
from typing import Any

from pydantic import BaseModel, Field

from fluxibly.agent.base import AgentConfig, AgentResponse, BaseAgent
from fluxibly.agent.utils.data import stage_data_files
from fluxibly.config.env import load_env
from fluxibly.logging import Logger

logger = Logger(component="runner")


class RunnerResult(BaseModel):
    """Result from Runner.run() — wraps AgentResponse with session info."""

    response: AgentResponse = Field(
        ..., description="The final agent response"
    )
    session_id: str = Field(..., description="The session ID used/generated")


class Runner:
    """Top-level execution loop for agents.

    Handles agent resolution, handoff chains, and session sandbox lifecycle.
    """

    @staticmethod
    async def run(
        agent: str | AgentConfig | BaseAgent,
        messages: list[dict[str, Any]],
        context: dict[str, Any] | None = None,
        max_turns: int = 20,
        session_id: str | None = None,
        sandbox_root: str | None = None,
        sandbox: bool = True,
    ) -> RunnerResult:
        """Run the agent loop until completion.

        Agent resolution:
            1. If ``str`` → ``create_agent(path_or_name)``
            2. If ``AgentConfig`` → ``create_agent(config)``
            3. If ``BaseAgent`` → use directly

        Session lifecycle:
            4. Resolve sandbox_root: arg → env ``FLUXIBLY_SANDBOX_DIR`` → ``/tmp/fluxibly/``
            5. Create or reuse session directory
            6. Set ``agent.session_id`` + ``agent.sandbox_dir``

        Execution loop:
            7. Call ``current_agent.forward()``
            8. If ``handoff_to`` is set → propagate session, swap agent, repeat
            9. No handoff → return ``RunnerResult``

        Cleanup:
            10. Delete session directory in ``finally`` block

        Args:
            agent: Config path/name, AgentConfig, or BaseAgent instance.
            messages: Conversation history.
            context: Optional runtime context dict.
            max_turns: Maximum handoff chain length (safety limit).
            session_id: Optional session ID to reuse.
            sandbox_root: Optional override for sandbox root directory.

        Returns:
            RunnerResult with final AgentResponse and session_id.
        """
        from fluxibly.agent.registry import create_agent

        # Step 0: Ensure .env is loaded (for monitoring config, API keys, etc.)
        load_env()

        # Step 1: Resolve agent
        if isinstance(agent, str):
            current_agent = create_agent(agent)
        elif isinstance(agent, AgentConfig):
            current_agent = create_agent(agent)
        elif isinstance(agent, BaseAgent):
            current_agent = agent
        else:
            raise TypeError(
                f"Runner.run() expects str, AgentConfig, or BaseAgent, "
                f"got {type(agent)}"
            )

        # Step 2: Resolve session
        resolved_session_id = session_id or str(uuid.uuid4())
        sandbox_dir: str | None = None

        if sandbox:
            resolved_sandbox_root = (
                sandbox_root
                or os.environ.get("FLUXIBLY_SANDBOX_DIR")
                or "/tmp/fluxibly/"
            )
            sandbox_dir = str(
                Path(resolved_sandbox_root) / resolved_session_id
            )
            Path(sandbox_dir).mkdir(parents=True, exist_ok=True)

        # Step 3: Set session on agent
        current_agent.session_id = resolved_session_id
        current_agent.sandbox_dir = sandbox_dir

        if context is None:
            context = {}

        # Stage data files if provided in context
        if sandbox_dir:
            data_items = context.get("data")
            if data_items:
                staged = stage_data_files(data_items, sandbox_dir)
                context["_staged_data"] = staged

        logger.info(
            "Runner started: agent={agent}, session={session}",
            agent=current_agent.config.name,
            session=resolved_session_id,
        )

        try:
            # Step 4: Execution loop
            turn = 0
            while turn < max_turns:
                turn += 1
                logger.info(
                    "Turn {turn}/{max}: agent={agent}",
                    turn=turn,
                    max=max_turns,
                    agent=current_agent.config.name,
                )

                response = await current_agent.forward(messages, context)

                # Check for handoff
                if response.handoff_to is not None:
                    next_agent = response.handoff_to
                    logger.info(
                        "Handoff: {from_agent} → {to_agent}",
                        from_agent=current_agent.config.name,
                        to_agent=next_agent.config.name,
                    )

                    # Propagate session to receiving agent
                    next_agent.session_id = resolved_session_id
                    next_agent.sandbox_dir = sandbox_dir
                    next_agent._sandbox_session = getattr(
                        current_agent, "_sandbox_session", None
                    )

                    current_agent = next_agent
                    continue

                # No handoff — we're done
                logger.info(
                    "Runner completed: agent={agent}, turns={turns}",
                    agent=current_agent.config.name,
                    turns=turn,
                )
                return RunnerResult(
                    response=response,
                    session_id=resolved_session_id,
                )

            # Exceeded max_turns
            logger.warning(
                "Runner hit max_turns ({max}), returning last response",
                max=max_turns,
            )
            return RunnerResult(
                response=response,  # type: ignore[possibly-undefined]
                session_id=resolved_session_id,
            )

        finally:
            # Step 5: Flush monitoring so all traces are committed before exit
            await Runner._flush_monitoring(current_agent)

            # Step 6: Cleanup sandbox
            if sandbox_dir:
                sandbox_path = Path(sandbox_dir)
                if sandbox_path.exists():
                    try:
                        shutil.rmtree(sandbox_path)
                        logger.debug(
                            "Cleaned up sandbox: {path}", path=sandbox_dir
                        )
                    except OSError as e:
                        logger.warning(
                            "Failed to clean up sandbox {path}: {err}",
                            path=sandbox_dir,
                            err=e,
                        )

    @staticmethod
    async def _flush_monitoring(agent: BaseAgent) -> None:
        """Flush the monitoring writer so all records are committed.

        Called in Runner.run()'s finally block to ensure traces are not
        left IN_PROGRESS when the process exits.
        """
        collector = getattr(agent, "_collector", None)
        if collector is not None and collector._initialized:
            try:
                await collector._writer.flush()
            except Exception as e:
                logger.warning(
                    "Failed to flush monitoring: {err}",
                    err=e,
                )

    @staticmethod
    def run_sync(
        agent: str | AgentConfig | BaseAgent,
        messages: list[dict[str, Any]],
        context: dict[str, Any] | None = None,
        max_turns: int = 20,
        session_id: str | None = None,
        sandbox_root: str | None = None,
        sandbox: bool = True,
    ) -> RunnerResult:
        """Synchronous wrapper for Runner.run()."""
        return asyncio.run(
            Runner.run(
                agent=agent,
                messages=messages,
                context=context,
                max_turns=max_turns,
                session_id=session_id,
                sandbox_root=sandbox_root,
                sandbox=sandbox,
            )
        )
