"""Monitoring utilities for agent forward() calls.

Provides reusable helpers for setting up, finalizing, and recording
monitoring traces and spans around agent execution. Any agent
implementation can use these instead of rolling its own lifecycle
management.
"""

from __future__ import annotations

import asyncio
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from fluxibly.monitoring.collector import MonitoringCollector
from fluxibly.monitoring.context import SpanContext, TraceContext
from fluxibly.monitoring.utils import extract_user_input
from fluxibly.schema.tools import ToolCall, ToolResult

if TYPE_CHECKING:
    from fluxibly.agent.base import AgentConfig, AgentResponse
    from fluxibly.llm.base import LLMConfig
    from fluxibly.schema.response import LLMResponse


@dataclass
class AgentMonitoringSetup:
    """Monitoring context created by setup_agent_monitoring().

    Bundles the collector, trace/span contexts, and depth so that
    a single ``mon`` argument carries all monitoring state.
    """

    collector: MonitoringCollector
    trace_ctx: TraceContext
    span_ctx: SpanContext
    is_root: bool
    depth: int


async def setup_agent_monitoring(
    collector: MonitoringCollector,
    messages: list[dict[str, Any]],
    context: dict[str, Any],
    agent_class_name: str,
    agent_name: str,
) -> AgentMonitoringSetup:
    """Set up monitoring trace and agent span for a forward() call.

    Creates a trace if this is the top-level (root) call, creates an agent
    span, and injects monitoring context into ``context["_monitoring"]`` for
    child components.
    """
    monitoring_ctx = context.get("_monitoring", {})
    trace_ctx = monitoring_ctx.get("trace_ctx")
    parent_span_id = monitoring_ctx.get("parent_span_id")
    depth = monitoring_ctx.get("depth", 0)

    is_root = trace_ctx is None
    if is_root:
        trace_ctx = await collector.start_trace(
            entry_input=extract_user_input(messages),
            service_name=collector.config.service_name,
            environment=collector.config.environment,
            tags=collector.config.tags,
        )

    span_ctx = await collector.start_span(
        trace_ctx=trace_ctx,
        span_type="agent",
        component_class=agent_class_name,
        component_name=agent_name,
        parent_span_id=parent_span_id,
        depth=depth,
    )

    child_depth = depth + 1
    context["_monitoring"] = {
        "trace_ctx": trace_ctx,
        "parent_span_id": span_ctx.span_id,
        "depth": child_depth,
    }

    return AgentMonitoringSetup(
        collector=collector,
        trace_ctx=trace_ctx,
        span_ctx=span_ctx,
        is_root=is_root,
        depth=child_depth,
    )


async def finalize_agent_monitoring(
    setup: AgentMonitoringSetup,
    result: AgentResponse,
    config: AgentConfig,
    context: dict[str, Any],
    messages: list[dict[str, Any]],
    iteration_count: int = 0,
    llm_call_count: int = 0,
    tool_names: list[dict[str, Any]] | None = None,
    tool_interactions: list[dict[str, Any]] | None = None,
) -> None:
    """Record agent details and close monitoring spans/trace on success."""
    await setup.collector.record_agent_details(
        span_ctx=setup.span_ctx,
        agent_response=result,
        config=config,
        context=context,
        input_messages=messages,
        iteration_count=iteration_count,
        total_llm_calls=llm_call_count,
        tool_names=tool_names,
        tool_interactions=tool_interactions,
    )
    await setup.collector.end_span(setup.span_ctx, status="completed")

    if setup.is_root:
        await setup.collector.end_trace(
            setup.trace_ctx,
            final_output=result.content.output.output_text,
            status="completed",
        )


async def fail_agent_monitoring(
    setup: AgentMonitoringSetup,
    error: Exception,
) -> None:
    """Close monitoring spans/trace on failure."""
    await setup.collector.end_span(
        setup.span_ctx, status="failed", error=str(error)
    )
    if setup.is_root:
        await setup.collector.end_trace(
            setup.trace_ctx,
            final_output="",
            status="failed",
            error_message=str(error),
        )


