import gradio as gr
from gradio import ChatMessage
from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated
import operator
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
from langchain_openai import ChatOpenAI
import asyncio
import time
import os
from pprint import pprint

# 상태 정의
class ChatState(TypedDict):
    messages: Annotated[list[BaseMessage], operator.add]
    current_response: str
    agent_type: str
    context: dict
    step_info: str

# LLM 초기화
llm = ChatOpenAI(
    model="gpt-3.5-turbo",
    temperature=0.7,
)

def step1_analyzer_node(state: ChatState) -> ChatState:
    """1단계: 메시지 분석 노드"""
    last_message = state["messages"][-1].content

    # 분석 시뮬레이션 (실제로는 더 복잡한 로직)
    analysis = {
        "length": len(last_message),
        "has_question": "?" in last_message,
        "language": "korean" if any(ord(c) > 127 for c in last_message) else "english",
        "sentiment": "positive" if any(word in last_message.lower() for word in ["좋", "감사", "고마워", "thanks", "good"]) else "neutral"
    }

    return {
        "messages": [],
        "current_response": "",
        "agent_type": "",
        "context": {"analysis": analysis},
        "step_info": f"📊 **1단계 완료** - 메시지 분석\n- 길이: {analysis['length']}자\n- 질문 포함: {'예' if analysis['has_question'] else '아니오'}\n- 언어: {analysis['language']}\n- 감정: {analysis['sentiment']}"
    }

def step2_classifier_node(state: ChatState) -> ChatState:
    """2단계: 의도 분류 노드"""
    last_message = state["messages"][-1].content.lower()
    analysis = state["context"]["analysis"]

    # 의도 분류 로직
    if any(word in last_message for word in ["코드", "프로그래밍", "python", "개발", "함수", "클래스"]):
        agent_type = "programmer"
        confidence = 0.9
    elif any(word in last_message for word in ["날씨", "뉴스", "정보", "검색", "찾아"]):
        agent_type = "informer"
        confidence = 0.8
    elif any(word in last_message for word in ["계산", "수학", "더하기", "빼기", "곱하기", "나누기"]):
        agent_type = "calculator"
        confidence = 0.95
    elif any(word in last_message for word in ["창작", "시", "소설", "이야기", "글"]):
        agent_type = "creative"
        confidence = 0.85
    else:
        agent_type = "general"
        confidence = 0.7

    context = state["context"]
    context["classification"] = {
        "agent_type": agent_type,
        "confidence": confidence
    }

    return {
        "messages": [],
        "current_response": "",
        "agent_type": agent_type,
        "context": context,
        "step_info": f"🎯 **2단계 완료** - 의도 분류\n- 분류 결과: {agent_type}\n- 신뢰도: {confidence:.1%}\n- 다음 단계: {'전문 처리' if confidence > 0.8 else '일반 처리'}"
    }

def step3_context_enricher_node(state: ChatState) -> ChatState:
    """3단계: 컨텍스트 강화 노드"""
    agent_type = state["agent_type"]

    # 에이전트 타입별 컨텍스트 강화
    enriched_context = {
        "programmer": {
            "system_prompt": "당신은 경험이 풍부한 시니어 개발자입니다. 코드 예시와 함께 명확하고 실용적인 답변을 제공하세요.",
            "tools": ["코드 실행", "문서 검색", "베스트 프랙티스"],
            "style": "기술적이고 정확한"
        },
        "informer": {
            "system_prompt": "당신은 정보 전문가입니다. 정확하고 최신의 정보를 구조화된 형태로 제공하세요.",
            "tools": ["웹 검색", "팩트 체크", "데이터 분석"],
            "style": "객관적이고 상세한"
        },
        "calculator": {
            "system_prompt": "당신은 수학 전문가입니다. 계산 과정을 단계별로 설명하고 정확한 답을 제공하세요.",
            "tools": ["수식 계산", "그래프 생성", "통계 분석"],
            "style": "논리적이고 체계적인"
        },
        "creative": {
            "system_prompt": "당신은 창작 전문가입니다. 상상력이 풍부하고 감성적인 콘텐츠를 제작하세요.",
            "tools": ["스토리텔링", "시각적 묘사", "감정 표현"],
            "style": "창의적이고 감성적인"
        },
        "general": {
            "system_prompt": "당신은 친근하고 도움이 되는 AI 어시스턴트입니다. 자연스럽고 이해하기 쉬운 답변을 제공하세요.",
            "tools": ["일반 대화", "정보 제공", "문제 해결"],
            "style": "친근하고 자연스러운"
        }
    }

    context = state["context"]
    context["enriched"] = enriched_context.get(agent_type, enriched_context["general"])

    return {
        "messages": [],
        "current_response": "",
        "agent_type": agent_type,
        "context": context,
        "step_info": f"🔧 **3단계 완료** - 컨텍스트 강화\n- 에이전트: {agent_type}\n- 스타일: {context['enriched']['style']}\n- 활용 도구: {', '.join(context['enriched']['tools'][:2])}"
    }

