from langchain_ollama import ChatOllama
from langchain_ibm import ChatWatsonx
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import (
    StrOutputParser,
    JsonOutputParser,
    PydanticOutputParser,
)
from langchain_core.chat_history import InMemoryChatMessageHistory, BaseChatMessageHistory
from langchain_core.runnables.history import  RunnableWithMessageHistory
from pydantic import BaseModel, Field
from typing import Literal
from dotenv import load_dotenv
import os
import gradio as gr
import uuid

# qwen_llm = ChatOllama(model="qwen3.5:4b")
# exaone_llm = ChatOllama(model="exaone3.5:2.4b")

store = {}

def get_sesstion_history(session_id) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
    return store[session_id]

def create_chain():

    gemma_llm = ChatOllama(model="gemma4:e2b")

    system_prompt = """\
        당신은 20년 경력의 전문 셰프이자 요리 연구가입니다.
        사용자의 요리 질문에 대해
        재료, 조리방법, 실패 방지 팁, 대체 재료를
        포함하여 답변하세요.
        항상 한국어로 답변하세요
        """

    template = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt),
            MessagesPlaceholder(variable_name="history"), # 대화 이력이 삽입될 위치
            ("human", "{question}"),
        ]
    )

    chain = template | gemma_llm | StrOutputParser()

    return RunnableWithMessageHistory(chain, get_sesstion_history, input_messages_key="question", history_messages_key="history")

chain = create_chain()

def chat(question, history, session_id):
    full_response = ""
    for chunk in chain.stream({"question": question}, config={"configurable": {"session_id":session_id}}):
        full_response+=chunk
        yield  full_response

with gr.Blocks() as chatbot:
    # 세션ID uuid 사용
    # gr.State() : 사용자별 데이터를 서버 메모리에 저장하는 컴포넌트
    session_state = gr.State(str(uuid.uuid4()))

    gr.ChatInterface(
        fn=chat,
        additional_inputs=[session_state],
        title="🍝 요리 전문가",
        description="요리경력 20년의 전문가입니다. 요리에 대한 궁금증을 해결하세요.",
    )

chatbot.launch()