from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, AIMessagePromptTemplate
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.memory import ConversationBufferMemory
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_core.messages import SystemMessage
from .config import GEMINI_API_KEY

from .database import Memory


chat_llm = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash",
    temperature=1.3,
    google_api_key=GEMINI_API_KEY,
    streaming=True
)

# memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
memory = Memory(memory_key="chat_history")


chat_template = ChatPromptTemplate.from_messages(
    [
        SystemMessage(
            content=(
                "You are Coach GPH, an expert Geophysicist with over 20 years of experience "
                "in the oil and gas industry and 10 years teaching and researching at Caltech. Your role is to answer technical questions clearly "
                "and concisely, in simple, easy-to-understand terms.\n\n"
                "Always follow these rules in your responses:\n"
                "1. Use Markdown formatting: headings, bold, italics, bullet points, and code blocks where appropriate.\n"
                "2. Break explanations into steps or numbered lists for clarity.\n"
                "3. Provide practical examples when explaining concepts.\n"
                "4. Keep a professional, friendly, and helpful tone.\n"
                "5. Avoid overly technical jargon unless necessary, and explain any technical terms you use.\n"
                "6. Ensure each response is self-contained and understandable even to someone with basic geophysics knowledge."
            )
        ),
        HumanMessagePromptTemplate.from_template("{question}"),
        AIMessagePromptTemplate.from_template("{chat_history}")
    ]
)

chain = (
    {"question": RunnablePassthrough(), "chat_history": memory.load_memory_variables}
    | chat_template
    | chat_llm
    | StrOutputParser()
)

def get_chain():
    return chain, memory