"""
A2A协议兼容的Search Agent服务

实现Google Agent2Agent协议规范
支持从请求传入 API key，也支持从环境变量获取
"""
import asyncio
import json
import uuid
import os
from typing import Optional, Dict, Any, AsyncGenerator
from datetime import datetime
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, Request, Response
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from loguru import logger

from agent import SearchAgentWrapper
from config import get_config, AgentConfig, A2AConfig

# 环境变量配置
SERVICE_HOST = os.getenv("SERVICE_HOST", "0.0.0.0")
SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8080"))
POD_NAME = os.getenv("POD_NAME", "search-agent-a2a")
TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "search_agent_A2A")

# ============== A2A 协议数据模型 ==============


class A2APart(BaseModel):
    """A2A消息部分"""
    kind: str = "text"
    text: Optional[str] = None
    data: Optional[Dict[str, Any]] = None
    mime_type: Optional[str] = None


class A2AMessage(BaseModel):
    """A2A消息"""
    role: str
    parts: list[A2APart]
    messageId: str = Field(default_factory=lambda: uuid.uuid4().hex)


class A2AMessageSendParams(BaseModel):
    """A2A发送消息参数"""
    message: A2AMessage
    configuration: Optional[Dict[str, Any]] = None
    api_key: Optional[str] = Field(None, description="LiteLLM API密钥（可选，优先使用，否则从环境变量获取）")
    model: Optional[str] = Field(None, description="模型名称（可选，优先使用，否则从环境变量获取）")


class A2ARequest(BaseModel):
    """A2A JSON-RPC请求"""
    jsonrpc: str = "2.0"
    id: str
    method: str
    params: Optional[Dict[str, Any]] = None


class A2AArtifact(BaseModel):
    """A2A响应工件"""
    artifactId: str = Field(default_factory=lambda: uuid.uuid4().hex)
    name: str = "response"
    parts: list[A2APart]


class A2ATaskStatus(BaseModel):
    """A2A任务状态"""
    state: str  # submitted, working, input-required, completed, failed, canceled
    timestamp: str = Field(default_factory=lambda: datetime.utcnow().isoformat() + "Z")
    message: Optional[str] = None


class A2ATask(BaseModel):
    """A2A任务"""
    kind: str = "task"
    id: str = Field(default_factory=lambda: uuid.uuid4().hex)
    contextId: str = Field(default_factory=lambda: uuid.uuid4().hex)
    status: A2ATaskStatus
    artifacts: Optional[list[A2AArtifact]] = None


class A2AResponse(BaseModel):
    """A2A JSON-RPC响应"""
    jsonrpc: str = "2.0"
    id: str
    result: Optional[A2ATask] = None
    error: Optional[Dict[str, Any]] = None


class A2AStreamEvent(BaseModel):
    """A2A流式事件"""
    kind: str
    taskId: str
    contextId: str
    data: Optional[Dict[str, Any]] = None


# ============== Agent Card ==============


class AgentSkill(BaseModel):
    """Agent技能"""
    id: str
    name: str
    description: str
    inputSchema: Optional[Dict[str, Any]] = None
    outputSchema: Optional[Dict[str, Any]] = None


class AgentCapabilities(BaseModel):
    """Agent能力"""
    text: bool = True
    streaming: bool = True
    push_notifications: bool = False
    forms: bool = False
    files: bool = False


class AgentCard(BaseModel):
    """A2A Agent Card - 描述Agent能力"""
    name: str
    description: str
    version: str
    url: str
    capabilities: AgentCapabilities
    skills: list[AgentSkill]
    authentication: Optional[Dict[str, Any]] = None


# ============== A2A Server ==============


