from typing import TypedDict, List

class HyperGraphState(TypedDict):
    question: str
    matched_entities: List[str]
    hyperedges: List[dict]
    context: str
    answer: str

# 하이퍼그래프: 하이퍼엣지 ID → 엔티티 집합 + 관계 설명
# 실제 구현에서는 DB에 저장하며, 여기서는 dict로 시뮬레이션
HYPERGRAPH = {
    "he_001": {
        "entities": ["Transformer", "Vaswani", "Google Brain", "Attention Mechanism", "NeurIPS 2017"],
        "description": "Vaswani et al.이 Google Brain에서 Attention Mechanism 기반의 Transformer를 제안하여 NeurIPS 2017에 발표",
    },
    "he_002": {
        "entities": ["BERT", "Devlin", "Google AI", "Transformer", "사전학습"],
        "description": "Devlin et al.이 Google AI에서 Transformer 기반의 BERT 사전학습 모델을 제안",
    },
    "he_003": {
        "entities": ["GPT-3", "OpenAI", "Few-shot Learning", "Transformer", "언어 모델"],
        "description": "OpenAI가 Transformer 기반의 GPT-3 대규모 언어 모델로 Few-shot Learning 능력을 입증",
    },
}

from langchain.chat_models import init_chat_model
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

llm = init_chat_model("gpt-4o-mini", temperature=0)

def extract_query_entities(state: HyperGraphState) -> HyperGraphState:
    """질문에서 핵심 엔티티를 추출합니다."""
    question = state["question"]

    prompt = ChatPromptTemplate.from_messages([
        ("system", (
            "질문에서 핵심 엔티티(인물, 조직, 기술, 개념 등)를 추출하세요.\n"
            "쉼표로 구분하여 출력하세요."
        )),
        ("human", "{question}"),
    ])

    chain = prompt | llm | StrOutputParser()
    result = chain.invoke({"question": question})
    entities = [e.strip() for e in result.split(",") if e.strip()]
    return {"matched_entities": entities}

def find_hyperedges(state: HyperGraphState) -> HyperGraphState:
    """매칭된 엔티티를 포함하는 하이퍼엣지를 탐색합니다."""
    matched_entities = set(state["matched_entities"])
    relevant_edges = []

    for edge_id, edge_data in HYPERGRAPH.items():
        # 쿼리 엔티티와 하이퍼엣지의 엔티티가 겹치면 관련 있음
        overlap = matched_entities & set(edge_data["entities"])
        if overlap:
            relevant_edges.append({
                "id": edge_id,
                "entities": edge_data["entities"],
                "description": edge_data["description"],
                "overlap": list(overlap),
            })

    return {"hyperedges": relevant_edges}

def expand_context(state: HyperGraphState) -> HyperGraphState:
    """하이퍼엣지의 모든 엔티티와 관계를 컨텍스트로 조립합니다."""
    hyperedges = state["hyperedges"]

    if not hyperedges:
        return {"context": "관련 정보를 찾을 수 없습니다."}

    context_parts = []
    all_entities = set()

    for edge in hyperedges:
        all_entities.update(edge["entities"])
        desc = edge["description"]
        ents = ", ".join(edge["entities"])
        context_parts.append(f"[관계] {desc}\n  관련 엔티티: {ents}")

    context = f"발견된 엔티티: {', '.join(all_entities)}\n\n" + "\n\n".join(context_parts)
    return {"context": context}

def generate(state: HyperGraphState) -> HyperGraphState:
    """하이퍼그래프 컨텍스트를 기반으로 답변을 생성합니다."""
    question = state["question"]
    context = state["context"]

    prompt = ChatPromptTemplate.from_messages([
        ("system", (
            "다음은 하이퍼그래프에서 검색한 엔티티 관계 정보입니다.\n"
            "이 정보를 참고하여 질문에 답변하세요.\n\n{context}"
        )),
        ("human", "{question}"),
    ])

    chain = prompt | llm | StrOutputParser()
    answer = chain.invoke({"context": context, "question": question})
    return {"answer": answer}

from langgraph.graph import StateGraph, START, END

workflow = StateGraph(HyperGraphState)

# 노드 추가
workflow.add_node("extract_query_entities", extract_query_entities)
workflow.add_node("find_hyperedges", find_hyperedges)
workflow.add_node("expand_context", expand_context)
workflow.add_node("generate", generate)

# 엣지 연결
workflow.add_edge(START, "extract_query_entities")
workflow.add_edge("extract_query_entities", "find_hyperedges")
workflow.add_edge("find_hyperedges", "expand_context")
workflow.add_edge("expand_context", "generate")
workflow.add_edge("generate", END)

# 컴파일
app = workflow.compile()

result = app.invoke({
    "question": "Transformer 아키텍처가 후속 연구에 미친 영향과 관련 연구자들은?"
})
print(result["answer"])