"""
다중 에이전트 협업 시스템 예제

이 예제는 LangGraph를 활용하여 여러 전문 에이전트가 협업하여 복잡한 문제를 해결하는 시스템을 구현합니다.
각 에이전트는 특정 전문 영역을 담당하며, 그래프 구조에 따라 상호작용하여 최종 결과를 도출합니다.

요구사항:
- 필요한 패키지: langchain, langgraph, dotenv
- OpenAI API 키 필요 (.env 파일에 설정)

실행 방법:
$ python multi_agent_collaboration.py
"""

import operator
import os
from typing import Annotated, Dict, List, Sequence, TypedDict, Union

from dotenv import load_dotenv
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph
from langgraph.checkpoint import MemorySaver


# --- 1. 환경 설정 ---
load_dotenv()
if not os.getenv("OPENAI_API_KEY"):
    raise ValueError("OPENAI_API_KEY가 설정되지 않았습니다. .env 파일을 확인하세요.")


# --- 2. 협업 에이전트 상태 정의 ---
class TeamState(TypedDict):
    """팀 협업을 위한 상태 클래스"""
    messages: Annotated[Sequence[BaseMessage], operator.add]  # 메시지 히스토리
    next: str  # 다음에 어떤 에이전트로 전환할지 지정
    task_status: Dict[str, str]  # 각 작업의 상태 추적
    artifacts: Dict[str, str]  # 에이전트들이 생성한 결과물 저장


# --- 3. 각 에이전트 역할 정의 ---
TEAM_ROLES = {
    "project_manager": """당신은 프로젝트 매니저입니다. 팀 전체를 조율하고 작업 방향을 결정합니다.
    
    당신의 임무:
    1. 사용자의 요청을 분석하여 필요한 작업을 식별하고 단계별로 분해합니다.
    2. 각 전문가에게 명확한 지시와 우선순위를 제공합니다.
    3. 각 전문가의 결과물을 검토하고 필요시 추가 작업을 요청합니다.
    4. 최종 결과물이 사용자의 요구사항을 충족하는지 확인합니다.
    
    항상 명확하고 구체적인 지시사항을 제공하세요.
    
    결과물을 검토할 때는 다음 에이전트를 지정하세요:
    - 데이터 분석이 필요하면: "data_scientist"
    - 알고리즘 설계가 필요하면: "algorithm_expert" 
    - 최종 보고서 작성이 필요하면: "technical_writer"
    - 모든 작업이 완료되었으면: "final"
    """,
    
    "data_scientist": """당신은 데이터 사이언티스트입니다. 데이터를 분석하고 통찰력을 제공합니다.
    
    당신의 임무:
    1. 주어진 데이터셋에 대한 탐색적 분석을 수행합니다.
    2. 데이터의 패턴, 이상점, 분포 특성을 파악합니다.
    3. 문제 해결에 필요한 핵심 인사이트를 도출합니다.
    4. 적합한 데이터 전처리 방법과 분석 기법을 제안합니다.
    
    분석 결과는 항상 "데이터 분석 결과: " 접두사로 시작하는 명확한 요약을 포함해야 합니다.
    분석 완료 후, 메시지 끝에 "next: project_manager"를 추가하여 프로젝트 매니저에게 결과를 전달하세요.
    """,
    
    "algorithm_expert": """당신은 알고리즘 전문가입니다. 효율적인 알고리즘과 모델을 설계합니다.
    
    당신의 임무:
    1. 문제에 가장 적합한 알고리즘과 모델을 제안합니다.
    2. 제안된 접근 방식의 시간/공간 복잡도를 분석합니다.
    3. 알고리즘의 강점과 약점, 한계를 설명합니다.
    4. 가능한 대안적 접근법과 그 장단점도 제시합니다.
    
    설계 결과는 항상 "알고리즘 설계 결과: " 접두사로 시작하는 명확한 요약을 포함해야 합니다.
    설계 완료 후, 메시지 끝에 "next: project_manager"를 추가하여 프로젝트 매니저에게 결과를 전달하세요.
    """,
    
    "technical_writer": """당신은 기술 작가입니다. 복잡한 기술 내용을 명확하게 문서화합니다.
    
    당신의 임무:
    1. 팀 내 다른 전문가들의 기술적 결과물을 이해하기 쉽게 정리합니다.
    2. 전체 프로젝트의 결과를 구조화된 보고서 형태로 작성합니다.
    3. 기술적 내용과 비즈니스 가치를 균형있게 표현합니다.
    4. 시각화와 예시를 통해 복잡한 개념을 쉽게 설명합니다.
    
    작성 결과는 항상 "## 최종 보고서"로 시작하는 마크다운 형식으로 작성하세요.
    보고서 작성 후, 메시지 끝에 "next: project_manager"를 추가하여 프로젝트 매니저의 최종 검토를 받으세요.
    """
}


