"""
FastAPI服务 - Facebook搜索智能Agent
提供搜索API接口和MCP协议支持
"""

import json
import uuid
from typing import Dict, Any, Optional, AsyncGenerator
from fastapi import FastAPI, HTTPException, Request, Header, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel, Field
from datetime import datetime
import sys
from typing import Optional
from loguru import logger

# 支持相对导入和绝对导入
try:
    from .config import Config
    from .agent import FacebookAgent
    from .models.schemas import SearchRequest, SearchResponse
    from .mcp_server import search_facebook, initialize_agent
except ImportError:
    # 如果相对导入失败，尝试绝对导入
    from config import Config
    from agent import FacebookAgent
    from models.schemas import SearchRequest, SearchResponse
    from mcp_server import search_facebook, initialize_agent


# ==================== FastAPI应用 ====================

# 创建FastAPI应用
app = FastAPI(
    title="Facebook搜索智能Agent API",
    description="提供Facebook内容搜索服务和MCP协议支持，基于Pydantic AI框架和LiteLLM Gateway",
    version="1.0.0"
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 全局变量
config: Optional[Config] = None
agent: Optional[FacebookAgent] = None

# MCP 工具映射
TOOL_MAP = {
    'search_facebook': search_facebook,
}

# Session 管理
sessions: Dict[str, Dict[str, Any]] = {}


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


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


async def handle_mcp_request(request_data: Dict[str, Any], session_id: Optional[str] = None, api_key: Optional[str] = None) -> Dict[str, Any]:
    """处理 MCP 请求"""
    method = request_data.get("method")
    params = request_data.get("params", {})
    request_id = request_data.get("id")
    
    # 对于 tools/call 方法，需要验证 API key
    if method == "tools/call":
        if not api_key or api_key.strip() == "" or api_key.strip() == "sk":
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "error": {
                    "code": -32001,
                    "message": "缺少 API key。请在请求头中提供 'api-key' 或 'Authorization: Bearer <token>'。"
                }
            }
    
    try:
        if method == "initialize":
            # 初始化会话
            if not session_id:
                session_id = str(uuid.uuid4())
                sessions[session_id] = {
                    "initialized": True,
                    "capabilities": {}
                }
            
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {
                        "tools": {},
                        "resources": {}
                    },
                    "serverInfo": {
                        "name": "Facebook搜索Agent",
                        "version": "1.0.0"
                    }
                }
            }
        
        elif method == "tools/list":
            # 列出所有工具
            tools = [
                {
                    "name": "search_facebook",
                    "description": "搜索Facebook内容，返回相关帖子和AI生成的总结。支持搜索关键词，返回帖子标题、链接、摘要、作者、点赞数等信息，并提供AI生成的总结。",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "搜索关键词，例如：technology news、travel tips、food recipes等"
                            },
                            "limit": {
                                "type": "integer",
                                "description": "返回结果数量，默认5，最大20",
                                "default": 5,
                                "minimum": 1,
                                "maximum": 20
                            }
                        },
                        "required": ["query"]
                    }
                }
            ]
            
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {
                    "tools": tools
                }
            }
        
        elif method == "tools/call":
            # 调用工具
            tool_name = params.get("name")
            arguments = params.get("arguments", {})
            
            if tool_name not in TOOL_MAP:
                raise ValueError(f"Tool '{tool_name}' not found")
            
            # 如果提供了 API key，临时更新配置
            if api_key and agent is not None:
                old_api_key = agent.config.litellm_api_key
                agent.config.litellm_api_key = api_key
                # 重新创建 LLM 客户端以使用新的 API key
                try:
                    from .clients.litellm_client import LiteLLMClient
                except ImportError:
                    from clients.litellm_client import LiteLLMClient
                agent.deps.llm_client = LiteLLMClient(agent.config)
            
            try:
                # 获取工具函数
                tool_func = TOOL_MAP[tool_name]
                
                # 调用工具（异步）
                result = await tool_func(**arguments)
            finally:
                # 恢复原来的 API key
                if api_key and agent is not None and 'old_api_key' in locals():
                    agent.config.litellm_api_key = old_api_key
                    try:
                        from .clients.litellm_client import LiteLLMClient
                    except ImportError:
                        from clients.litellm_client import LiteLLMClient
                    agent.deps.llm_client = LiteLLMClient(agent.config)
            
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {
                    "content": [
                        {
                            "type": "text",
                            "text": str(result)
                        }
                    ]
                }
            }
        
        elif method == "ping":
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {}
            }
        
        else:
            raise ValueError(f"Unknown method: {method}")
    
    except Exception as e:
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": -32603,
                "message": str(e)
            }
        }


