"""MonitoringCollector — core instrumentation for Fluxibly.

Manages trace/span lifecycle and records Agent, LLM, and Tool telemetry.
Uses a singleton pattern keyed by DSN to share DB resources.
"""

from __future__ import annotations

from datetime import datetime, timezone
from decimal import Decimal
from typing import TYPE_CHECKING, Any
from uuid import uuid4

import loguru

from fluxibly.monitoring.config import MonitoringConfig
from fluxibly.monitoring.context import SpanContext, TraceContext
from fluxibly.monitoring.db.models import (
    AgentCallDetailRecord,
    LLMCallDetailRecord,
    SpanRecord,
    ToolCallDetailRecord,
    TraceRecord,
)
from fluxibly.monitoring.utils import (
    sanitize_dict,
    truncate_text,
)
from fluxibly.monitoring.writer import AsyncDBWriter

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

logger = loguru.logger.bind(component="monitoring")

# Singleton cache keyed by DSN
_collectors: dict[str, MonitoringCollector] = {}


class MonitoringCollector:
    """Core instrumentation that wraps Agent/LLM operations.

    Not called directly by users — integrated automatically when
    MonitoringConfig is set on AgentConfig or LLMConfig.
    """

    def __init__(self, config: MonitoringConfig) -> None:
        self.config = config
        self._repo: Any = None  # MonitoringRepository, set on init
        self._writer = AsyncDBWriter(
            batch_size=config.writer_batch_size,
            flush_interval_ms=config.writer_flush_interval_ms,
        )
        self._initialized = False
        self._seq_counters: dict[str, int] = (
            {}
        )  # parent_span_id -> next sequence number

    @classmethod
    def get_or_create(cls, config: MonitoringConfig) -> MonitoringCollector:
        """Get or create a collector instance, cached by DSN."""
        dsn = config.dsn
        if dsn not in _collectors:
            _collectors[dsn] = cls(config)
        return _collectors[dsn]

    async def ensure_initialized(self) -> None:
        """Lazily initialize the DB pool and repository."""
        if self._initialized:
            return

        try:
            from fluxibly.monitoring.db.connection import get_pool
            from fluxibly.monitoring.db.repository import MonitoringRepository

            pool = await get_pool(self.config.to_connect_kwargs())
            self._repo = MonitoringRepository(pool)
            await self._writer.start()
            self._initialized = True
        except Exception as e:
            logger.error(
                f"Monitoring initialization failed (continuing without monitoring): {e}"
            )
            self._initialized = False

    # ══════════════════════════════════════════════════════════════
    # Trace Management
    # ══════════════════════════════════════════════════════════════

    async def start_trace(
        self,
        entry_input: str,
        service_name: str = "default",
        environment: str = "development",
        tags: dict[str, str] | None = None,
    ) -> TraceContext:
        """Create a new trace for a top-level request."""
        await self.ensure_initialized()

        trace_id = str(uuid4())
        now = datetime.now(timezone.utc)

        trace = TraceRecord(
            trace_id=trace_id,
            service_name=service_name,
            environment=environment,
            tags=tags or {},
            entry_input=truncate_text(entry_input),
            status="in_progress",
            created_at=now,
            updated_at=now,
        )

        if self._repo:
            await self._writer.write(self._repo.insert_trace, trace)

        return TraceContext(
            trace_id=trace_id,
            service_name=service_name,
            environment=environment,
            tags=tags or {},
            created_at=now,
        )

    async def end_trace(
        self,
        trace_ctx: TraceContext,
        final_output: str,
        status: str,
        error_message: str | None = None,
    ) -> None:
        """Finalize a trace with summary statistics."""
        if not self._repo:
            return

        # Flush pending writes first so span/detail data is available for aggregation
        await self._writer.flush()

        duration_ms = int(
            (datetime.now(timezone.utc) - trace_ctx.created_at).total_seconds()
            * 1000
        )

        # Aggregate totals from spans and LLM details
        totals = await self._repo.aggregate_trace_totals(trace_ctx.trace_id)

        await self._writer.write(
            self._repo.update_trace,
            trace_ctx.trace_id,
            final_output=truncate_text(final_output),
            total_duration_ms=duration_ms,
            total_cost=totals.get("total_cost"),
            total_input_tokens=totals.get("total_input_tokens", 0),
            total_output_tokens=totals.get("total_output_tokens", 0),
            total_tokens=totals.get("total_tokens", 0),
            span_count=totals.get("span_count", 0),
            status=status,
            error_message=error_message,
        )
        # Flush the trace update immediately for visibility
        await self._writer.flush()

        # Clean up sequence counters for this trace to prevent memory leak
        keys_to_remove = [
            k for k in self._seq_counters if k == trace_ctx.trace_id
        ]
        for k in keys_to_remove:
            del self._seq_counters[k]

    # ══════════════════════════════════════════════════════════════
    # Span Management
    # ══════════════════════════════════════════════════════════════

    async def start_span(
        self,
        trace_ctx: TraceContext,
        span_type: str,
        component_class: str,
        component_name: str,
        parent_span_id: str | None = None,
        model: str | None = None,
        provider: str | None = None,
        depth: int = 0,
    ) -> SpanContext:
        """Create a new span under a trace."""
        span_id = str(uuid4())
        now = datetime.now(timezone.utc)

        # Assign sequence number under parent
        parent_key = parent_span_id or trace_ctx.trace_id
        seq = self._seq_counters.get(parent_key, 0)
        self._seq_counters[parent_key] = seq + 1

        span = SpanRecord(
            span_id=span_id,
            trace_id=trace_ctx.trace_id,
            parent_span_id=parent_span_id,
            span_type=span_type,
            component_class=component_class,
            component_name=component_name,
            model=model,
            provider=provider,
            sequence_number=seq,
            started_at=now,
            status="in_progress",
            depth=depth,
        )

        if self._repo:
            await self._writer.write(self._repo.insert_span, span)

        return SpanContext(
            span_id=span_id,
            trace_id=trace_ctx.trace_id,
            parent_span_id=parent_span_id,
            span_type=span_type,
            component_class=component_class,
            component_name=component_name,
            depth=depth,
            sequence_number=seq,
            started_at=now,
        )

    async def end_span(
        self,
        span_ctx: SpanContext,
        status: str,
        error: str | None = None,
    ) -> None:
        """Finalize a span with timing and status."""
        if not self._repo:
            return

        now = datetime.now(timezone.utc)
        duration_ms = int((now - span_ctx.started_at).total_seconds() * 1000)

        await self._writer.write(
            self._repo.update_span,
            span_ctx.span_id,
            ended_at=now,
            duration_ms=duration_ms,
            status=status,
            error_message=error,
        )

    # ══════════════════════════════════════════════════════════════
    # Detail Recording
    # ══════════════════════════════════════════════════════════════

    async def record_llm_details(
        self,
        span_ctx: SpanContext,
        llm_response: LLMResponse,
        config: LLMConfig,
        input_messages: list[dict[str, Any]] | None = None,
        tool_names: list[dict[str, Any]] | None = None,
    ) -> None:
        """Record full LLM call details after forward() completes."""
        if not self._repo:
            return

        usage = llm_response.metadata.usage
        excluded = (
            frozenset(self.config.excluded_fields)
            if self.config.mask_api_keys
            else frozenset()
        )

        # Extract system prompt from messages
        system_prompt = None
        if input_messages:
            for msg in input_messages:
                if msg.get("role") == "system":
                    content = msg.get("content", "")
                    system_prompt = (
                        content if isinstance(content, str) else str(content)
                    )
                    break

        # Sanitize messages if capturing
        stored_messages = None
        if self.config.capture_input and input_messages:
            stored_messages = [
                sanitize_dict(m, excluded) for m in input_messages
            ]

        # Output content
        stored_content = None
        if self.config.capture_output:
            stored_content = [
                item.model_dump(mode="json")
                for item in llm_response.output.content
            ]

        # Config snapshot
        config_snapshot = config.model_dump(mode="json")
        if self.config.mask_api_keys:
            config_snapshot = sanitize_dict(config_snapshot, excluded)

        # Raw response
        raw_response = None
        if (
            self.config.capture_raw_response
            and llm_response.metadata.raw_response
        ):
            raw_response = llm_response.metadata.raw_response

        detail = LLMCallDetailRecord(
            span_id=span_ctx.span_id,
            model=llm_response.metadata.model,
            provider=llm_response.metadata.provider,
            api_type=config.api_type,
            input_messages=stored_messages,
            input_message_count=len(input_messages) if input_messages else 0,
            system_prompt=truncate_text(system_prompt),
            output_text=(
                truncate_text(llm_response.output.output_text)
                if self.config.capture_output
                else None
            ),
            output_content=stored_content,
            reasoning_text=truncate_text(llm_response.output.reasoning),
            stop_reason=llm_response.metadata.stop_reason,
            input_tokens=usage.input_tokens,
            output_tokens=usage.output_tokens,
            reasoning_tokens=usage.reasoning_tokens,
            cached_tokens=usage.cached_tokens,
            total_tokens=usage.total_tokens,
            input_cost=(
                Decimal(str(usage.input_cost))
                if usage.input_cost is not None
                else None
            ),
            output_cost=(
                Decimal(str(usage.output_cost))
                if usage.output_cost is not None
                else None
            ),
            reasoning_cost=(
                Decimal(str(usage.reasoning_cost))
                if usage.reasoning_cost is not None
                else None
            ),
            cached_cost=(
                Decimal(str(usage.cached_cost))
                if usage.cached_cost is not None
                else None
            ),
            total_cost=(
                Decimal(str(usage.total_cost))
                if usage.total_cost is not None
                else None
            ),
            temperature=config.temperature,
            max_output_tokens=config.max_output_tokens,
            top_p=config.top_p,
            reasoning_enabled=bool(
                config.reasoning and config.reasoning.enabled
            ),
            reasoning_effort=(
                config.reasoning.effort if config.reasoning else None
            ),
            streaming=config.streaming,
            response_id=llm_response.metadata.id,
            config_snapshot=config_snapshot,
            raw_response=raw_response,
            tool_names=tool_names,
        )

        await self._writer.write(self._repo.insert_llm_detail, detail)

    async def record_agent_details(
        self,
        span_ctx: SpanContext,
        agent_response: AgentResponse,
        config: AgentConfig,
        context: dict[str, Any] | None = None,
        input_messages: list[dict[str, Any]] | None = None,
        iteration_count: int = 0,
        total_llm_calls: int = 0,
        tool_names: list[dict[str, Any]] | None = None,
        tool_interactions: list[dict[str, Any]] | None = None,
    ) -> None:
        """Record full Agent call details after forward() completes."""
        if not self._repo:
            return

        excluded = (
            frozenset(self.config.excluded_fields)
            if self.config.mask_api_keys
            else frozenset()
        )

        # Sanitize context
        stored_context = None
        if self.config.capture_context and context:
            # Remove _monitoring from stored context
            ctx_copy = {k: v for k, v in context.items() if k != "_monitoring"}
            stored_context = sanitize_dict(ctx_copy, excluded)

        # Sanitize messages
        stored_messages = None
        if self.config.capture_input and input_messages:
            stored_messages = [
                sanitize_dict(m, excluded) for m in input_messages
            ]

        # Output content (full structured LLMResponse)
        stored_output_content = None
        if self.config.capture_output:
            stored_output_content = agent_response.content.model_dump(
                mode="json"
            )

        # Config snapshot
        config_snapshot = config.model_dump(mode="json")
        if self.config.mask_api_keys:
            config_snapshot = sanitize_dict(config_snapshot, excluded)

        detail = AgentCallDetailRecord(
            span_id=span_ctx.span_id,
            agent_class=span_ctx.component_class,
            agent_name=config.name,
            system_prompt=truncate_text(config.system_prompt),
            user_prompt_template=truncate_text(config.user_prompt),
            input_messages=stored_messages,
            input_message_count=len(input_messages) if input_messages else 0,
            output_text=(
                truncate_text(agent_response.content.output.output_text)
                if self.config.capture_output
                else None
            ),
            output_content=stored_output_content,
            context=stored_context,
            tool_count=(
                len(tool_names) if tool_names else len(config.tools or [])
            ),
            tool_names=tool_names,
            sub_agent_count=len(config.agents or []),
            iteration_count=iteration_count,
            total_llm_calls=total_llm_calls,
            config_snapshot=config_snapshot,
            tool_interactions=tool_interactions,
        )

        await self._writer.write(self._repo.insert_agent_detail, detail)

    async def record_tool_call(
        self,
        span_ctx: SpanContext,
        tool_call: ToolCall,
        result: ToolResult,
        duration_ms: int,
        iteration: int = 1,
        sequence_in_batch: int = 0,
    ) -> None:
        """Record a single tool call execution."""
        if not self._repo:
            return

        # Determine tool type
        tool_type = "function"
        if tool_call.name.startswith("agent_"):
            tool_type = "agent_delegation"
        elif tool_call.type != "function":
            tool_type = tool_call.type

        # Sanitize arguments
        arguments = None
        if self.config.capture_tool_io:
            args = tool_call.arguments
            if isinstance(args, str):
                import json

                try:
                    args = json.loads(args)
                except (json.JSONDecodeError, TypeError):
                    args = {"raw": args}
            if isinstance(args, dict):
                excluded = (
                    frozenset(self.config.excluded_fields)
                    if self.config.mask_api_keys
                    else frozenset()
                )
                arguments = sanitize_dict(args, excluded)
            else:
                arguments = {"value": str(args)}

        result_content = None
        if self.config.capture_tool_io:
            content = result.content
            result_content = truncate_text(
                content if isinstance(content, str) else str(content)
            )

        now = datetime.now(timezone.utc)

        detail = ToolCallDetailRecord(
            span_id=span_ctx.span_id,
            tool_call_id=tool_call.id,
            tool_name=tool_call.name,
            tool_type=tool_type,
            arguments=arguments,
            result_content=result_content,
            is_error=result.is_error,
            started_at=now,
            ended_at=now,
            duration_ms=duration_ms,
            iteration_number=iteration,
            sequence_in_batch=sequence_in_batch,
        )

        await self._writer.write(self._repo.insert_tool_detail, detail)

    # ══════════════════════════════════════════════════════════════
    # Lifecycle
    # ══════════════════════════════════════════════════════════════

    async def shutdown(self) -> None:
        """Gracefully shutdown the collector."""
        await self._writer.stop()
        self._initialized = False

    @classmethod
    async def shutdown_all(cls) -> None:
        """Shutdown all cached collectors."""
        for collector in _collectors.values():
            await collector.shutdown()
        _collectors.clear()
