"""
Echo Agent - 简单的回显测试代理
用于测试 Agent Manager 的部署功能
"""
import os
import sys
import logging
from datetime import datetime
from typing import Optional, Any

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import uvicorn

# 添加 common 模块路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 导入回调工具
try:
    from common.agent_callback_utils import AgentCallbackHandler, CallbackContextManager
    CALLBACK_ENABLED = True
except ImportError:
    CALLBACK_ENABLED = False
    AgentCallbackHandler = None
    CallbackContextManager = None

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 环境变量
SERVICE_HOST = os.getenv("SERVICE_HOST", "0.0.0.0")
SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8000"))
POD_NAME = os.getenv("POD_NAME", "echo-agent")
USER_ID = os.getenv("USER_ID", "")

# FastAPI 应用
app = FastAPI(
    title="Echo Agent",
    description="简单的回显测试代理，用于验证 Agent Manager 部署功能",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 回调处理器
callback_handler: Optional[AgentCallbackHandler] = None


# ==================== 请求/响应模型 ====================

class EchoRequest(BaseModel):
    """回显请求"""
    message: str = Field(..., description="要回显的消息")
    user_id: Optional[str] = Field(None, description="用户ID（用于计费回调）")
    metadata: Optional[dict] = Field(None, description="附加元数据")


class EchoResponse(BaseModel):
    """回显响应"""
    message: str
    echo: str
    pod_name: str
    metadata: Optional[dict] = None
    timestamp: str


class HealthResponse(BaseModel):
    """健康检查响应"""
    status: str
    pod_name: str
    version: str
    callback_enabled: bool
    timestamp: str


# ==================== 生命周期 ====================

@app.on_event("startup")
async def startup_event():
    """应用启动时初始化回调处理器"""
    global callback_handler
    
    if CALLBACK_ENABLED:
        callback_handler = AgentCallbackHandler(
            agent_name=POD_NAME,
            user_id=USER_ID
        )
        logger.info(f"回调处理器已初始化: callback_url={callback_handler.callback_url}")
    else:
        logger.warning("回调模块未加载，计费回调功能不可用")


# ==================== API 端点 ====================

@app.get("/health", response_model=HealthResponse)
@app.get("/", response_model=HealthResponse)
async def health_check():
    """健康检查"""
    return HealthResponse(
        status="healthy",
        pod_name=POD_NAME,
        version="1.0.0",
        callback_enabled=CALLBACK_ENABLED,
        timestamp=datetime.utcnow().isoformat()
    )


@app.post("/echo", response_model=EchoResponse)
async def echo(request: EchoRequest):
    """回显消息"""
    logger.info(f"Echo: {request.message}")
    
    # 使用回调上下文管理器
    if CALLBACK_ENABLED and callback_handler and request.user_id:
        with CallbackContextManager(
            handler=callback_handler,
            user_id=request.user_id,
            request_id=f"echo-{int(datetime.utcnow().timestamp())}"
        ) as ctx:
            ctx.add_tool("echo")
            
            response = EchoResponse(
                message=request.message,
                echo=f"[Echo from {POD_NAME}] {request.message}",
                pod_name=POD_NAME,
                metadata=request.metadata,
                timestamp=datetime.utcnow().isoformat()
            )
    else:
        response = EchoResponse(
            message=request.message,
            echo=f"[Echo from {POD_NAME}] {request.message}",
            pod_name=POD_NAME,
            metadata=request.metadata,
            timestamp=datetime.utcnow().isoformat()
        )
    
    return response


@app.get("/echo")
async def echo_get(message: str = "Hello", user_id: Optional[str] = None):
    """GET 方式回显"""
    return await echo(EchoRequest(message=message, user_id=user_id))


@app.get("/info")
async def get_info():
    """获取 Agent 信息"""
    return {
        "agent_name": "Echo Agent",
        "pod_name": POD_NAME,
        "version": "1.0.0",
        "callback_enabled": CALLBACK_ENABLED,
        "callback_url": callback_handler.callback_url if callback_handler else None,
        "capabilities": ["echo", "health_check"],
        "environment": {
            "SERVICE_HOST": SERVICE_HOST,
            "SERVICE_PORT": SERVICE_PORT
        },
        "timestamp": datetime.utcnow().isoformat()
    }


# ==================== 主入口 ====================

def main():
    """主函数"""
    logger.info(f"启动 Echo Agent - {POD_NAME}")
    logger.info(f"服务地址: {SERVICE_HOST}:{SERVICE_PORT}")
    logger.info(f"回调功能: {'已启用' if CALLBACK_ENABLED else '未启用'}")
    
    uvicorn.run(app, host=SERVICE_HOST, port=SERVICE_PORT, log_level="info")


if __name__ == "__main__":
    main()