def setup_logger():
    """配置日志"""
    logger.remove()
    logger.add(
        sys.stderr,
        level=config.log_level if config else "INFO",
        format="<green>{time:HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{message}</cyan>"
    )


@app.on_event("startup")
async def startup_event():
    """应用启动时初始化"""
    global config, agent
    
    try:
        # 加载配置
        config = Config.from_env()
        config.validate()
        
        # 配置日志
        setup_logger()
        
        # 创建Agent
        agent = FacebookAgent(config)
        
        logger.info("=" * 60)
        logger.info("Facebook搜索智能Agent API 启动成功")
        logger.info("=" * 60)
        logger.info(f"LiteLLM Gateway: {config.litellm_gateway_url}")
        logger.info(f"模型: {config.litellm_model}")
        logger.info(f"Facebook API Host: {config.facebook_api_host}")
        logger.info("MCP服务器已就绪，支持HTTP/SSE传输")
        
        # 初始化MCP Agent（如果还没有初始化）
        try:
            initialize_agent()
        except Exception:
            pass  # 如果已经初始化，忽略错误
        
    except Exception as e:
        logger.error(f"启动失败: {e}")
        raise


@app.get("/", tags=["健康检查"])
async def root():
    """根路径 - 服务信息"""
    return {
        "service": "Facebook搜索智能Agent API",
        "status": "running",
        "version": "1.0.0",
        "timestamp": datetime.now().isoformat(),
        "endpoints": {
            "api": {
                "search": "/search",
                "config": "/config",
                "health": "/health"
            },
            "mcp": {
                "http": "/mcp",
                "sse": "/mcp/sse"
            }
        },
        "tools": list(TOOL_MAP.keys())
    }


@app.get("/health", tags=["健康检查"])
async def health_check():
    """健康检查接口"""
    return {
        "status": "healthy",
        "service": "Facebook搜索智能Agent API + MCP Server",
        "agent_initialized": agent is not None,
        "tools_count": len(TOOL_MAP),
        "timestamp": datetime.now().isoformat()
    }


# ==================== API Key 验证 ====================

async def verify_api_key(
    api_key: Optional[str] = Header(None, alias="api-key"),
    authorization: Optional[str] = Header(None)
) -> str:
    """
    验证 API Key
    
    支持从以下位置获取 API key:
    1. api-key 请求头
    2. Authorization: Bearer <token> 请求头
    
    如果没有提供 API key，返回 401 错误
    """
    # 从 api-key 请求头获取
    if api_key:
        if not api_key.strip() or api_key.strip() == "sk":
            raise HTTPException(
                status_code=401,
                detail="无效的 API key。请提供有效的 API key。"
            )
        return api_key.strip()
    
    # 从 Authorization 请求头获取
    if authorization:
        if authorization.startswith("Bearer "):
            api_key = authorization[7:].strip()
        else:
            api_key = authorization.strip()
        
        if not api_key or api_key == "sk":
            raise HTTPException(
                status_code=401,
                detail="无效的 API key。请提供有效的 API key。"
            )
        return api_key
    
    # 如果没有提供 API key，返回错误
    raise HTTPException(
        status_code=401,
        detail="缺少 API key。请在请求头中提供 'api-key' 或 'Authorization: Bearer <token>'。"
    )


@app.post("/search", response_model=SearchResponse, tags=["搜索"])
async def search(request: SearchRequest, api_key: str = Depends(verify_api_key)):
    """
    执行Facebook搜索
    
    - **query**: 搜索关键词（必填）
    - **limit**: 返回结果数量（可选，默认5，最大20）
    
    返回：
    - 搜索结果列表
    - AI生成的总结
    - 统计信息
    """
    if agent is None:
        raise HTTPException(status_code=503, detail="服务未初始化")
    
    try:
        # 如果提供了 API key，临时更新配置
        if api_key:
            old_api_key = agent.config.litellm_api_key
            agent.config.litellm_api_key = api_key
            # 重新创建 LLM 客户端以使用新的 API key
            try:
                from .clients.litellm_client import LiteLLMClient
            except ImportError:
                from clients.litellm_client import LiteLLMClient
            agent.deps.llm_client = LiteLLMClient(agent.config)
        
        try:
            # 执行搜索
            response = await agent.search(request)
        finally:
            # 恢复原来的 API key
            if api_key and 'old_api_key' in locals():
                agent.config.litellm_api_key = old_api_key
                try:
                    from .clients.litellm_client import LiteLLMClient
                except ImportError:
                    from clients.litellm_client import LiteLLMClient
                agent.deps.llm_client = LiteLLMClient(agent.config)
        
        if not response.success:
            raise HTTPException(status_code=500, detail=response.message or "搜索失败")
        
        return response
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"搜索失败: {e}")
        raise HTTPException(status_code=500, detail=f"搜索失败: {str(e)}")


