"""
Tool Call Logger for Autonomous Agents

Logs all tool calls to the monitoring database (dwd.agent_tool_calls)
for Grafana dashboard visualization.
"""

from __future__ import annotations

import asyncio
import logging
import os
import time
from datetime import datetime, timezone
from typing import Any

logger = logging.getLogger(__name__)

# Tool category mapping
TOOL_CATEGORIES = {
    # Health tools
    "check_pipeline_health": "health",
    "check_single_pipeline": "health",
    "check_layer_health": "health",
    "get_system_health": "health",
    # ETL tools
    "run_etl": "etl",
    "run_multiple_etls": "etl",
    "list_available_modules": "etl",
    # Alert tools
    "send_alert": "alert",
    "request_human_approval": "alert",
    "get_audit_log": "alert",
    # Circuit breaker tools
    "get_circuit_status": "circuit",
    "list_circuits": "circuit",
    "list_open_circuits": "circuit",
    "reset_circuit": "circuit",
    "configure_circuit": "circuit",
    # Anomaly tools
    "detect_anomalies": "anomaly",
    "get_anomaly_history": "anomaly",
    # Healing tools
    "heal_circuit": "healing",
    "heal_stale_data": "healing",
    "heal_quality_issue": "healing",
    "get_healing_history": "healing",
    # Quality tools
    "list_quality_rules": "quality",
    "check_data_quality": "quality",
    "get_quality_statistics": "quality",
    "add_quality_rule": "quality",
    # Lineage tools
    "list_lineage_graphs": "lineage",
    "get_lineage_graph": "lineage",
    "trace_upstream": "lineage",
    "trace_downstream": "lineage",
    # Contract tools
    "list_contracts": "contract",
    "get_contract": "contract",
    "check_contract_compliance": "contract",
    "create_contract": "contract",
    # Metadata tools
    "list_schemas": "metadata",
    "get_schema": "metadata",
    "list_tags": "metadata",
    "get_metadata_statistics": "metadata",
}


async def log_tool_call(
    execution_id: str,
    tool_name: str,
    input_params: dict[str, Any] | None = None,
    output_result: Any = None,
    success: bool = True,
    duration_ms: float = 0.0,
    error_message: str | None = None,
) -> None:
    """
    Log a tool call to the monitoring database.

    Args:
        execution_id: Parent execution ID (from scheduler)
        tool_name: Name of the tool called
        input_params: Tool input parameters
        output_result: Tool output result
        success: Whether the call succeeded
        duration_ms: Call duration in milliseconds
        error_message: Error message if failed
    """
    try:
        import asyncpg
        import json

        db_url = os.environ.get(
            "TIMESCALEDB_URL",
            "postgresql://sweetnight:sweetnight123@localhost:5433/sweetnight_geo"
        )
        # asyncpg requires plain postgresql:// scheme
        db_url = db_url.replace("postgresql+asyncpg://", "postgresql://")

        # Get tool category
        tool_category = TOOL_CATEGORIES.get(tool_name, "other")

        # Serialize params and result to JSON
        input_json = None
        output_json = None

        if input_params:
            try:
                input_json = json.dumps(input_params, default=str)
            except Exception:
                input_json = str(input_params)

        if output_result is not None:
            try:
                # Truncate large results
                result_str = json.dumps(output_result, default=str)
                if len(result_str) > 10000:
                    result_str = result_str[:10000] + "...(truncated)"
                output_json = result_str
            except Exception:
                output_json = str(output_result)[:10000]

        conn = await asyncpg.connect(db_url)
        try:
            await conn.execute(
                """
                INSERT INTO dwd.agent_tool_calls (
                    execution_id, tool_name, tool_category,
                    input_params, output_result, success,
                    duration_ms, error_message, called_at
                ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
                """,
                execution_id,
                tool_name,
                tool_category,
                input_json,
                output_json,
                success,
                duration_ms,
                error_message,
                datetime.now(timezone.utc),
            )
            logger.debug(f"Logged tool call: {tool_name}")
        finally:
            await conn.close()

    except Exception as e:
        # Don't fail agent execution due to logging errors
        logger.warning(f"Failed to log tool call: {e}")


class ToolCallTracker:
    """
    Context manager for tracking tool call duration and logging.

    Usage:
        async with ToolCallTracker(execution_id, "check_health", {"module": "ga4"}) as tracker:
            result = await some_tool_function()
            tracker.set_result(result)
    """

    def __init__(
        self,
        execution_id: str,
        tool_name: str,
        input_params: dict[str, Any] | None = None,
    ):
        self.execution_id = execution_id
        self.tool_name = tool_name
        self.input_params = input_params
        self.result = None
        self.success = True
        self.error_message = None
        self._start_time = None

    async def __aenter__(self):
        self._start_time = time.perf_counter()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        duration_ms = (time.perf_counter() - self._start_time) * 1000

        if exc_val:
            self.success = False
            self.error_message = str(exc_val)

        await log_tool_call(
            execution_id=self.execution_id,
            tool_name=self.tool_name,
            input_params=self.input_params,
            output_result=self.result,
            success=self.success,
            duration_ms=duration_ms,
            error_message=self.error_message,
        )

        # Don't suppress exceptions
        return False

    def set_result(self, result: Any) -> None:
        """Set the result of the tool call."""
        self.result = result

    def set_error(self, error: str) -> None:
        """Mark the call as failed with an error message."""
        self.success = False
        self.error_message = error