async def monitored_llm_forward(
    collector: MonitoringCollector,
    llm: Any,
    llm_config: LLMConfig,
    trace_ctx: TraceContext,
    parent_span_ctx: SpanContext,
    working_messages: list[dict[str, Any]],
    depth: int,
    tool_names: list[dict[str, Any]] | None = None,
) -> LLMResponse:
    """Execute LLM.forward() wrapped in a monitoring span."""
    llm_span = await collector.start_span(
        trace_ctx=trace_ctx,
        span_type="llm",
        component_class=llm.__class__.__name__,
        component_name=llm_config.model,
        parent_span_id=parent_span_ctx.span_id,
        model=llm_config.model,
        provider=getattr(llm, "_provider", None),
        depth=depth,
    )

    try:
        llm_response = llm.forward()

        await collector.record_llm_details(
            span_ctx=llm_span,
            llm_response=llm_response,
            config=llm_config,
            input_messages=working_messages,
            tool_names=tool_names,
        )
        await collector.end_span(llm_span, status="completed")

        return llm_response

    except Exception as e:
        await collector.end_span(llm_span, status="failed", error=str(e))
        raise


# ══════════════════════════════════════════════════════════════
# Tool Call Monitoring
# ══════════════════════════════════════════════════════════════


async def monitored_tool_call(
    collector: MonitoringCollector,
    tool_call: ToolCall,
    execute_fn: Callable[[ToolCall], Awaitable[ToolResult]],
    parent_span_ctx: SpanContext,
    iteration: int,
    sequence_in_batch: int,
) -> ToolResult:
    """Execute a single tool call with timing and monitoring recording."""
    start_time = time.monotonic()
    result = await execute_fn(tool_call)
    duration_ms = int((time.monotonic() - start_time) * 1000)

    await collector.record_tool_call(
        span_ctx=parent_span_ctx,
        tool_call=tool_call,
        result=result,
        duration_ms=duration_ms,
        iteration=iteration,
        sequence_in_batch=sequence_in_batch,
    )

    return result


async def monitored_tool_calls_batch(
    tool_calls: list[ToolCall],
    execute_fn: Callable[[ToolCall], Awaitable[ToolResult]],
    collector: MonitoringCollector,
    parent_span_ctx: SpanContext,
    iteration: int,
    max_workers: int = 1,
) -> list[ToolResult]:
    """Execute tool calls with monitoring and concurrency control."""
    if max_workers <= 1 or len(tool_calls) == 1:
        results: list[ToolResult] = []
        for seq, tc in enumerate(tool_calls):
            results.append(
                await monitored_tool_call(
                    collector,
                    tc,
                    execute_fn,
                    parent_span_ctx,
                    iteration,
                    seq,
                )
            )
        return results

    semaphore = asyncio.Semaphore(max_workers)

    async def _run(tc: ToolCall, seq: int) -> ToolResult:
        async with semaphore:
            return await monitored_tool_call(
                collector,
                tc,
                execute_fn,
                parent_span_ctx,
                iteration,
                seq,
            )

    raw_results: list[ToolResult | BaseException] = await asyncio.gather(
        *[_run(tc, i) for i, tc in enumerate(tool_calls)],
        return_exceptions=True,
    )

    # Convert exceptions to error ToolResults
    final: list[ToolResult] = []
    for i, r in enumerate(raw_results):
        if isinstance(r, BaseException):
            final.append(
                ToolResult(
                    tool_call_id=tool_calls[i].id,
                    content=f"Tool execution failed: {r}",
                    is_error=True,
                )
            )
        else:
            final.append(r)
    return final


# ══════════════════════════════════════════════════════════════
# Unified (optional monitoring) helper
# ══════════════════════════════════════════════════════════════


async def llm_forward(
    llm: Any,
    llm_config: LLMConfig,
    mon: AgentMonitoringSetup | None,
    working_messages: list[dict[str, Any]],
    tool_names: list[dict[str, Any]] | None = None,
) -> LLMResponse:
    """Execute LLM.forward(), optionally wrapped in monitoring."""
    if mon:
        return await monitored_llm_forward(
            mon.collector,
            llm,
            llm_config,
            mon.trace_ctx,
            mon.span_ctx,
            working_messages,
            mon.depth,
            tool_names=tool_names,
        )
    return llm.forward()
