"""
Azure Blob Storage AI Agent - MCP (Model Context Protocol) 版本
使用 MCP 协议实现智能文件操作功能
"""
import os
import logging
import json
from typing import Optional, Dict, Any, List
from datetime import datetime
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from azure.storage.blob import BlobServiceClient, ContainerClient
import uvicorn
import asyncio

# 配置日志
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", "8080"))
POD_NAME = os.getenv("POD_NAME", "azure-blob-agent-mcp")
TEMPLATE_TYPE = os.getenv("TEMPLATE_TYPE", "azure_blob_agent_mcp")
AGENT_FRAMEWORK = os.getenv("AGENT_FRAMEWORK", "mcp")

# 工具配置 (从环境变量传入的 JSON)
TOOLS_CONFIG = json.loads(os.getenv("TOOLS_CONFIG", "{}"))
TOOL_ENDPOINT = os.getenv("TOOL_ENDPOINT", "")
TOOL_API_KEY = os.getenv("TOOL_API_KEY", "")

# 模型配置
MODEL_PROVIDER = os.getenv("MODEL_PROVIDER", "openai")
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4")
MODEL_API_KEY = os.getenv("MODEL_API_KEY", "")
MODEL_ENDPOINT = os.getenv("MODEL_ENDPOINT", "https://api.openai.com/v1")

# 存储配置
AZURE_STORAGE_CONNECTION_STRING = os.getenv("AZURE_STORAGE_CONNECTION_STRING", "")
STORAGE_ACCOUNT_NAME = os.getenv("STORAGE_ACCOUNT_NAME", "")

# 用户标识
USER_ID = os.getenv("USER_ID", "")
TENANT_ID = os.getenv("TENANT_ID", "")
NAMESPACE = os.getenv("NAMESPACE", "ai-agents")

# 全局存储客户端
blob_service_client: Optional[BlobServiceClient] = None
connection_string: Optional[str] = None

# MCP 工具注册表
mcp_tools: Dict[str, Any] = {}

# FastAPI应用
app = FastAPI(
    title="Azure Blob Storage AI Agent (MCP)",
    description="基于 MCP 协议的智能 Azure Blob 存储管理代理",
    version="1.0.0"
)


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

class ConnectRequest(BaseModel):
    """连接请求"""
    connection_string: str = Field(..., description="Azure Storage连接字符串")


class MCPToolRequest(BaseModel):
    """MCP 工具调用请求"""
    tool_name: str = Field(..., description="工具名称")
    parameters: Dict[str, Any] = Field(default_factory=dict, description="工具参数")


class MCPQueryRequest(BaseModel):
    """MCP 查询请求"""
    query: str = Field(..., description="自然语言查询或操作指令")
    container_name: Optional[str] = Field(None, description="指定容器名称")
    context: Optional[Dict] = Field(default_factory=dict, description="上下文信息")


class HealthResponse(BaseModel):
    """健康检查响应"""
    status: str
    connected: bool
    framework: str
    user_id: Optional[str] = None
    namespace: Optional[str] = None
    connection_info: Optional[Dict] = None


# ==================== MCP 工具定义 ====================