def step4_response_generator_node(state: ChatState) -> ChatState:
    """4단계: 응답 생성 노드"""
    enriched_context = state["context"]["enriched"]
    system_prompt = enriched_context["system_prompt"]

    # 시스템 프롬프트와 함께 메시지 구성
    messages = [HumanMessage(content=system_prompt)] + state["messages"]

    try:
        response = llm.invoke(messages)

        # 에이전트 타입에 따른 아이콘 설정
        icons = {
            "programmer": "💻",
            "informer": "📚",
            "calculator": "🔢",
            "creative": "🎨",
            "general": "💬"
        }

        icon = icons.get(state["agent_type"], "💬")
        final_response = f"{icon} **[{state['agent_type'].upper()}]**\n\n{response.content}"

        return {
            "messages": [response],
            "current_response": final_response,
            "agent_type": state["agent_type"],
            "context": state["context"],
            "step_info": f"✅ **4단계 완료** - 응답 생성\n- 최종 응답 준비됨\n- 응답 길이: {len(response.content)}자"
        }

    except Exception as e:
        error_msg = f"❌ 응답 생성 중 오류 발생: {str(e)}"
        return {
            "messages": [AIMessage(content=error_msg)],
            "current_response": error_msg,
            "agent_type": state["agent_type"],
            "context": state["context"],
            "step_info": f"❌ **4단계 실패** - 오류 발생\n- 오류: {str(e)}"
        }

def should_continue_to_classifier(state: ChatState) -> str:
    return "classifier"

def should_continue_to_enricher(state: ChatState) -> str:
    return "enricher"

def should_continue_to_generator(state: ChatState) -> str:
    return "generator"

def should_end(state: ChatState) -> str:
    return END

# LangGraph 워크플로우 생성
def create_enhanced_workflow():
    workflow = StateGraph(ChatState)

    # 4개 노드 추가
    workflow.add_node("analyzer", step1_analyzer_node)
    workflow.add_node("classifier", step2_classifier_node)
    workflow.add_node("enricher", step3_context_enricher_node)
    workflow.add_node("generator", step4_response_generator_node)

    # 시작점 설정
    workflow.set_entry_point("analyzer")

    # 순차적 엣지 추가
    workflow.add_conditional_edges(
        "analyzer",
        should_continue_to_classifier,
        {"classifier": "classifier"}
    )

    workflow.add_conditional_edges(
        "classifier",
        should_continue_to_enricher,
        {"enricher": "enricher"}
    )

    workflow.add_conditional_edges(
        "enricher",
        should_continue_to_generator,
        {"generator": "generator"}
    )

    workflow.add_conditional_edges(
        "generator",
        should_end,
        {END: END}
    )

    return workflow.compile()

# 글로벌 워크플로우 인스턴스
enhanced_workflow = create_enhanced_workflow()