class A2ASearchAgentServer:
    """A2A协议Search Agent服务器"""
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        model: Optional[str] = None
    ):
        """
        初始化A2A Search Agent服务器
        
        Args:
            api_key: LiteLLM API密钥（可选，优先使用，否则从环境变量获取）
            model: 模型名称（可选，优先使用，否则从环境变量获取）
        """
        # 获取配置（用于服务初始化，实际处理请求时使用请求中的api_key）
        # 注意：这里的api_key和model仅用于服务启动验证，实际请求时会使用请求中的api_key
        self.llm_config, self.agent_config, self.a2a_config = get_config(api_key, model)
        
        # 任务存储
        self.tasks: Dict[str, A2ATask] = {}
        
        # 创建FastAPI应用
        self.app = self._create_app()
    
    def _create_app(self) -> FastAPI:
        """创建FastAPI应用"""
        
        @asynccontextmanager
        async def lifespan(app: FastAPI):
            logger.info("A2A Search Agent服务启动", agent_name=self.agent_config.name)
            yield
            # 注意：每个请求创建的Agent实例在请求结束时已关闭，这里不需要额外清理
            logger.info("A2A Search Agent服务关闭")
        
        app = FastAPI(
            title=f"{self.agent_config.name} - A2A Agent",
            description=self.agent_config.description,
            version=self.agent_config.version,
            lifespan=lifespan
        )
        
        # CORS中间件
        app.add_middleware(
            CORSMiddleware,
            allow_origins=["*"],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )
        
        # 注册路由
        self._register_routes(app)
        
        return app
    
    def _get_agent(self, api_key: str, model: Optional[str] = None) -> SearchAgentWrapper:
        """
        获取Agent实例
        
        Args:
            api_key: 用户的API密钥（必需，从请求参数中获取）
            model: 模型名称（可选，从环境变量获取）
        
        Returns:
            SearchAgentWrapper实例
        """
        # api_key必须提供（来自请求），model从环境变量获取（如果未提供）
        if not model:
            model = os.getenv("MODEL_NAME") or os.getenv("LLM_MODEL") or os.getenv("LITELLM_MODEL")
        
        # 创建新的配置和Agent（每次都创建新的，使用请求中的api_key）
        llm_config, agent_config, _ = get_config(api_key=api_key, model=model)
        return SearchAgentWrapper(
            litellm_config=llm_config,
            agent_config=agent_config,
            api_key=api_key,
            model=model
        )
    
    def _register_routes(self, app: FastAPI):
        """注册A2A协议路由"""
        
        @app.get("/")
        async def root():
            """服务根路径"""
            return {
                "name": self.agent_config.name,
                "version": self.agent_config.version,
                "protocol": "A2A",
                "status": "running",
                "pod_name": POD_NAME,
                "template_type": TEMPLATE_TYPE
            }
        
        @app.get("/health")
        async def health_check():
            """健康检查"""
            return {
                "status": "healthy",
                "pod_name": POD_NAME,
                "template_type": TEMPLATE_TYPE,
                "configured": self.llm_config.api_key is not None,
                "timestamp": datetime.utcnow().isoformat()
            }
        
        @app.get("/.well-known/agent.json")
        async def get_agent_card(request: Request):
            """获取Agent Card (A2A发现协议)"""
            base_url = str(request.base_url).rstrip("/")
            
            card = AgentCard(
                name=self.agent_config.name,
                description=self.agent_config.description,
                version=self.agent_config.version,
                url=base_url,
                capabilities=AgentCapabilities(
                    text=True,
                    streaming=self.agent_config.enable_streaming,
                    push_notifications=False
                ),
                skills=[
                    AgentSkill(
                        id="intelligent-search",
                        name="智能搜索",
                        description="理解用户查询意图，自动规划搜索策略，从多个来源获取信息并生成高质量、有来源引用的答案"
                    )
                ]
            )
            return card.model_dump()
        
        @app.post("/message/send")
        async def send_message(request: Request):
            """A2A message/send 端点"""
            body = await request.json()
            
            # 解析JSON-RPC请求
            try:
                rpc_request = A2ARequest(**body)
            except Exception as e:
                return JSONResponse({
                    "jsonrpc": "2.0",
                    "id": body.get("id", "unknown"),
                    "error": {
                        "code": -32600,
                        "message": f"Invalid Request: {str(e)}"
                    }
                })
            
            # 处理 message/send 方法
            if rpc_request.method == "message/send":
                return await self._handle_message_send(rpc_request)
            elif rpc_request.method == "message/stream":
                return await self._handle_message_stream(rpc_request)
            else:
                return JSONResponse({
                    "jsonrpc": "2.0",
                    "id": rpc_request.id,
                    "error": {
                        "code": -32601,
                        "message": f"Method not found: {rpc_request.method}"
                    }
                })
        
        @app.post("/message/stream")
        async def stream_message(request: Request):
            """A2A message/stream 端点 (SSE流式响应)"""
            body = await request.json()
            
            try:
                rpc_request = A2ARequest(**body)
            except Exception as e:
                return JSONResponse({
                    "jsonrpc": "2.0",
                    "id": body.get("id", "unknown"),
                    "error": {
                        "code": -32600,
                        "message": f"Invalid Request: {str(e)}"
                    }
                })
            
            return await self._handle_message_stream(rpc_request)
        
        @app.get("/tasks/{task_id}")
        async def get_task(task_id: str):
            """获取任务状态"""
            if task_id not in self.tasks:
                raise HTTPException(status_code=404, detail="Task not found")
            return self.tasks[task_id].model_dump()
    
    async def _handle_message_send(self, request: A2ARequest) -> JSONResponse:
        """处理 message/send 请求"""
        params = request.params or {}
        message_data = params.get("message", {})
        
        # 提取API key（必须从请求参数中获取，等同于API格式版本的llm_api_key）
        api_key = params.get("api_key")
        if not api_key:
            return JSONResponse({
                "jsonrpc": "2.0",
                "id": request.id,
                "error": {
                    "code": -32602,
                    "message": "Invalid params: api_key is required"
                }
            })
        
        # 提取model（从环境变量获取，不支持在请求中传递，与API格式版本保持一致）
        # 支持多种环境变量名称：MODEL_NAME（优先）、LLM_MODEL（AKS部署）、LITELLM_MODEL
        model = os.getenv("MODEL_NAME") or os.getenv("LLM_MODEL") or os.getenv("LITELLM_MODEL")
        if not model:
            return JSONResponse({
                "jsonrpc": "2.0",
                "id": request.id,
                "error": {
                    "code": -32000,
                    "message": "Model not configured: MODEL_NAME or LLM_MODEL environment variable is required"
                }
            })
        
        # 提取用户消息文本
        user_text = ""
        parts = message_data.get("parts", [])
        for part in parts:
            if part.get("kind") == "text":
                user_text += part.get("text", "")
        
        if not user_text:
            return JSONResponse({
                "jsonrpc": "2.0",
                "id": request.id,
                "error": {
                    "code": -32602,
                    "message": "Invalid params: no text content found"
                }
            })
        
        # 创建任务
        task_id = uuid.uuid4().hex
        context_id = params.get("contextId", uuid.uuid4().hex)
        
        task = A2ATask(
            id=task_id,
            contextId=context_id,
            status=A2ATaskStatus(state="working")
        )
        self.tasks[task_id] = task
        
        try:
            # 获取Agent实例（使用请求中的api_key和环境变量中的model）
            agent = self._get_agent(api_key=api_key, model=model)
            
            # 调用Agent获取响应
            logger.info("处理搜索消息", task_id=task_id, message_preview=user_text[:50])
            
            response = await agent.search(query=user_text)
            
            # 如果创建了新Agent（使用了请求中的api_key），关闭它
            await agent.close()
            
            # 构建答案文本（包含来源信息）
            answer_parts = [response.answer.content]
            
            if response.answer.sources:
                answer_parts.append("\n\n## 来源")
                for i, source in enumerate(response.answer.sources, 1):
                    answer_parts.append(f"{i}. [{source.title}]({source.url})")
            
            answer_text = "\n".join(answer_parts)
            
            # 更新任务状态
            task.status = A2ATaskStatus(state="completed")
            task.artifacts = [
                A2AArtifact(
                    name="response",
                    parts=[A2APart(kind="text", text=answer_text)]
                )
            ]
            self.tasks[task_id] = task
            
            return JSONResponse({
                "jsonrpc": "2.0",
                "id": request.id,
                "result": task.model_dump()
            })
            
        except Exception as e:
            logger.error("处理搜索消息失败", error=str(e))
            task.status = A2ATaskStatus(state="failed", message=str(e))
            self.tasks[task_id] = task
            
            return JSONResponse({
                "jsonrpc": "2.0",
                "id": request.id,
                "error": {
                    "code": -32000,
                    "message": f"Agent error: {str(e)}"
                }
            })
    
    async def _handle_message_stream(self, request: A2ARequest) -> StreamingResponse:
        """处理 message/stream 请求 (SSE)"""
        params = request.params or {}
        message_data = params.get("message", {})
        
        # 提取API key（必须从请求参数中获取）
        api_key = params.get("api_key")
        if not api_key:
            # 对于流式响应，需要通过SSE发送错误
            async def error_generator():
                error_event = {
                    "kind": "task-error",
                    "taskId": "unknown",
                    "contextId": "unknown",
                    "data": {
                        "error": "Invalid params: api_key is required"
                    }
                }
                yield f"data: {json.dumps(error_event)}\n\n"
            return StreamingResponse(
                error_generator(),
                media_type="text/event-stream"
            )
        
        # 提取model（从环境变量获取）
        model = os.getenv("MODEL_NAME") or os.getenv("LLM_MODEL") or os.getenv("LITELLM_MODEL")
        if not model:
            async def error_generator():
                error_event = {
                    "kind": "task-error",
                    "taskId": "unknown",
                    "contextId": "unknown",
                    "data": {
                        "error": "Model not configured: MODEL_NAME or LLM_MODEL environment variable is required"
                    }
                }
                yield f"data: {json.dumps(error_event)}\n\n"
            return StreamingResponse(
                error_generator(),
                media_type="text/event-stream"
            )
        
        # 提取用户消息
        user_text = ""
        parts = message_data.get("parts", [])
        for part in parts:
            if part.get("kind") == "text":
                user_text += part.get("text", "")
        
        task_id = uuid.uuid4().hex
        context_id = params.get("contextId", uuid.uuid4().hex)
        
        async def event_generator() -> AsyncGenerator[str, None]:
            """生成SSE事件流"""
            agent = None
            try:
                # 获取Agent实例（使用请求中的api_key和环境变量中的model）
                agent = self._get_agent(api_key=api_key, model=model)
                
                # 发送任务开始事件
                start_event = {
                    "kind": "task-start",
                    "taskId": task_id,
                    "contextId": context_id
                }
                yield f"data: {json.dumps(start_event)}\n\n"
                
                # 执行搜索（SearchAgent不支持流式，所以发送完整结果）
                response = await agent.search(query=user_text)
                
                # 构建答案文本
                answer_parts = [response.answer.content]
                
                if response.answer.sources:
                    answer_parts.append("\n\n## 来源")
                    for i, source in enumerate(response.answer.sources, 1):
                        answer_parts.append(f"{i}. [{source.title}]({source.url})")
                
                answer_text = "\n".join(answer_parts)
                
                # 发送完整答案（作为增量发送，以便显示进度）
                # 将答案分成小块发送以模拟流式效果
                chunk_size = 100
                for i in range(0, len(answer_text), chunk_size):
                    chunk = answer_text[i:i + chunk_size]
                    delta_event = {
                        "kind": "artifact-delta",
                        "taskId": task_id,
                        "contextId": context_id,
                        "data": {
                            "kind": "text",
                            "text": chunk
                        }
                    }
                    yield f"data: {json.dumps(delta_event)}\n\n"
                    # 添加小延迟以模拟真实流式效果
                    await asyncio.sleep(0.01)
                
                # 发送完成事件
                complete_event = {
                    "kind": "task-complete",
                    "taskId": task_id,
                    "contextId": context_id,
                    "data": {
                        "status": "completed",
                        "artifacts": [{
                            "name": "response",
                            "parts": [{"kind": "text", "text": answer_text}]
                        }]
                    }
                }
                yield f"data: {json.dumps(complete_event)}\n\n"
                
            except Exception as e:
                # 发送错误事件
                error_event = {
                    "kind": "task-error",
                    "taskId": task_id,
                    "contextId": context_id,
                    "data": {
                        "error": str(e)
                    }
                }
                yield f"data: {json.dumps(error_event)}\n\n"
            finally:
                # 如果创建了新Agent（使用了请求中的api_key），关闭它
                if agent:
                    await agent.close()
        
        return StreamingResponse(
            event_generator(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no"
            }
        )
    
    def run(self, host: Optional[str] = None, port: Optional[int] = None):
        """运行服务器"""
        import uvicorn
        
        host = host or self.agent_config.host
        port = port or self.agent_config.port
        
        logger.info(f"启动A2A Search Agent服务", host=host, port=port)
        uvicorn.run(self.app, host=host, port=port)


def create_app(api_key: Optional[str] = None, model: Optional[str] = None) -> FastAPI:
    """
    创建FastAPI应用（用于uvicorn启动）
    
    使用方式:
        uvicorn a2a_server:app --host 0.0.0.0 --port 8080
    
    或设置环境变量后:
        export LITELLM_API_KEY="your-key"
        export MODEL_NAME="your-model"
        uvicorn a2a_server:app --host 0.0.0.0 --port 8080
    """
    server = A2ASearchAgentServer(api_key=api_key, model=model)
    return server.app


# uvicorn 启动入口
# 环境变量: LITELLM_API_KEY, MODEL_NAME (或 LITELLM_MODEL)
# 注意: app 只在 main.py 中创建，避免导入时立即执行验证
# app = create_app()