import asyncio
import gradio as gr
from mcp import ClientSession
from mcp.client.sse import sse_client
from langchain_mcp_adapters.tools import load_mcp_tools
from langgraph.prebuilt import create_react_agent
import traceback
from typing import List, Tuple, Optional
from pprint import pprint

class MCPChatbot:
    def __init__(self):
        self.session: Optional[ClientSession] = None
        self.agent = None
        self.tools = None
        self.sse_url = "http://127.0.0.1:7860/gradio_api/mcp/sse"
        self.is_initialized = False

    async def initialize(self):
        """MCP 연결 및 에이전트 초기화"""
        try:
            # SSE 클라이언트 연결
            self.sse_client_context = sse_client(self.sse_url)
            read, write = await self.sse_client_context.__aenter__()

            # 클라이언트 세션 생성
            self.session_context = ClientSession(read, write)
            self.session = await self.session_context.__aenter__()

            # 세션 초기화
            await self.session.initialize()

            # 도구 로드
            self.tools = await load_mcp_tools(self.session)
            tool_names = [tool.name for tool in self.tools]
            print(f"로드된 도구들: {tool_names}")

            # 에이전트 생성
            self.agent = create_react_agent("openai:gpt-4o", self.tools)
            self.is_initialized = True

            return f"✅ 초기화 완료! 사용 가능한 도구: {', '.join(tool_names)}"

        except Exception as e:
            error_msg = f"❌ 초기화 실패: {str(e)}\n{traceback.format_exc()}"
            print(error_msg)
            return error_msg

    async def cleanup(self):
        """리소스 정리"""
        try:
            if hasattr(self, 'session_context') and self.session_context:
                await self.session_context.__aexit__(None, None, None)
            if hasattr(self, 'sse_client_context') and self.sse_client_context:
                await self.sse_client_context.__aexit__(None, None, None)
            self.is_initialized = False
        except Exception as e:
            print(f"정리 중 오류: {e}")

    async def chat(self, message: str) -> str:
        """메시지 처리 및 응답 생성"""
        if not self.is_initialized:
            return "❌ 먼저 '초기화' 버튼을 클릭해주세요!"

        try:
            # 에이전트에게 메시지 전달
            response = await self.agent.ainvoke({
                "messages": [{"role": "user", "content": message}]
            })
            pprint(response)

            # 응답에서 마지막 메시지 추출
            if "messages" in response and response["messages"]:
                last_message = response["messages"][-1]
                if hasattr(last_message, 'content'):
                    return last_message.content
                elif isinstance(last_message, dict) and 'content' in last_message:
                    return last_message['content']
                else:
                    return str(last_message)
            else:
                return "응답을 받지 못했습니다."

        except Exception as e:
            error_msg = f"❌ 오류 발생: {str(e)}\n{traceback.format_exc()}"
            print(error_msg)
            return error_msg

# 전역 챗봇 인스턴스
chatbot = MCPChatbot()

# 비동기 함수들을 동기 함수로 래핑
def initialize_chatbot():
    """챗봇 초기화"""
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

    return loop.run_until_complete(chatbot.initialize())

def process_message(message: str, history: List[Tuple[str, str]]) -> Tuple[List[Tuple[str, str]], str]:
    """메시지 처리"""
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

    response = loop.run_until_complete(chatbot.chat(message))

    # 히스토리에 대화 추가
    history.append((message, response))

    return history, ""

def cleanup_chatbot():
    """챗봇 정리"""
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

    loop.run_until_complete(chatbot.cleanup())
    return "정리 완료"

# Gradio 인터페이스 생성
def create_interface():
    with gr.Blocks(title="MCP 챗봇", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🤖 MCP 도구 연동 챗봇")
        gr.Markdown("MCP(Model Context Protocol) 도구들을 사용할 수 있는 AI 챗봇입니다.")

        with gr.Row():
            with gr.Column(scale=3):
                # 채팅 인터페이스
                chatbot_interface = gr.Chatbot(
                    label="대화",
                    height=500,
                    show_copy_button=True
                )

                with gr.Row():
                    msg_input = gr.Textbox(
                        placeholder="메시지를 입력하세요... (예: (3 + 5) x 12는 뭐야?)",
                        label="메시지",
                        scale=4
                    )
                    send_btn = gr.Button("전송", variant="primary", scale=1)

                # 예시 메시지들
                gr.Examples(
                    examples=[
                        "(3 + 5) x 12는 뭐야?",
                        "1234의 소인수분해를 해줘",
                        "오늘 날씨는 어때?",
                        "안녕하세요!"
                    ],
                    inputs=msg_input
                )

            with gr.Column(scale=1):
                # 제어 패널
                gr.Markdown("### 🛠️ 제어 패널")

                init_btn = gr.Button("🚀 초기화", variant="secondary")
                init_status = gr.Textbox(
                    label="초기화 상태",
                    value="초기화가 필요합니다",
                    interactive=False
                )

                cleanup_btn = gr.Button("🧹 정리", variant="secondary")
                cleanup_status = gr.Textbox(
                    label="정리 상태",
                    interactive=False
                )

                # 도움말
                gr.Markdown("""
                ### 📝 사용법
                1. **초기화** 버튼을 클릭하여 MCP 서버에 연결
                2. 초기화가 완료되면 메시지를 입력하여 대화 시작
                3. 수학 계산, 날씨 정보 등 다양한 질문 가능
                4. 사용 완료 후 **정리** 버튼으로 연결 종료

                ### 🔧 사용 가능한 기능
                - 수학 계산 및 소인수분해
                - 날씨 정보 조회
                - 이미지 처리 (방향 확인, 세피아 필터)
                - 기타 MCP 도구들
                """)

        # 이벤트 핸들러 설정
        init_btn.click(
            initialize_chatbot,
            outputs=init_status
        )

        cleanup_btn.click(
            cleanup_chatbot,
            outputs=cleanup_status
        )

        # 메시지 전송 처리
        def handle_submit(message, history):
            if not message.strip():
                return history, ""
            return process_message(message, history)

        send_btn.click(
            handle_submit,
            inputs=[msg_input, chatbot_interface],
            outputs=[chatbot_interface, msg_input]
        )

        msg_input.submit(
            handle_submit,
            inputs=[msg_input, chatbot_interface],
            outputs=[chatbot_interface, msg_input]
        )

    return demo

# 애플리케이션 실행
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        share=False,
    )