class MCPTool:
    """MCP 工具基类"""
    
    def __init__(self, name: str, description: str, parameters_schema: Dict):
        self.name = name
        self.description = description
        self.parameters_schema = parameters_schema
    
    async def execute(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """执行工具"""
        raise NotImplementedError
    
    def to_mcp_spec(self) -> Dict:
        """转换为 MCP 工具规范"""
        return {
            "name": self.name,
            "description": self.description,
            "inputSchema": {
                "type": "object",
                "properties": self.parameters_schema,
                "required": list(self.parameters_schema.keys())
            }
        }


class ListContainersTool(MCPTool):
    """列出所有容器工具"""
    
    def __init__(self):
        super().__init__(
            name="list_containers",
            description="列出 Azure Blob Storage 中的所有容器",
            parameters_schema={}
        )
    
    async def execute(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        global blob_service_client
        
        if not blob_service_client:
            return {"error": "未连接到 Azure Blob Storage"}
        
        try:
            containers = blob_service_client.list_containers()
            container_list = []
            for container in containers:
                container_list.append({
                    "name": container.name,
                    "last_modified": str(container.last_modified)
                })
            
            return {
                "success": True,
                "containers": container_list,
                "count": len(container_list)
            }
        except Exception as e:
            logger.error(f"列出容器失败: {str(e)}")
            return {"error": str(e)}


class ListBlobsTool(MCPTool):
    """列出容器中的 blob 工具"""
    
    def __init__(self):
        super().__init__(
            name="list_blobs",
            description="列出指定容器中的所有文件",
            parameters_schema={
                "container_name": {
                    "type": "string",
                    "description": "容器名称"
                }
            }
        )
    
    async def execute(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        global blob_service_client
        
        if not blob_service_client:
            return {"error": "未连接到 Azure Blob Storage"}
        
        container_name = parameters.get("container_name")
        if not container_name:
            return {"error": "缺少参数: container_name"}
        
        try:
            container_client = blob_service_client.get_container_client(container_name)
            blobs = container_client.list_blobs()
            
            blob_list = []
            total_size = 0
            for blob in blobs:
                blob_info = {
                    "name": blob.name,
                    "size": blob.size,
                    "size_mb": round(blob.size / (1024 * 1024), 2),
                    "content_type": blob.content_settings.content_type if blob.content_settings else "unknown",
                    "last_modified": str(blob.last_modified)
                }
                blob_list.append(blob_info)
                total_size += blob.size
            
            return {
                "success": True,
                "container": container_name,
                "blobs": blob_list,
                "count": len(blob_list),
                "total_size_mb": round(total_size / (1024 * 1024), 2)
            }
        except Exception as e:
            logger.error(f"列出 blob 失败: {str(e)}")
            return {"error": str(e)}


class GetBlobInfoTool(MCPTool):
    """获取 blob 信息工具"""
    
    def __init__(self):
        super().__init__(
            name="get_blob_info",
            description="获取特定文件的详细信息",
            parameters_schema={
                "container_name": {
                    "type": "string",
                    "description": "容器名称"
                },
                "blob_name": {
                    "type": "string",
                    "description": "文件名称"
                }
            }
        )
    
    async def execute(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        global blob_service_client
        
        if not blob_service_client:
            return {"error": "未连接到 Azure Blob Storage"}
        
        container_name = parameters.get("container_name")
        blob_name = parameters.get("blob_name")
        
        if not container_name or not blob_name:
            return {"error": "缺少参数: container_name 或 blob_name"}
        
        try:
            blob_client = blob_service_client.get_blob_client(container_name, blob_name)
            properties = blob_client.get_blob_properties()
            
            return {
                "success": True,
                "blob_name": blob_name,
                "container": container_name,
                "size": properties.size,
                "size_mb": round(properties.size / (1024 * 1024), 2),
                "content_type": properties.content_settings.content_type if properties.content_settings else "unknown",
                "creation_time": str(properties.creation_time),
                "last_modified": str(properties.last_modified),
                "etag": properties.etag,
                "metadata": properties.metadata if properties.metadata else {}
            }
        except Exception as e:
            logger.error(f"获取 blob 信息失败: {str(e)}")
            return {"error": str(e)}


class SearchBlobsTool(MCPTool):
    """搜索 blob 工具"""
    
    def __init__(self):
        super().__init__(
            name="search_blobs",
            description="在容器中搜索包含关键字的文件",
            parameters_schema={
                "container_name": {
                    "type": "string",
                    "description": "容器名称"
                },
                "keyword": {
                    "type": "string",
                    "description": "搜索关键字"
                }
            }
        )
    
    async def execute(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        global blob_service_client
        
        if not blob_service_client:
            return {"error": "未连接到 Azure Blob Storage"}
        
        container_name = parameters.get("container_name")
        keyword = parameters.get("keyword")
        
        if not container_name or not keyword:
            return {"error": "缺少参数: container_name 或 keyword"}
        
        try:
            container_client = blob_service_client.get_container_client(container_name)
            blobs = container_client.list_blobs()
            
            matched_blobs = []
            for blob in blobs:
                if keyword.lower() in blob.name.lower():
                    matched_blobs.append({
                        "name": blob.name,
                        "size": blob.size,
                        "size_kb": round(blob.size / 1024, 2),
                        "last_modified": str(blob.last_modified)
                    })
            
            return {
                "success": True,
                "container": container_name,
                "keyword": keyword,
                "results": matched_blobs,
                "count": len(matched_blobs)
            }
        except Exception as e:
            logger.error(f"搜索 blob 失败: {str(e)}")
            return {"error": str(e)}


class GetStorageStatsTool(MCPTool):
    """获取存储统计工具"""
    
    def __init__(self):
        super().__init__(
            name="get_storage_stats",
            description="获取存储的统计信息，包括容器数量、文件数量、总大小等",
            parameters_schema={}
        )
    
    async def execute(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        global blob_service_client
        
        if not blob_service_client:
            return {"error": "未连接到 Azure Blob Storage"}
        
        try:
            containers = list(blob_service_client.list_containers())
            total_containers = len(containers)
            total_blobs = 0
            total_size = 0
            
            container_stats = []
            for container in containers:
                container_client = blob_service_client.get_container_client(container.name)
                blobs = list(container_client.list_blobs())
                blob_count = len(blobs)
                container_size = sum(blob.size for blob in blobs)
                
                total_blobs += blob_count
                total_size += container_size
                
                container_stats.append({
                    "name": container.name,
                    "blobs": blob_count,
                    "size_mb": round(container_size / (1024 * 1024), 2)
                })
            
            return {
                "success": True,
                "total_containers": total_containers,
                "total_blobs": total_blobs,
                "total_size_mb": round(total_size / (1024 * 1024), 2),
                "container_stats": container_stats
            }
        except Exception as e:
            logger.error(f"获取统计信息失败: {str(e)}")
            return {"error": str(e)}


# ==================== MCP 工具注册 ====================

def register_tools():
    """注册所有 MCP 工具"""
    global mcp_tools
    
    tools = [
        ListContainersTool(),
        ListBlobsTool(),
        GetBlobInfoTool(),
        SearchBlobsTool(),
        GetStorageStatsTool()
    ]
    
    for tool in tools:
        mcp_tools[tool.name] = tool
    
    logger.info(f"✅ 注册了 {len(mcp_tools)} 个 MCP 工具")


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

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """健康检查"""
    global blob_service_client, connection_string
    
    connected = blob_service_client is not None
    
    connection_info = None
    if connected:
        try:
            account_info = blob_service_client.get_account_information()
            connection_info = {
                "account_kind": account_info.get('account_kind', 'unknown'),
                "sku_name": account_info.get('sku_name', 'unknown'),
                "connected_at": str(datetime.now())
            }
        except Exception as e:
            logger.error(f"获取账户信息失败: {str(e)}")
    
    return HealthResponse(
        status="healthy" if connected else "not_connected",
        connected=connected,
        framework=AGENT_FRAMEWORK,
        user_id=USER_ID,
        namespace=NAMESPACE,
        connection_info=connection_info
    )


@app.post("/connect")
async def connect_to_storage(request: ConnectRequest):
    """连接到 Azure Blob Storage"""
    global blob_service_client, connection_string
    
    try:
        blob_service_client = BlobServiceClient.from_connection_string(
            request.connection_string
        )
        
        account_info = blob_service_client.get_account_information()
        connection_string = request.connection_string
        
        logger.info(f"✅ 成功连接到 Azure Blob Storage (User: {USER_ID})")
        
        return {
            "status": "connected",
            "message": "成功连接到 Azure Blob Storage",
            "framework": AGENT_FRAMEWORK,
            "user_id": USER_ID,
            "account_info": {
                "account_kind": account_info.get('account_kind'),
                "sku_name": account_info.get('sku_name')
            }
        }
    except Exception as e:
        logger.error(f"❌ 连接失败: {str(e)}")
        blob_service_client = None
        connection_string = None
        raise HTTPException(status_code=400, detail=f"连接失败: {str(e)}")


@app.get("/mcp/tools")
async def list_mcp_tools():
    """列出所有可用的 MCP 工具"""
    if not blob_service_client:
        raise HTTPException(
            status_code=400,
            detail="未连接到 Azure Blob Storage，请先调用 /connect"
        )
    
    tools_spec = [tool.to_mcp_spec() for tool in mcp_tools.values()]
    
    return {
        "tools": tools_spec,
        "count": len(tools_spec),
        "framework": AGENT_FRAMEWORK
    }


@app.post("/mcp/call")
async def call_mcp_tool(request: MCPToolRequest):
    """调用 MCP 工具"""
    if not blob_service_client:
        raise HTTPException(
            status_code=400,
            detail="未连接到 Azure Blob Storage，请先调用 /connect"
        )
    
    tool_name = request.tool_name
    if tool_name not in mcp_tools:
        raise HTTPException(
            status_code=404,
            detail=f"工具 '{tool_name}' 不存在"
        )
    
    try:
        tool = mcp_tools[tool_name]
        result = await tool.execute(request.parameters)
        
        return {
            "tool": tool_name,
            "result": result,
            "timestamp": str(datetime.now())
        }
    except Exception as e:
        logger.error(f"工具调用失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"工具调用失败: {str(e)}")


@app.post("/query")
async def query_storage(request: MCPQueryRequest):
    """使用自然语言查询存储 (简化版 - 实际应集成 LLM)"""
    if not blob_service_client:
        raise HTTPException(
            status_code=400,
            detail="未连接到 Azure Blob Storage，请先调用 /connect"
        )
    
    try:
        query = request.query.lower()
        result = None
        
        # 简单的规则匹配 (实际应使用 LLM 进行意图识别)
        if "容器" in query and ("列出" in query or "显示" in query or "有哪些" in query):
            tool = mcp_tools["list_containers"]
            result = await tool.execute({})
        elif "统计" in query or "有多少" in query or "占用" in query:
            tool = mcp_tools["get_storage_stats"]
            result = await tool.execute({})
        elif request.container_name:
            if "文件" in query or "blob" in query.lower():
                tool = mcp_tools["list_blobs"]
                result = await tool.execute({"container_name": request.container_name})
        
        if result:
            return {
                "status": "success",
                "query": request.query,
                "result": result,
                "framework": AGENT_FRAMEWORK
            }
        else:
            return {
                "status": "info",
                "query": request.query,
                "message": "未能匹配到合适的工具，请使用 /mcp/tools 查看可用工具",
                "available_tools": list(mcp_tools.keys())
            }
    except Exception as e:
        logger.error(f"查询执行失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")


@app.get("/")
async def root():
    """根端点"""
    return {
        "service": "Azure Blob Storage AI Agent",
        "version": "1.0.0",
        "framework": AGENT_FRAMEWORK,
        "pod_name": POD_NAME,
        "template": TEMPLATE_TYPE,
        "user_id": USER_ID,
        "namespace": NAMESPACE,
        "connected": blob_service_client is not None,
        "tools_count": len(mcp_tools),
        "endpoints": {
            "health": "/health",
            "connect": "POST /connect",
            "list_tools": "GET /mcp/tools",
            "call_tool": "POST /mcp/call",
            "query": "POST /query"
        }
    }


# ==================== 主函数 ====================

def init_storage_connection():
    """启动时初始化存储连接"""
    global blob_service_client, connection_string
    
    if AZURE_STORAGE_CONNECTION_STRING:
        try:
            logger.info("检测到环境变量中的连接字符串，尝试连接...")
            blob_service_client = BlobServiceClient.from_connection_string(
                AZURE_STORAGE_CONNECTION_STRING
            )
            
            account_info = blob_service_client.get_account_information()
            connection_string = AZURE_STORAGE_CONNECTION_STRING
            
            logger.info(f"✅ 成功连接到 Azure Blob Storage")
            logger.info(f"   - Account Kind: {account_info.get('account_kind')}")
            logger.info(f"   - SKU: {account_info.get('sku_name')}")
        except Exception as e:
            logger.error(f"❌ 启动时连接失败: {str(e)}")
            logger.info("💡 提示: 可以稍后通过 /connect API 手动连接")
            blob_service_client = None
            connection_string = None
    else:
        logger.info("💡 未设置 AZURE_STORAGE_CONNECTION_STRING，需通过 /connect API 手动连接")


def main():
    """启动服务"""
    logger.info(f"🚀 启动 Azure Blob Storage AI Agent (MCP)")
    logger.info(f"   - Framework: {AGENT_FRAMEWORK}")
    logger.info(f"   - Pod名称: {POD_NAME}")
    logger.info(f"   - 模板类型: {TEMPLATE_TYPE}")
    logger.info(f"   - User ID: {USER_ID}")
    logger.info(f"   - Namespace: {NAMESPACE}")
    logger.info(f"   - 模型: {MODEL_NAME} @ {MODEL_PROVIDER}")
    logger.info(f"   - 服务地址: http://{SERVICE_HOST}:{SERVICE_PORT}")
    
    # 注册 MCP 工具
    register_tools()
    
    # 初始化存储连接
    init_storage_connection()
    
    uvicorn.run(
        app,
        host=SERVICE_HOST,
        port=SERVICE_PORT,
        log_level="info"
    )


if __name__ == "__main__":
    main()