"""
Agent回调工具 - 用于向Agent Manager回调运行时长记录
"""
import os
import time
import logging
import requests
from typing import Optional, List
from datetime import datetime, timezone

logger = logging.getLogger(__name__)


class AgentCallbackHandler:
    """Agent回调处理器"""

    def __init__(
        self,
        agent_name: Optional[str] = None,
        user_id: Optional[str] = None,
        callback_url: Optional[str] = None
    ):
        self.agent_name = agent_name or os.getenv("POD_NAME", "unknown-agent")
        self.user_id = user_id or os.getenv("USER_ID", "")
        self.callback_url = callback_url or os.getenv(
            "AGENT_CALLBACK_URL",
            "http://mcp-server.taiji-ai.svc.cluster.local:8000/api/v1/billing/agent-callback"
        )

        self.start_time: Optional[datetime] = None
        self.tools_used: List[str] = []
        self.request_id: Optional[str] = None

        logger.info(
            "AgentCallbackHandler initialized: agent=%s, callback_url=%s",
            self.agent_name,
            self.callback_url,
        )

    def start_request(self, request_id: Optional[str] = None, user_id: Optional[str] = None):
        self.start_time = datetime.now(timezone.utc)
        self.tools_used = []
        self.request_id = request_id or f"req-{int(time.time())}"

        if user_id:
            self.user_id = user_id

        logger.info("Request started: request_id=%s, user_id=%s", self.request_id, self.user_id)

    def add_tool_used(self, tool_name: str):
        if tool_name not in self.tools_used:
            self.tools_used.append(tool_name)
            logger.debug("Tool used: %s", tool_name)

    def end_request(self, tools_used: Optional[List[str]] = None) -> bool:
        if not self.start_time:
            logger.warning("Cannot end request: no start time recorded")
            return False

        if not self.user_id:
            logger.warning("Cannot send callback: user_id not set")
            return False

        end_time = datetime.now(timezone.utc)
        running_time = (end_time - self.start_time).total_seconds()
        final_tools_used = tools_used if tools_used is not None else self.tools_used

        success = self._send_callback(
            running_time_seconds=int(running_time),
            start_time=self.start_time,
            end_time=end_time,
            tools_used=final_tools_used
        )

        self.start_time = None
        self.tools_used = []
        self.request_id = None

        return success

    def _send_callback(
        self,
        running_time_seconds: int,
        start_time: datetime,
        end_time: datetime,
        tools_used: List[str]
    ) -> bool:
        try:
            payload = {
                "agentName": self.agent_name,
                "userId": self.user_id,
                "podRunningTimeSeconds": running_time_seconds,
                "toolsUsed": tools_used,
                "startTime": start_time.isoformat(),
                "endTime": end_time.isoformat(),
                "requestId": self.request_id
            }

            logger.info("Sending callback: %s", payload)

            response = requests.post(
                self.callback_url,
                json=payload,
                timeout=5
            )

            if response.status_code == 200:
                logger.info("Callback sent successfully: %s", response.json())
                return True

            logger.error("Callback failed with status %s: %s", response.status_code, response.text)
            return False

        except requests.exceptions.RequestException as e:
            logger.error("Failed to send callback: %s", str(e))
            return False
        except Exception as e:
            logger.error("Unexpected error sending callback: %s", str(e))
            return False


class CallbackContextManager:
    """回调上下文管理器 - 使用with语句自动处理开始和结束"""

    def __init__(
        self,
        handler: AgentCallbackHandler,
        request_id: Optional[str] = None,
        user_id: Optional[str] = None,
        tools_used: Optional[List[str]] = None
    ):
        self.handler = handler
        self.request_id = request_id
        self.user_id = user_id
        self.tools_used = tools_used or []

    def __enter__(self):
        self.handler.start_request(
            request_id=self.request_id,
            user_id=self.user_id
        )
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.handler.end_request(tools_used=self.tools_used)
        return False

    def add_tool(self, tool_name: str):
        self.handler.add_tool_used(tool_name)
        if tool_name not in self.tools_used:
            self.tools_used.append(tool_name)