# 단기 메모리 관리 예시 (Python)
class ShortTermMemory:
    def __init__(self, max_tokens=4000):
        self.messages = []
        self.max_tokens = max_tokens
    
    def add_message(self, role, content):
        self.messages.append({"role": role, "content": content})
        # 토큰 초과 시 오래된 메시지 제거
        while self._count_tokens() > self.max_tokens:
            # 시스템 메시지는 유지, 가장 오래된 대화 제거
            if len(self.messages) > 1:
                self.messages.pop(1)
    
    def _count_tokens(self):
        # 간단한 토큰 추정 (실제로는 tiktoken 등 사용)
        return sum(len(m["content"].split()) * 1.3 
                   for m in self.messages)

# RAG 기반 메모리 검색 예시
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings

class LongTermMemory:
    def __init__(self):
        self.embeddings = OpenAIEmbeddings(
            model="text-embedding-3-small"
        )
        self.vectorstore = Chroma(
            embedding_function=self.embeddings,
            persist_directory="./memory_db"
        )
    
    def store(self, content: str, metadata: dict = {}):
        """기억 저장"""
        self.vectorstore.add_texts(
            texts=[content],
            metadatas=[metadata]
        )
    
    def retrieve(self, query: str, k: int = 5):
        """관련 기억 검색"""
        results = self.vectorstore.similarity_search_with_score(
            query=query,
            k=k
        )
        # 유사도 점수 기준 필터링 (0.7 이상만)
        return [doc for doc, score in results if score > 0.7]
    
    def retrieve_with_mmr(self, query: str, k: int = 5):
        """MMR로 다양성 확보한 검색"""
        # Maximum Marginal Relevance
        # 유사도 높으면서도 서로 다양한 결과 반환
        return self.vectorstore.max_marginal_relevance_search(
            query=query, k=k, fetch_k=20
        )

# 메모리 컨트롤러 패턴
class MemoryController:
    def __init__(self):
        self.short_term = ShortTermMemory(max_tokens=4000)
        self.long_term = LongTermMemory()
        self.importance_threshold = 0.7
    
    def process_interaction(self, user_input, agent_response):
        # 1. 단기 메모리에 추가
        self.short_term.add_message("user", user_input)
        self.short_term.add_message("assistant", agent_response)
        
        # 2. 중요도 평가 (LLM 활용)
        importance = self._evaluate_importance(
            user_input, agent_response
        )
        
        # 3. 중요한 정보는 장기 메모리에 저장
        if importance >= self.importance_threshold:
            memory_text = f"사용자: {user_input}\n응답: {agent_response}"
            self.long_term.store(
                content=memory_text,
                metadata={
                    "timestamp": datetime.now().isoformat(),
                    "importance": importance,
                    "type": "interaction"
                }
            )
    
    def get_context_for_query(self, query: str) -> str:
        # 장기 메모리에서 관련 정보 검색
        memories = self.long_term.retrieve(query, k=3)
        
        # 단기 메모리의 최근 대화
        recent = self.short_term.get_recent(n=10)
        
        return self._format_context(memories, recent)
    
    def _evaluate_importance(self, user_input, response) -> float:
        # 실제로는 LLM을 사용해 중요도 평가
        # 여기서는 간단한 휴리스틱 사용
        important_keywords = ["기억해", "중요", "반드시", "이름", "날짜"]
        score = sum(1 for kw in important_keywords 
                   if kw in user_input or kw in response)
        return min(score / len(important_keywords), 1.0)

import numpy as np

def cosine_similarity(vec_a, vec_b):
    """두 벡터의 코사인 유사도 계산"""
    dot_product = np.dot(vec_a, vec_b)
    norm_a = np.linalg.norm(vec_a)
    norm_b = np.linalg.norm(vec_b)
    return dot_product / (norm_a * norm_b)

# 예시
query_vec = embed("오늘 날씨 어때?")
memory_vec = embed("어제 비가 많이 왔어요")
similarity = cosine_similarity(query_vec, memory_vec)
# similarity ≈ 0.72 (날씨 관련이라 유사도 높음)

import math
from datetime import datetime

def temporal_decay_score(base_score, timestamp, decay_factor=0.99):
    """시간 경과에 따른 점수 감쇠"""
    hours_elapsed = (datetime.now() - timestamp).total_seconds() / 3600
    # 지수 감쇠 함수
    temporal_weight = math.exp(-decay_factor * hours_elapsed / 24)
    return base_score * temporal_weight

# 예시: 1주일 전 기억은 현재 기억의 약 93% 점수
# 1달 전 기억은 약 74% 점수
# 1년 전 기억은 약 7% 점수

# 메모리 압축 예시
async def consolidate_memories(memories: list, llm) -> str:
    """여러 기억을 하나의 요약으로 압축"""
    memory_texts = "\n".join([m.content for m in memories])
    
    prompt = f"""다음 기억들을 분석하여 핵심 인사이트를 추출해줘:

{memory_texts}

다음 형식으로 요약해줘:
- 사용자 특성: [핵심 특성]
- 주요 관심사: [관심 분야]
- 중요 사실: [기억해야 할 사실들]"""
    
    summary = await llm.generate(prompt)
    return summary

