"""Agent runner — executes AI agents as ARQ jobs with safety checks.

Safety features:
- Max executions per agent per hour (rate limiting)
- Max duration per execution (asyncio timeout)
- Auto-stop on infinite loop (same tool called 5x consecutively)
- Budget limit per agent (track cumulative cost_usd, stop if over budget)
"""

from __future__ import annotations

import asyncio
import logging
from datetime import UTC, datetime
from typing import Any

from sqlalchemy import func, select

from app.core.db import get_session_factory

logger = logging.getLogger(__name__)

# Track consecutive tool calls per agent run to detect infinite loops
_consecutive_tool_calls: dict[str, dict[str, int]] = {}  # run_id -> {tool_name: count}


async def run_agent(
    ctx: dict[str, Any],
    agent_id: str,
    trigger_type: str = "manual",
    trigger_data: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """ARQ job function. Loads AgentDefinition from DB, checks rate limits,
    gathers context, calls LLM via LiteLLM, executes tool calls via ToolRegistry,
    saves result to AgentRun. Returns result dict.

    Safety checks:
    1. Rate limit: max_executions_per_hour
    2. Max duration: max_duration_seconds (asyncio.timeout)
    3. Infinite loop: same tool 5x consecutively
    4. Budget limit: cumulative cost_usd
    """
    from app.plugins.builtins.automation.models import AgentDefinition, AgentRun

    factory = get_session_factory()

    # Load agent definition
    async with factory() as db:
        result = await db.execute(
            select(AgentDefinition).where(AgentDefinition.id == agent_id)
        )
        agent = result.scalar_one_or_none()

    if agent is None:
        logger.error("AgentDefinition %s not found", agent_id)
        return {"error": f"AgentDefinition {agent_id} not found", "status": "failed"}

    if not agent.is_active:
        logger.warning("AgentDefinition %s is inactive", agent_id)
        return {"error": "Agent is inactive", "status": "skipped"}

    # ── Safety Check 1: Rate Limit ──
    if agent.max_executions_per_hour:
        async with factory() as db:
            one_hour_ago = datetime.now(UTC)
            count_result = await db.execute(
                select(func.count())
                .select_from(AgentRun)
                .where(
                    AgentRun.agent_id == agent.id,
                    AgentRun.created_at >= one_hour_ago,
                )
            )
            recent_runs = count_result.scalar() or 0
            if recent_runs >= agent.max_executions_per_hour:
                logger.warning(
                    "Rate limit hit for agent %s: %d runs in last hour (max %d)",
                    agent.id, recent_runs, agent.max_executions_per_hour,
                )
                return {"error": "Rate limit exceeded", "status": "rate_limited"}

    # ── Safety Check 2: Budget Limit ──
    if agent.budget_limit_usd > 0:
        async with factory() as db:
            cost_result = await db.execute(
                select(func.coalesce(func.sum(AgentRun.cost_usd), 0.0))
                .where(AgentRun.agent_id == agent.id)
            )
            total_cost = float(cost_result.scalar() or 0.0)
            if total_cost >= agent.budget_limit_usd:
                logger.warning(
                    "Budget limit hit for agent %s: $%.4f total cost (limit $%.2f)",
                    agent.id, total_cost, agent.budget_limit_usd,
                )
                return {"error": "Budget limit exceeded", "status": "budget_exceeded"}

    # Gather context
    context_data: dict[str, Any] = {}
    if trigger_type == "proactive" or agent.mode == "proactive":
        # Collect context data (recent contacts, mails, events)
        try:
            from app.services.contact_service import list_contacts
            async with factory() as db:
                contacts = await list_contacts(db, agent.tenant_id, page=1, page_size=10)
                context_data["recent_contacts"] = contacts.get("items", [])
        except Exception:
            logger.warning("Failed to collect contacts for proactive context")

        try:
            from app.core.cache import get_cached_mail_summary
            context_data["recent_mails"] = get_cached_mail_summary(agent.tenant_id) or []
        except Exception:
            logger.warning("Failed to collect mails for proactive context")

        try:
            from app.services.workflow_service import list_instances
            async with factory() as db:
                events = await list_instances(db, agent.tenant_id, page=1, page_size=10)
                context_data["recent_events"] = events.get("items", [])
        except Exception:
            logger.warning("Failed to collect events for proactive context")
    else:
        # Reactive mode: use trigger_data as context
        context_data = trigger_data or {}

    # Call LLM via LiteLLM
    result_data: dict[str, Any] = {
        "agent_id": str(agent.id),
        "agent_name": agent.name,
        "trigger_type": trigger_type,
        "status": "running",
        "llm_response": None,
        "tool_calls": [],
        "cost_usd": 0.0,
        "error": None,
    }

    # ── Safety Check 3: Max Duration ──
    max_duration = agent.max_duration_seconds or 300

    try:
        async def _run_llm() -> None:
            """Inner coroutine for LLM call with tool execution."""
            from app.ai.llm_client import LLMClient

            llm = LLMClient(
                model=agent.model or None,
                api_key=agent.api_key or None,
                api_base=agent.api_base or None,
                provider=agent.provider or None,
            )

            # Build system prompt from agent configuration
            system_prompt = agent.system_prompt or "You are a helpful AI assistant."
            user_prompt = f"Context: {context_data}"

            # Use litellm directly for more control
            import litellm

            litellm_model = agent.model or "gpt-4o"
            if agent.provider and agent.provider != "openai":
                litellm_model = f"{agent.provider}/{litellm_model}"

            kwargs: dict[str, Any] = {
                "model": litellm_model,
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                "temperature": 0.3,
                "max_tokens": agent.max_tokens or 1000,
            }

            if agent.api_key:
                kwargs["api_key"] = agent.api_key
            if agent.api_base:
                kwargs["api_base"] = agent.api_base

            response = await litellm.acompletion(**kwargs)
            content = response.choices[0].message.content

            # Track cost
            if hasattr(response, "usage") and response.usage:
                usage = response.usage
                # Estimate cost (simplified)
                input_cost = (usage.prompt_tokens or 0) * 0.00001 / 1000
                output_cost = (usage.completion_tokens or 0) * 0.00003 / 1000
                result_data["cost_usd"] = round(input_cost + output_cost, 6)

            result_data["llm_response"] = content

            # Execute tool calls if LLM returned function calls
            if hasattr(response.choices[0].message, "tool_calls") and response.choices[0].message.tool_calls:
                from app.plugins.builtins.ai_assistant.tool_registry import get_tool_registry

                registry = get_tool_registry()
                tool_call_count: dict[str, int] = {}
                for tc in response.choices[0].message.tool_calls:
                    tool_name = tc.function.name
                    # ── Safety Check 4: Infinite Loop Detection ──
                    tool_call_count[tool_name] = tool_call_count.get(tool_name, 0) + 1
                    if tool_call_count[tool_name] >= 5:
                        logger.warning(
                            "Infinite loop detected for agent %s: tool '%s' called %d times consecutively",
                            agent.id, tool_name, tool_call_count[tool_name],
                        )
                        result_data["error"] = f"Infinite loop detected: tool '{tool_name}' called 5+ times consecutively"
                        result_data["status"] = "failed"
                        return

                    tool = registry.get(tool_name)
                    if tool:
                        try:
                            import json
                            args = json.loads(tc.function.arguments)
                            tool_result = await tool.handler(
                                arguments=args,
                                context={"tenant_id": str(agent.tenant_id)},
                            )
                            result_data["tool_calls"].append({
                                "tool": tool_name,
                                "arguments": args,
                                "result": tool_result,
                            })
                        except Exception as e:
                            result_data["tool_calls"].append({
                                "tool": tool_name,
                                "error": str(e),
                            })

                result_data["status"] = "completed"

        # Run with timeout
        try:
            await asyncio.wait_for(_run_llm(), timeout=max_duration)
        except asyncio.TimeoutError:
            logger.warning(
                "Agent %s execution timed out after %d seconds",
                agent.id, max_duration,
            )
            result_data["status"] = "timed_out"
            result_data["error"] = f"Execution timed out after {max_duration} seconds"

    except Exception as e:
        logger.exception("Agent run failed for %s", agent.id)
        result_data["status"] = "failed"
        result_data["error"] = str(e)

    # Save result to AgentRun
    try:
        async with factory() as db:
            run = AgentRun(
                tenant_id=agent.tenant_id,
                agent_id=agent.id,
                trigger_type=trigger_type,
                status=result_data["status"],
                input_data=context_data,
                output_data=result_data.get("llm_response"),
                tool_calls=result_data.get("tool_calls", []),
                cost_usd=result_data.get("cost_usd", 0.0),
                error_message=result_data.get("error"),
                duration_seconds=None,
            )
            db.add(run)
            await db.flush()
            result_data["run_id"] = str(run.id)
    except Exception as e:
        logger.exception("Failed to save AgentRun for %s", agent.id)
        result_data["save_error"] = str(e)

    return result_data