# schemas.py

from pydantic import BaseModel, Field
from typing import List, Literal


# ✅ Tool 1: 사용자 입력 → 키워드 추출

class ExtractKeywordsInput(BaseModel):
    user_input: str

class ExtractKeywordsOutput(BaseModel):
    keywords: List[str]


# ✅ Tool 2: 키워드 → 유사 키워드 검색

class SearchSimilarKeywordsInput(BaseModel):
    keywords: List[str]

class SimilarKeyword(BaseModel):
    keyword: str
    score: float

class SearchSimilarKeywordsOutput(BaseModel):
    similar_keywords: List[SimilarKeyword]


# ✅ Tool 3: 유사 키워드 → 가중 평균 벡터 생성

class ComputeUserEmbeddingInput(BaseModel):
    similar_keywords: List[SimilarKeyword]

class ComputeUserEmbeddingOutput(BaseModel):
    vector: List[float]


# ✅ Tool 4: 사용자 벡터 → 장소 추천

class RecommendPlacesInput(BaseModel):
    vector: List[float]

class RecommendedPlace(BaseModel):
    place_id: int
    similarity_score: float

class RecommendPlacesOutput(BaseModel):
    data: List[RecommendedPlace]

# tools.py

from langchain.tools import StructuredTool
from schemas import (
    ExtractKeywordsInput, ExtractKeywordsOutput,
    SearchSimilarKeywordsInput, SearchSimilarKeywordsOutput,
    ComputeUserEmbeddingInput, ComputeUserEmbeddingOutput,
    RecommendPlacesInput, RecommendPlacesOutput
)
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import GoogleGenerativeAI
import json
import faiss
import numpy as np
import pickle

# 🔧 공통 임베딩 모델
embedding_model = HuggingFaceEmbeddings(model_name="snunlp/KR-SBERT-V40K-klueNLI-augSTS")

# 🔧 Gemini LLM (예시)
llm = GoogleGenerativeAI(
    model="models/gemini-2.0-flash-lite",
    google_api_key="YOUR_GEMINI_API_KEY"
)


# ✅ Tool 1: 키워드 추출
def extract_keywords_fn(input: ExtractKeywordsInput) -> ExtractKeywordsOutput:
    prompt = f"""
    다음 사용자 입력에서 장소 추천을 위한 키워드를 JSON 형식으로 추출하세요.

    입력: "{input.user_input}"

    출력 예시:
    {{
        "keywords": ["데이트", "감성", "야경"]
    }}
    """
    response = llm(prompt)
    parsed = json.loads(response)
    return ExtractKeywordsOutput(**parsed)


# ✅ Tool 2: 유사 키워드 검색
def search_similar_keywords_fn(input: SearchSimilarKeywordsInput) -> SearchSimilarKeywordsOutput:
    index = faiss.read_index("keyword_vectors.index")
    with open("keyword_meta.pkl", "rb") as f:
        keyword_meta = pickle.load(f)

    vectors = embedding_model.embed_documents(input.keywords)
    query = np.mean(vectors, axis=0).astype("float32").reshape(1, -1)

    D, I = index.search(query, 5)
    similar = [{"keyword": keyword_meta[i], "score": float(D[0][j])} for j, i in enumerate(I[0])]
    return SearchSimilarKeywordsOutput(similar_keywords=similar)


# ✅ Tool 3: 가중 평균 벡터 생성
def compute_user_embedding_fn(input: ComputeUserEmbeddingInput) -> ComputeUserEmbeddingOutput:
    vectors = [embedding_model.embed_query(kw.keyword) for kw in input.similar_keywords]
    sims = [kw.score for kw in input.similar_keywords]

    weighted = sum(np.array(vec) * sim for vec, sim in zip(vectors, sims))
    avg_vector = weighted / sum(sims)

    return ComputeUserEmbeddingOutput(vector=avg_vector.tolist())


# ✅ Tool 4: 장소 추천
def recommend_places_fn(input: RecommendPlacesInput) -> RecommendPlacesOutput:
    index = faiss.read_index("place_vectors.index")
    with open("place_meta.pkl", "rb") as f:
        place_meta = pickle.load(f)

    vec = np.array(input.vector).astype("float32").reshape(1, -1)
    D, I = index.search(vec, 5)

    recommended = [
        {"place_id": int(place_meta[i]["id"]), "similarity_score": float(D[0][j])}
        for j, i in enumerate(I[0])
    ]

    return RecommendPlacesOutput(data=recommended)

# agent.py

from langchain.agents import AgentExecutor, Tool, initialize_agent
from langchain.agents.agent_types import AgentType
from langchain_community.llms import GoogleGenerativeAI
from langchain.tools import StructuredTool

from tools import (
    extract_keywords_fn,
    search_similar_keywords_fn,
    compute_user_embedding_fn,
    recommend_places_fn
)
from schemas import (
    ExtractKeywordsInput, ExtractKeywordsOutput,
    SearchSimilarKeywordsInput, SearchSimilarKeywordsOutput,
    ComputeUserEmbeddingInput, ComputeUserEmbeddingOutput,
    RecommendPlacesInput, RecommendPlacesOutput
)

# ✅ LLM (Gemini)
llm = GoogleGenerativeAI(
    model="models/gemini-pro",
    google_api_key="YOUR_GEMINI_API_KEY"
)

# ✅ Tool 목록 (구조화된 입력/출력 스키마 적용)
tools = [
    StructuredTool.from_function(
        name="extract_keywords",
        description="사용자 입력에서 장소 추천에 필요한 키워드를 추출합니다.",
        func=extract_keywords_fn,
        args_schema=ExtractKeywordsInput,
        return_schema=ExtractKeywordsOutput
    ),
    StructuredTool.from_function(
        name="search_similar_keywords",
        description="추천 키워드를 기반으로 벡터 DB에서 유사 키워드를 검색합니다.",
        func=search_similar_keywords_fn,
        args_schema=SearchSimilarKeywordsInput,
        return_schema=SearchSimilarKeywordsOutput
    ),
    StructuredTool.from_function(
        name="compute_user_embedding",
        description="유사 키워드와 유사도 점수를 기반으로 가중 평균 임베딩 벡터를 계산합니다.",
        func=compute_user_embedding_fn,
        args_schema=ComputeUserEmbeddingInput,
        return_schema=ComputeUserEmbeddingOutput
    ),
    StructuredTool.from_function(
        name="recommend_places",
        description="사용자 임베딩 벡터를 기반으로 장소 벡터 DB에서 추천 결과를 반환합니다.",
        func=recommend_places_fn,
        args_schema=RecommendPlacesInput,
        return_schema=RecommendPlacesOutput
    )
]

# ✅ Agent 생성
agent_executor = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT,
    verbose=True
)

# run.py

from agent import agent_executor

if __name__ == "__main__":
    # 📝 테스트용 사용자 입력
    user_query = "저녁에 분위기 좋은 데이트 장소 추천해줘"

    print("\n==============================")
    print("🔍 사용자 요청:", user_query)
    print("==============================\n")

    # 🧠 LangChain Agent 실행
    result = agent_executor.run(user_query)

    print("\n==============================")
    print("🏆 최종 추천 결과 (구조화된 JSON):")
    print(result)
    print("==============================")