# 압축 트리거: 같은 주제의 기억이 10개 이상 쌓이면 실행

# LangChain 하이브리드 메모리 구현
from langchain.memory import (
    ConversationSummaryBufferMemory,
    VectorStoreRetrieverMemory
)
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.chat_models import ChatOpenAI

# 단기 메모리: 최근 대화 + 요약
short_term = ConversationSummaryBufferMemory(
    llm=ChatOpenAI(model="gpt-4o-mini"),
    max_token_limit=2000,
    return_messages=True
)

# 장기 메모리: 벡터 DB 기반
vectorstore = Chroma(
    embedding_function=OpenAIEmbeddings(),
    persist_directory="./long_term_memory"
)
long_term = VectorStoreRetrieverMemory(
    retriever=vectorstore.as_retriever(
        search_kwargs={"k": 3}
    )
)

# 두 메모리를 결합한 컨텍스트 생성
def get_combined_context(query: str) -> str:
    # 장기 메모리에서 관련 기억 검색
    relevant_memories = long_term.load_memory_variables(
        {"input": query}
    )
    
    # 단기 메모리에서 최근 대화 가져오기
    recent_history = short_term.load_memory_variables({})
    
    context = f"""관련 기억:
{relevant_memories.get('history', '없음')}

최근 대화:
{recent_history.get('history', '없음')}"""
    
    return context

# LangGraph 메모리 에이전트 기본 구조
from langgraph.graph import StateGraph, END
from typing import TypedDict, List, Annotated
import operator

# 에이전트 상태 정의 (메모리 포함)
class AgentState(TypedDict):
    messages: Annotated[List, operator.add]  # 대화 히스토리
    memories: List[str]                       # 검색된 장기 기억
    current_query: str                        # 현재 쿼리
    importance_score: float                   # 현재 상호작용 중요도

# 메모리 검색 노드
def retrieve_memories(state: AgentState) -> AgentState:
    query = state["current_query"]
    # 벡터 DB에서 관련 기억 검색
    memories = memory_db.search(query, k=3)
    return {"memories": memories}

# LLM 추론 노드
def agent_reasoning(state: AgentState) -> AgentState:
    context = "\n".join(state["memories"])
    messages = state["messages"]
    # 메모리 컨텍스트를 포함한 LLM 호출
    response = llm.invoke(
        system=f"관련 기억:\n{context}",
        messages=messages
    )
    return {"messages": [response]}

# 메모리 저장 노드
def save_memory(state: AgentState) -> AgentState:
    if state["importance_score"] > 0.7:
        interaction = f"Q: {state['current_query']}\nA: {state['messages'][-1].content}"
        memory_db.store(interaction)
    return state

# 그래프 구성
workflow = StateGraph(AgentState)
workflow.add_node("retrieve", retrieve_memories)
workflow.add_node("reason", agent_reasoning)
workflow.add_node("save", save_memory)

workflow.set_entry_point("retrieve")
workflow.add_edge("retrieve", "reason")
workflow.add_edge("reason", "save")
workflow.add_edge("save", END)

app = workflow.compile()

# Mem0 사용 예시
from mem0 import Memory

m = Memory()

# 메모리 추가 (자동으로 중요 정보 추출)
m.add(
    "저는 파이썬 개발자이고 머신러닝에 관심이 많아요. "
    "특히 NLP 분야를 공부하고 있어요.",
    user_id="user_123"
)

# 관련 메모리 검색
memories = m.search(
    query="사용자의 기술 스택은?",
    user_id="user_123"
)
# 결과: [{"memory": "파이썬 개발자, 머신러닝/NLP 관심"}]

# 메모리 업데이트 (자동 충돌 해결)
m.add(
    "최근에 취직해서 이제 풀타임 ML 엔지니어로 일하고 있어요",
    user_id="user_123"
)
# 자동으로 "학생" → "ML 엔지니어"로 업데이트

# 메타 메모리 구현 예시
class MetaMemory:
    def __init__(self):
        self.memory_index = {}  # 메모리 목록과 신뢰도
    
    def register_memory(self, memory_id, content, confidence=1.0):
        self.memory_index[memory_id] = {
            "content": content,
            "confidence": confidence,
            "access_count": 0,
            "last_accessed": datetime.now(),
            "source": "user_interaction"
        }
    
    def get_memory_confidence(self, memory_id) -> float:
        """특정 기억의 신뢰도 반환"""
        if memory_id in self.memory_index:
            return self.memory_index[memory_id]["confidence"]
        return 0.0
    
    def should_verify(self, memory_id, threshold=0.8) -> bool:
        """신뢰도가 낮으면 검증 필요 여부 반환"""
        confidence = self.get_memory_confidence(memory_id)
        return confidence < threshold
    
    def decay_confidence(self, memory_id, decay_rate=0.01):
        """시간이 지남에 따라 신뢰도 감소"""
        if memory_id in self.memory_index:
            current = self.memory_index[memory_id]["confidence"]
            self.memory_index[memory_id]["confidence"] = max(
                0.1, current * (1 - decay_rate)
            )