"""
Work Assistant Agent - LangChain ReAct Agent with MCP Support

LangGraph를 사용한 ReAct 패턴 에이전트 구현
"""

import os
import json
from typing import List, Dict, Any, Optional, AsyncGenerator, Literal, TYPE_CHECKING
from dataclasses import dataclass, field

from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage, ToolMessage
from langchain_core.tools import BaseTool
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langgraph.prebuilt import create_react_agent

from .mcp_client import MCPClientManager, create_mcp_client


@dataclass
class AgentConfig:
    """에이전트 설정"""
    model_provider: Literal["openai", "anthropic"] = "openai"
    model_name: str = "gpt-4o"
    temperature: float = 0.0
    max_tokens: int = 4096
    system_prompt: str = ""
    mcp_config_path: Optional[str] = None
    streaming: bool = True


@dataclass
class AgentContext:
    """에이전트 컨텍스트 (세션별 상태)"""
    tenant_id: str = ""
    user_uid: str = ""
    user_email: str = ""
    user_name: str = ""
    conversation_history: List[BaseMessage] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)


class WorkAssistantAgent:
    """
    Work Assistant Agent
    
    MCP 서버의 도구들을 활용하는 LangChain ReAct 에이전트.
    Cursor IDE의 agent처럼 MCP를 등록하여 사용할 수 있습니다.
    
    사용 예:
        agent = WorkAssistantAgent(config)
        await agent.initialize()
        
        context = AgentContext(tenant_id="uengine", user_uid="xxx")
        response = await agent.run("휴가신청 해줘", context)
    """
    
    DEFAULT_SYSTEM_PROMPT = """당신은 업무를 도와주는 AI 어시스턴트입니다.

사용자의 요청을 분석하고, 적절한 도구를 사용하여 작업을 수행합니다.

【도구 사용 가이드라인】

1. 프로세스 생성 요청 시 ("프로세스 만들어줘", "워크플로우 생성해줘" 등)
   → generate_process 도구만 호출하고 즉시 종료

2. 프로세스 실행 요청 시
   → get_process_list → get_process_detail → get_form_fields → execute_process

3. 프로세스 실행 결과 조회 시
   → get_process_list → get_instance_list → get_todolist

4. 조직도 조회 시
   → get_organization 직접 호출

【중요: 응답 형식】

모든 작업 완료 후, 반드시 아래 JSON 형식으로 최종 응답을 작성하세요:

```json
{
  "action": "작업유형",
  "status": "success 또는 error",
  "message": "사용자에게 보여줄 친절한 메시지",
  "data": { 실행 결과 데이터 (선택) },
  "next_actions": ["후속 가능한 작업들"] 
}
```

action 유형:
- "process_created": 프로세스 생성 요청
- "process_executed": 프로세스 실행 완료
- "query_result": 조회 결과 반환
- "organization_info": 조직도 정보
- "general_response": 일반 응답
- "error": 오류 발생

예시:
- 프로세스 목록 조회: {"action": "query_result", "status": "success", "message": "총 5개의 프로세스가 있습니다.", "data": {"processes": [...]}}
- 실행 완료: {"action": "process_executed", "status": "success", "message": "휴가신청이 완료되었습니다.", "data": {"instance_id": "xxx"}}

도구 호출 시 tenant_id, user_uid, user_email 등 컨텍스트 정보를 반드시 활용하세요.
메시지는 한국어로 친절하게 작성하세요."""
    
    def __init__(self, config: Optional[AgentConfig] = None):
        """
        Args:
            config: 에이전트 설정
        """
        self.config = config or AgentConfig()
        self.llm: Optional[BaseChatModel] = None
        self.mcp_client: Optional[MCPClientManager] = None
        self.tools: List[BaseTool] = []
        self.agent: Any = None  # CompiledStateGraph
        self._initialized = False
    
    def _create_llm(self) -> BaseChatModel:
        """LLM 인스턴스 생성"""
        if self.config.model_provider == "openai":
            return ChatOpenAI(
                model=self.config.model_name,
                temperature=self.config.temperature,
                max_tokens=self.config.max_tokens,
                streaming=self.config.streaming
            )
        elif self.config.model_provider == "anthropic":
            return ChatAnthropic(
                model=self.config.model_name,
                temperature=self.config.temperature,
                max_tokens=self.config.max_tokens,
                streaming=self.config.streaming
            )
        else:
            raise ValueError(f"지원하지 않는 모델 제공자: {self.config.model_provider}")
    
    async def initialize(self) -> None:
        """에이전트 초기화 (MCP 연결 및 도구 로드)"""
        if self._initialized:
            return
        
        # LLM 생성
        self.llm = self._create_llm()
        
        # MCP 클라이언트 초기화 및 연결
        if self.config.mcp_config_path:
            self.mcp_client = MCPClientManager(self.config.mcp_config_path)
            await self.mcp_client.connect_all()
            self.tools = await self.mcp_client.get_langchain_tools()
            print(f"[Agent] {len(self.tools)}개의 MCP 도구 로드 완료")
        
        # 시스템 프롬프트 설정
        system_prompt = self.config.system_prompt or self.DEFAULT_SYSTEM_PROMPT
        
        # ReAct 에이전트 생성 (LangGraph 버전에 따라 파라미터가 다름)
        try:
            # 최신 버전: prompt 파라미터 사용
            self.agent = create_react_agent(
                model=self.llm,
                tools=self.tools,
                prompt=system_prompt
            )
        except TypeError:
            # 이전 버전: state_modifier 또는 기본값 사용
            try:
                self.agent = create_react_agent(
                    model=self.llm,
                    tools=self.tools,
                    state_modifier=system_prompt
                )
            except TypeError:
                # 파라미터 없이 생성
                self.agent = create_react_agent(
                    model=self.llm,
                    tools=self.tools
                )
        
        self._initialized = True
        print("[Agent] 초기화 완료")
    
    async def add_mcp_server(
        self,
        name: str,
        command: str,
        args: List[str] = None,
        env: Dict[str, str] = None
    ) -> None:
        """
        런타임에 MCP 서버 추가
        
        Args:
            name: 서버 이름
            command: 실행 명령어
            args: 명령어 인자
            env: 환경 변수
        """
        if not self.mcp_client:
            self.mcp_client = MCPClientManager()
        
        self.mcp_client.register_server(name, command, args, env)
        await self.mcp_client.connect_server(name)
        
        # 새 도구 추가
        new_tools = await self.mcp_client.get_langchain_tools(name)
        self.tools.extend(new_tools)
        
        # 에이전트 재생성
        if self.llm:
            system_prompt = self.config.system_prompt or self.DEFAULT_SYSTEM_PROMPT
            try:
                self.agent = create_react_agent(
                    model=self.llm,
                    tools=self.tools,
                    prompt=system_prompt
                )
            except TypeError:
                self.agent = create_react_agent(
                    model=self.llm,
                    tools=self.tools
                )
        
        print(f"[Agent] MCP 서버 '{name}' 추가됨 ({len(new_tools)}개 도구)")
    
    def _prepare_messages(
        self,
        user_input: str,
        context: AgentContext
    ) -> List[BaseMessage]:
        """대화 메시지 준비"""
        messages = []
        
        # 컨텍스트 정보를 시스템 메시지로 추가
        context_info = f"""현재 컨텍스트 정보:
- tenant_id: {context.tenant_id}
- user_uid: {context.user_uid}
- user_email: {context.user_email}
- user_name: {context.user_name}

도구 호출 시 위 정보를 적절히 사용하세요."""
        
        messages.append(SystemMessage(content=context_info))
        
        # 이전 대화 기록 추가
        messages.extend(context.conversation_history)
        
        # 현재 사용자 입력 추가
        messages.append(HumanMessage(content=user_input))
        
        return messages
    
    async def run(
        self,
        user_input: str,
        context: Optional[AgentContext] = None,
        stream: bool = False
    ) -> str | AsyncGenerator[str, None]:
        """
        에이전트 실행
        
        Args:
            user_input: 사용자 입력
            context: 에이전트 컨텍스트
            stream: 스트리밍 여부
            
        Returns:
            응답 문자열 또는 스트리밍 제너레이터
        """
        if not self._initialized:
            await self.initialize()
        
        context = context or AgentContext()
        messages = self._prepare_messages(user_input, context)
        
        if stream:
            return self._stream_response(messages, context)
        else:
            return await self._get_response(messages, context)
    
    async def _get_response(
        self,
        messages: List[BaseMessage],
        context: AgentContext
    ) -> str:
        """동기식 응답 생성"""
        input_state = {"messages": messages}
        
        result = await self.agent.ainvoke(input_state)
        
        # 마지막 AI 메시지 추출
        final_message = result["messages"][-1]
        response_text = final_message.content if hasattr(final_message, 'content') else str(final_message)
        
        # 대화 기록 업데이트
        context.conversation_history.append(HumanMessage(content=messages[-1].content))
        context.conversation_history.append(AIMessage(content=response_text))
        
        return response_text
    
    async def _stream_response(
        self,
        messages: List[BaseMessage],
        context: AgentContext
    ) -> AsyncGenerator[str, None]:
        """스트리밍 응답 생성"""
        input_state = {"messages": messages}
        full_response = ""
        
        async for event in self.agent.astream_events(input_state, version="v2"):
            kind = event["event"]
            
            if kind == "on_chat_model_stream":
                # LLM 토큰 스트리밍
                chunk = event["data"]["chunk"]
                if hasattr(chunk, "content") and chunk.content:
                    full_response += chunk.content
                    yield json.dumps({
                        "type": "token",
                        "content": chunk.content
                    })
            
            elif kind == "on_tool_start":
                # 도구 호출 시작
                tool_name = event["name"]
                tool_input = event["data"].get("input", {})
                yield json.dumps({
                    "type": "tool_start",
                    "tool": tool_name,
                    "input": tool_input
                })
            
            elif kind == "on_tool_end":
                # 도구 호출 종료
                tool_output = event["data"].get("output", "")
                yield json.dumps({
                    "type": "tool_end",
                    "output": str(tool_output)[:500]  # 출력 길이 제한
                })
        
        # 대화 기록 업데이트
        context.conversation_history.append(HumanMessage(content=messages[-1].content))
        context.conversation_history.append(AIMessage(content=full_response))
        
        yield json.dumps({
            "type": "done",
            "content": full_response
        })
    
    async def cleanup(self) -> None:
        """리소스 정리"""
        if self.mcp_client:
            await self.mcp_client.disconnect_all()
        self._initialized = False
        print("[Agent] 정리 완료")


# 편의 함수
async def create_agent(
    model_provider: str = "openai",
    model_name: str = "gpt-4o",
    mcp_config_path: Optional[str] = None,
    system_prompt: Optional[str] = None,
    **kwargs
) -> WorkAssistantAgent:
    """
    에이전트 생성 편의 함수
    
    사용 예:
        agent = await create_agent(
            model_provider="openai",
            model_name="gpt-4o",
            mcp_config_path="mcp_config.json"
        )
    """
    config = AgentConfig(
        model_provider=model_provider,
        model_name=model_name,
        mcp_config_path=mcp_config_path,
        system_prompt=system_prompt or "",
        **kwargs
    )
    
    agent = WorkAssistantAgent(config)
    await agent.initialize()
    return agent

