"""
MCP Client Manager

MCP 서버들을 등록하고 관리하며, LangChain Tool로 변환하는 모듈
"""

import os
import json
import asyncio
import subprocess
from typing import Dict, List, Any, Optional, Callable, Tuple
from dataclasses import dataclass, field
from contextlib import asynccontextmanager

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from langchain_core.tools import StructuredTool
from pydantic import BaseModel, Field, create_model


@dataclass
class MCPServerConfig:
    """MCP 서버 설정"""
    name: str
    command: str
    args: List[str] = field(default_factory=list)
    env: Dict[str, str] = field(default_factory=dict)
    enabled: bool = True


@dataclass
class MCPToolInfo:
    """MCP 도구 정보"""
    server_name: str
    name: str
    description: str
    input_schema: Dict[str, Any]


@dataclass
class MCPConnection:
    """MCP 연결 정보"""
    session: ClientSession
    cleanup: Callable
    task: asyncio.Task


class MCPClientManager:
    """
    MCP 클라이언트 매니저
    
    여러 MCP 서버를 등록하고 관리하며, 각 서버의 도구들을 LangChain Tool로 변환합니다.
    Cursor IDE의 MCP 설정 방식과 유사하게 JSON 설정 파일로 서버를 관리합니다.
    """
    
    def __init__(self, config_path: Optional[str] = None):
        """
        Args:
            config_path: MCP 서버 설정 파일 경로 (JSON)
        """
        self.config_path = config_path
        self.servers: Dict[str, MCPServerConfig] = {}
        self.sessions: Dict[str, ClientSession] = {}
        self.tools: Dict[str, MCPToolInfo] = {}
        self._connections: Dict[str, MCPConnection] = {}
        self._background_tasks: List[asyncio.Task] = []
        
        if config_path and os.path.exists(config_path):
            self.load_config(config_path)
    
    def load_config(self, config_path: str) -> None:
        """JSON 설정 파일에서 MCP 서버 설정 로드"""
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        
        mcp_servers = config.get("mcpServers", {})
        for name, server_config in mcp_servers.items():
            self.register_server(
                name=name,
                command=server_config.get("command", ""),
                args=server_config.get("args", []),
                env=server_config.get("env", {}),
                enabled=server_config.get("enabled", True)
            )
    
    def register_server(
        self,
        name: str,
        command: str,
        args: List[str] = None,
        env: Dict[str, str] = None,
        enabled: bool = True
    ) -> None:
        """
        MCP 서버 등록
        
        Args:
            name: 서버 이름 (고유 식별자)
            command: 실행 명령어 (예: "uvx", "npx", "python")
            args: 명령어 인자
            env: 환경 변수
            enabled: 활성화 여부
        """
        self.servers[name] = MCPServerConfig(
            name=name,
            command=command,
            args=args or [],
            env=env or {},
            enabled=enabled
        )
    
    async def _run_server_connection(
        self, 
        name: str, 
        server_params: StdioServerParameters,
        ready_event: asyncio.Event
    ) -> None:
        """백그라운드에서 MCP 서버 연결 유지"""
        try:
            async with stdio_client(server_params) as (read, write):
                async with ClientSession(read, write) as session:
                    await session.initialize()
                    self.sessions[name] = session
                    print(f"[MCP] 서버 '{name}' 연결 성공")
                    ready_event.set()
                    
                    # 연결 유지 (무한 대기)
                    while True:
                        await asyncio.sleep(3600)  # 1시간마다 체크
                        
        except asyncio.CancelledError:
            print(f"[MCP] 서버 '{name}' 연결 종료")
            raise
        except Exception as e:
            print(f"[MCP] 서버 '{name}' 연결 오류: {e}")
            ready_event.set()  # 오류 시에도 이벤트 설정하여 대기 해제
    
    async def connect_server(self, name: str, timeout: float = 30.0) -> Optional[ClientSession]:
        """
        특정 MCP 서버에 연결
        
        Args:
            name: 서버 이름
            timeout: 연결 대기 시간 (초)
            
        Returns:
            ClientSession 또는 None (연결 실패 시)
        """
        if name not in self.servers:
            print(f"[MCP] 서버 '{name}'을 찾을 수 없습니다.")
            return None
        
        server = self.servers[name]
        if not server.enabled:
            print(f"[MCP] 서버 '{name}'이 비활성화되어 있습니다.")
            return None
        
        if name in self.sessions:
            return self.sessions[name]
        
        try:
            # 환경 변수 병합
            env = {**os.environ, **server.env}
            
            # StdioServerParameters 생성
            server_params = StdioServerParameters(
                command=server.command,
                args=server.args,
                env=env
            )
            
            # 연결 준비 이벤트
            ready_event = asyncio.Event()
            
            # 백그라운드 태스크로 연결 시작
            task = asyncio.create_task(
                self._run_server_connection(name, server_params, ready_event)
            )
            self._background_tasks.append(task)
            
            # 연결 대기
            try:
                await asyncio.wait_for(ready_event.wait(), timeout=timeout)
            except asyncio.TimeoutError:
                print(f"[MCP] 서버 '{name}' 연결 시간 초과")
                task.cancel()
                return None
            
            return self.sessions.get(name)
            
        except Exception as e:
            print(f"[MCP] 서버 '{name}' 연결 실패: {e}")
            return None
    
    async def connect_all(self) -> Dict[str, ClientSession]:
        """모든 활성화된 서버에 연결"""
        for name, server in self.servers.items():
            if server.enabled and name not in self.sessions:
                await self.connect_server(name)
        return self.sessions
    
    async def disconnect_all(self) -> None:
        """모든 서버 연결 해제"""
        # 백그라운드 태스크 취소
        for task in self._background_tasks:
            if not task.done():
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
        
        self.sessions.clear()
        self._background_tasks.clear()
        print("[MCP] 모든 연결 해제됨")
    
    async def list_tools(self, server_name: Optional[str] = None) -> List[MCPToolInfo]:
        """
        MCP 서버의 도구 목록 조회
        
        Args:
            server_name: 특정 서버 이름 (None이면 모든 서버)
            
        Returns:
            MCPToolInfo 목록
        """
        tools = []
        
        sessions_to_query = {}
        if server_name:
            if server_name in self.sessions:
                sessions_to_query[server_name] = self.sessions[server_name]
        else:
            sessions_to_query = self.sessions
        
        for name, session in sessions_to_query.items():
            try:
                result = await session.list_tools()
                for tool in result.tools:
                    tool_info = MCPToolInfo(
                        server_name=name,
                        name=tool.name,
                        description=tool.description or "",
                        input_schema=tool.inputSchema or {}
                    )
                    tools.append(tool_info)
                    self.tools[f"{name}.{tool.name}"] = tool_info
            except Exception as e:
                print(f"[MCP] 서버 '{name}'의 도구 목록 조회 실패: {e}")
        
        return tools
    
    async def call_tool(
        self,
        server_name: str,
        tool_name: str,
        arguments: Dict[str, Any]
    ) -> Any:
        """
        MCP 도구 호출
        
        Args:
            server_name: 서버 이름
            tool_name: 도구 이름
            arguments: 인자
            
        Returns:
            도구 실행 결과
        """
        print(f"[MCP] 도구 호출 시작: {server_name}.{tool_name}")
        # print(f"[MCP] 인자: {arguments}")
        
        if server_name not in self.sessions:
            print(f"[MCP] 오류: 서버 '{server_name}'에 연결되어 있지 않습니다.")
            raise ValueError(f"서버 '{server_name}'에 연결되어 있지 않습니다.")
        
        try:
            session = self.sessions[server_name]
            result = await session.call_tool(tool_name, arguments)
            
            # 결과 처리
            if result.content:
                # TextContent 처리
                contents = []
                for content in result.content:
                    if hasattr(content, 'text'):
                        contents.append(content.text)
                    elif hasattr(content, 'data'):
                        contents.append(content.data)
                result_str = "\n".join(contents) if contents else ""
                print(f"[MCP] 도구 호출 성공: {tool_name}")
                # print(f"[MCP] 결과 (처음 500자): {result_str[:500] if result_str else 'empty'}")
                return result_str
            
            print(f"[MCP] 도구 호출 완료 (content 없음): {tool_name}")
            return result
        except Exception as e:
            print(f"[MCP] 도구 호출 실패: {tool_name}")
            print(f"[MCP] 오류: {e}")
            raise
    
    def _create_pydantic_model_from_schema(
        self,
        name: str,
        schema: Dict[str, Any]
    ) -> type[BaseModel]:
        """JSON Schema를 Pydantic 모델로 변환"""
        properties = schema.get("properties", {})
        required = set(schema.get("required", []))
        
        fields = {}
        for prop_name, prop_schema in properties.items():
            prop_type = prop_schema.get("type", "string")
            description = prop_schema.get("description", "")
            default = ... if prop_name in required else None
            
            # 타입 매핑
            type_mapping = {
                "string": str,
                "integer": int,
                "number": float,
                "boolean": bool,
                "array": list,
                "object": dict
            }
            python_type = type_mapping.get(prop_type, str)
            
            if default is None:
                python_type = Optional[python_type]
            
            fields[prop_name] = (python_type, Field(default=default, description=description))
        
        return create_model(name, **fields)
    
    async def get_langchain_tools(
        self,
        server_name: Optional[str] = None
    ) -> List[StructuredTool]:
        """
        MCP 도구들을 LangChain StructuredTool로 변환
        
        Args:
            server_name: 특정 서버 이름 (None이면 모든 서버)
            
        Returns:
            LangChain StructuredTool 목록
        """
        mcp_tools = await self.list_tools(server_name)
        langchain_tools = []
        
        for tool_info in mcp_tools:
            # Pydantic 모델 생성
            input_schema = tool_info.input_schema
            model_name = f"{tool_info.server_name}_{tool_info.name}_Input".replace("-", "_")
            args_schema = self._create_pydantic_model_from_schema(model_name, input_schema)
            
            # 클로저를 위한 변수 캡처
            _server_name = tool_info.server_name
            _tool_name = tool_info.name
            _self = self
            
            # 비동기 함수 생성
            async def tool_func(_server=_server_name, _tool=_tool_name, **kwargs) -> str:
                try:
                    result = await _self.call_tool(_server, _tool, kwargs)
                    return str(result) if result else "실행 완료"
                except Exception as e:
                    return f"오류: {str(e)}"
            
            # StructuredTool 생성
            langchain_tool = StructuredTool(
                name=f"{tool_info.server_name}__{tool_info.name}",
                description=f"[{tool_info.server_name}] {tool_info.description}",
                func=lambda **kwargs: asyncio.run(tool_func(**kwargs)),
                coroutine=tool_func,
                args_schema=args_schema
            )
            
            langchain_tools.append(langchain_tool)
        
        return langchain_tools


@asynccontextmanager
async def create_mcp_client(config_path: Optional[str] = None):
    """
    MCP 클라이언트 컨텍스트 매니저
    
    사용 예:
        async with create_mcp_client("config.json") as client:
            tools = await client.get_langchain_tools()
    """
    client = MCPClientManager(config_path)
    try:
        await client.connect_all()
        yield client
    finally:
        await client.disconnect_all()