# --- 4. 각 에이전트 노드 생성 함수 ---
def create_agent_node(agent_name: str, llm):
    """특정 역할을 담당하는 에이전트 노드를 생성"""
    
    def agent_fn(state: TeamState):
        # 이전 메시지 히스토리와 시스템 메시지 준비
        system_message = SystemMessage(content=TEAM_ROLES[agent_name])
        messages = [system_message] + list(state["messages"])
        
        # LLM 호출하여 응답 생성
        response = llm.invoke(messages)
        
        # 응답에서 다음 에이전트 결정 (기본값은 project_manager)
        next_agent = "project_manager"
        
        if agent_name == "project_manager":
            # 프로젝트 매니저는 next 필드를 명시적으로 지정할 수 있음
            if "next: data_scientist" in response.content.lower():
                next_agent = "data_scientist"
            elif "next: algorithm_expert" in response.content.lower():
                next_agent = "algorithm_expert"
            elif "next: technical_writer" in response.content.lower():
                next_agent = "technical_writer"
            elif "next: final" in response.content.lower() or "작업 완료" in response.content:
                next_agent = "final"
        else:
            # 전문가 에이전트들은 항상 결과를 매니저에게 전달
            if "데이터 분석 결과: " in response.content:
                # 데이터 분석 결과 저장
                result_start = response.content.find("데이터 분석 결과: ")
                result_end = response.content.find("next:", result_start) if "next:" in response.content else len(response.content)
                state["artifacts"]["data_analysis"] = response.content[result_start:result_end].strip()
            
            elif "알고리즘 설계 결과: " in response.content:
                # 알고리즘 설계 결과 저장
                result_start = response.content.find("알고리즘 설계 결과: ")
                result_end = response.content.find("next:", result_start) if "next:" in response.content else len(response.content)
                state["artifacts"]["algorithm_design"] = response.content[result_start:result_end].strip()
            
            elif "## 최종 보고서" in response.content:
                # 최종 보고서 저장
                state["artifacts"]["final_report"] = response.content
        
        # 상태 업데이트
        return {
            "messages": [response],
            "next": next_agent,
            "task_status": {**state["task_status"], agent_name: "completed"},
            "artifacts": state["artifacts"]
        }
    
    return agent_fn


# --- 5. 라우팅 함수 정의 ---
def router(state: TeamState) -> str:
    """현재 상태에 따라 다음에 실행할 에이전트를 결정"""
    return state["next"]


# --- 6. 그래프 구축 ---
def build_team_graph(temperature=0):
    # LLM 모델 설정
    llm = ChatOpenAI(model="gpt-4-turbo", temperature=temperature)
    
    # 각 에이전트 노드 생성
    nodes = {
        "project_manager": create_agent_node("project_manager", llm),
        "data_scientist": create_agent_node("data_scientist", llm),
        "algorithm_expert": create_agent_node("algorithm_expert", llm),
        "technical_writer": create_agent_node("technical_writer", llm),
    }
    
    # 팀 워크플로우 그래프 구축
    workflow = StateGraph(TeamState)
    
    # 노드 추가
    for name, node_fn in nodes.items():
        workflow.add_node(name, node_fn)
    
    # 엣지 및 조건부 라우팅 설정
    workflow.set_entry_point("project_manager")
    workflow.add_conditional_edges(
        "project_manager",
        router,
        {
            "data_scientist": "data_scientist",
            "algorithm_expert": "algorithm_expert", 
            "technical_writer": "technical_writer",
            "final": END
        }
    )
    
    # 모든 전문가 노드는 프로젝트 매니저로 돌아감
    for expert in ["data_scientist", "algorithm_expert", "technical_writer"]:
        workflow.add_edge(expert, "project_manager")
    
    return workflow.compile()


# --- 7. 메인 실행 함수 ---
def main():
    print("다중 에이전트 협업 시스템을 초기화합니다...")
    
    # 그래프 컴파일
    team_app = build_team_graph(temperature=0.1)
    
    # 영구 저장소 설정 (선택사항)
    memory_saver = MemorySaver()
    
    # 초기 메시지 설정 (사용자 요청)
    task_description = """
    전자상거래 웹사이트의 제품 추천 시스템을 개선하고 싶습니다. 
    현재 시스템은 단순히 인기 제품이나 최근 조회 제품을 기반으로 추천하고 있어 개인화가 부족합니다.
    사용자의 구매 이력, 검색 패턴, 체류 시간 등의 데이터를 활용해 머신러닝 기반의 개선된 추천 시스템을 설계해주세요.
    최종 결과물로 데이터 분석 방향, 적합한 알고리즘, 구현 방안이 포함된 기술 보고서가 필요합니다.
    """
    
    initial_state = {
        "messages": [HumanMessage(content=task_description)],
        "next": "project_manager",
        "task_status": {},
        "artifacts": {}
    }
    
    print("\n작업을 시작합니다...\n" + "="*50)
    
    # 스트리밍 실행 (각 단계별 출력)
    for i, state in enumerate(team_app.stream(initial_state, config={"checkpointer": memory_saver})):
        if i > 0:  # 초기 상태는 건너뜀
            last_message = state["messages"][-1].content
            current_agent = state["next"]
            
            print(f"\n{'='*50}")
            print(f"STEP {i}: {'최종 결과' if current_agent == 'final' else current_agent}")
            print(f"{'-'*50}")
            print(last_message[:1000] + ("..." if len(last_message) > 1000 else ""))
            print(f"{'='*50}\n")
    
    # 최종 결과 출력
    final_state = memory_saver.get_latest()
    if "final_report" in final_state["artifacts"]:
        print("\n\n최종 보고서:")
        print("="*50)
        print(final_state["artifacts"]["final_report"])
    
    print("\n협업 에이전트 작업이 완료되었습니다.")


if __name__ == "__main__":
    main() 