def stream_chatbot_response(message, history):
    """각 단계를 누적해서 실시간 표시"""
    if not message.strip():
        yield "", history
        return

    # 메시지 히스토리를 LangChain 메시지로 변환
    messages = []
    for human_msg, ai_msg in history:
        if human_msg and not "📊" in human_msg and not "🎯" in human_msg:
            messages.append(HumanMessage(content=human_msg))
        if ai_msg and not "📊" in ai_msg and not "🎯" in ai_msg and not "⚡" in ai_msg:
            messages.append(AIMessage(content=ai_msg))

    # 현재 메시지 추가
    messages.append(HumanMessage(content=message))

    # 초기 상태 설정
    initial_state = {
        "messages": messages,
        "current_response": "",
        "agent_type": "",
        "context": {},
        "step_info": ""
    }

    # 워크플로우를 스트림으로 실행
    current_history = history.copy()
    yield "", current_history
    time.sleep(0.3)

    # 각 노드를 순차적으로 실행하면서 중간 결과를 누적 표시
    for step_result in enhanced_workflow.stream(initial_state):
        node_name = list(step_result.keys())[0]
        node_result = step_result[node_name]
        pprint(step_result)
        if "step_info" in node_result and node_result["step_info"]:
            # 현재 단계 정보를 누적 리스트에 추가
            step_info = node_result["step_info"]
            pprint(step_info)
            # 누적된 모든 단계 정보를 표시
            current_history.append(ChatMessage( role="assistant",
                                        content=step_info,
                                        metadata={"title": f"{node_name}", "status": "done"}))
            yield "", current_history
            time.sleep(0.2)  # 시각적 효과를 위한 지연

    current_history.append(ChatMessage( role="assistant",
                                        content=step_result["generator"]["current_response"]))

    yield "", current_history


def clear_chat():
    """채팅 히스토리 초기화"""
    return []

# Gradio 인터페이스 생성
def create_enhanced_gradio_interface():
    with gr.Blocks(title="Enhanced LangGraph 챗봇", theme=gr.themes.Soft()) as demo:
        gr.Markdown(
            """
            # 🚀 Enhanced LangGraph + Gradio 챗봇

            **4단계 처리 과정을 실시간으로 확인할 수 있는 AI 챗봇**

            **처리 단계:**
            1. 📊 **메시지 분석** - 입력 메시지의 특성 분석
            2. 🎯 **의도 분류** - 사용자 의도에 따른 에이전트 선택
            3. 🔧 **컨텍스트 강화** - 전문 도메인별 컨텍스트 설정
            4. ✅ **응답 생성** - 최적화된 답변 생성

            **지원 에이전트:** 💻 프로그래머 | 📚 정보전문가 | 🔢 계산기 | 🎨 창작가 | 💬 일반대화
            """
        )

        # 챗봇 컴포넌트
        chatbot = gr.Chatbot(
            value=[],
            height=400,
            show_label=False,
            container=True,
            type="messages"
        )

        with gr.Row():
            msg = gr.Textbox(
                placeholder="메시지를 입력하세요... (각 처리 단계가 실시간으로 표시됩니다)",
                show_label=False,
                scale=4,
                container=False
            )
            submit_btn = gr.Button("전송", scale=1, variant="primary")
            clear_btn = gr.Button("초기화", scale=1, variant="secondary")

        # 상태 표시
        with gr.Row():
            gr.Markdown("💡 **팁**: 다양한 주제로 대화해보세요. 각 단계별 처리 과정을 확인할 수 있습니다!")

        # 이벤트 핸들러 - 스트리밍 방식으로 변경
        submit_btn.click(
            stream_chatbot_response,
            inputs=[msg, chatbot],
            outputs=[msg, chatbot]
        )

        msg.submit(
            stream_chatbot_response,
            inputs=[msg, chatbot],
            outputs=[msg, chatbot]
        )

        # 초기화 버튼
        clear_btn.click(
            clear_chat,
            outputs=[chatbot]
        )

        # 카테고리별 예제 질문들
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 💻 프로그래밍")
                gr.Examples(
                    examples=[
                        "Python으로 피보나치 수열 함수 만드는 방법?",
                        "딕셔너리와 리스트의 차이점을 알려줘",
                        "클래스와 객체에 대해 설명해줘"
                    ],
                    inputs=msg
                )

            with gr.Column():
                gr.Markdown("### 🔢 계산/수학")
                gr.Examples(
                    examples=[
                        "25 곱하기 37은 얼마야?",
                        "복리 계산 방법을 알려줘",
                        "삼각함수에 대해 설명해줘"
                    ],
                    inputs=msg
                )

            with gr.Column():
                gr.Markdown("### 🎨 창작")
                gr.Examples(
                    examples=[
                        "봄에 대한 짧은 시를 써줘",
                        "우주 여행 이야기를 만들어줘",
                        "창의적인 아이디어를 제안해줘"
                    ],
                    inputs=msg
                )

    return demo

demo = create_enhanced_gradio_interface()
if __name__ == "__main__":
    demo.launch()