@app.get("/config", tags=["配置"])
async def get_config():
    """获取当前配置信息（隐藏敏感信息）"""
    if config is None:
        raise HTTPException(status_code=503, detail="服务未初始化")
    
    return {
        "litellm_model": config.litellm_model,
        "facebook_api_host": config.facebook_api_host,
        "max_results": config.max_results,
        "log_level": config.log_level,
        "timeout": config.timeout
    }


# ==================== MCP 协议端点 ====================

@app.post("/mcp", tags=["MCP"])
async def mcp_http_endpoint(request: Request):
    """MCP HTTP 端点 - Streamable HTTP"""
    try:
        body = await request.json()
        session_id = request.headers.get("x-mcp-session-id")
        
        # 从请求头获取 API key
        api_key = request.headers.get("api-key") or request.headers.get("api_key")
        if not api_key:
            auth_header = request.headers.get("Authorization")
            if auth_header:
                if auth_header.startswith("Bearer "):
                    api_key = auth_header[7:]
                else:
                    api_key = auth_header
        
        response = await handle_mcp_request(body, session_id, api_key=api_key)
        
        # 如果创建了新会话，返回 session ID
        if "result" in response and isinstance(response["result"], dict):
            if "sessionId" not in response["result"] and session_id:
                response["result"]["sessionId"] = session_id
        
        return JSONResponse(
            content=response,
            headers={"x-mcp-session-id": session_id or ""}
        )
    
    except Exception as e:
        return JSONResponse(
            status_code=400,
            content={
                "jsonrpc": "2.0",
                "error": {
                    "code": -32700,
                    "message": f"Parse error: {str(e)}"
                }
            }
        )


@app.get("/mcp/sse", tags=["MCP"])
async def mcp_sse_endpoint(request: Request):
    """MCP SSE 端点 - Server-Sent Events"""
    session_id = request.headers.get("x-mcp-session-id") or str(uuid.uuid4())
    
    async def event_stream() -> AsyncGenerator[str, None]:
        # 发送初始连接消息
        yield f"data: {json.dumps({'type': 'connection', 'sessionId': session_id})}\n\n"
        
        # 发送工具列表
        tools = list(TOOL_MAP.keys())
        yield f"data: {json.dumps({'type': 'tools', 'tools': tools})}\n\n"
        
        # 保持连接
        import asyncio
        while True:
            await asyncio.sleep(30)  # 心跳
            yield f"data: {json.dumps({'type': 'ping'})}\n\n"
    
    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "x-mcp-session-id": session_id
        }
    )


@app.post("/mcp/sse", tags=["MCP"])
async def mcp_sse_post(request: Request):
    """MCP SSE POST 端点 - 处理 SSE 请求"""
    try:
        body = await request.json()
        session_id = request.headers.get("x-mcp-session-id") or str(uuid.uuid4())
        
        # 从请求头获取 API key
        api_key = request.headers.get("api-key") or request.headers.get("api_key")
        if not api_key:
            auth_header = request.headers.get("Authorization")
            if auth_header:
                if auth_header.startswith("Bearer "):
                    api_key = auth_header[7:]
                else:
                    api_key = auth_header
        
        async def response_stream() -> AsyncGenerator[str, None]:
            response = await handle_mcp_request(body, session_id, api_key=api_key)
            yield f"data: {json.dumps(response)}\n\n"
        
        return StreamingResponse(
            response_stream(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "x-mcp-session-id": session_id
            }
        )
    
    except Exception as e:
        return JSONResponse(
            status_code=400,
            content={
                "jsonrpc": "2.0",
                "error": {
                    "code": -32700,
                    "message": f"Parse error: {str(e)}"
                }
            }
        )


if __name__ == "__main__":
    import uvicorn
    
    # 运行服务
    uvicorn.run(
        "facebook_agent.api:app",
        host="0.0.0.0",
        port=8000,
        reload=True,
        log_level="info"
    )