import os
from langchain_core.prompts import ChatPromptTemplate
from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, MessagesPlaceholder
import yaml

current_dir = os.path.dirname(os.path.abspath(__file__))


# Function Agent
with open(f"{current_dir}/rag_template.yaml", "r") as yaml_file:
    rag_templates = yaml.safe_load(yaml_file)

sys_msg_template: str = rag_templates["sys_msg"]
human_msg_template: str = rag_templates["human_msg"]

rag_agent_prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(sys_msg_template),
    HumanMessagePromptTemplate.from_template(human_msg_template),
    MessagesPlaceholder(variable_name = "agent_scratchpad")
])


# Question Rephrase Agent
with open(f"{current_dir}/question_rephrase_template.yaml", "r") as yaml_file:
    rephrase_templates = yaml.safe_load(yaml_file)

sys_msg_template: str = rephrase_templates["sys_msg"]
human_msg_template: str = rephrase_templates["human_msg"]

rephrase_agent_prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(sys_msg_template),
    HumanMessagePromptTemplate.from_template(human_msg_template),
])


# Router Agent
with open(f"{current_dir}/router_template.yaml", "r") as yaml_file:
    router_templates = yaml.safe_load(yaml_file)

sys_msg_template: str = router_templates["sys_msg"]

router_agent_prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(sys_msg_template),
])


# Document Grader Agent
with open(f"{current_dir}/doc_grader_template.yaml", "r") as yaml_file:
    grader_templates = yaml.safe_load(yaml_file)

sys_msg_template: str = grader_templates["sys_msg"]
human_msg_template: str = grader_templates["human_msg"]

doc_grader_agent_prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(sys_msg_template),
    HumanMessagePromptTemplate.from_template(human_msg_template),
])


# Hallucination Grader Agent
with open(f"{current_dir}/hallucination_grader_template.yaml", "r") as yaml_file:
    hallucination_grader_templates = yaml.safe_load(yaml_file)

sys_msg_template: str = hallucination_grader_templates["sys_msg"]

hallucination_grader_agent_prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(sys_msg_template),
])


# Answer Grader Agent
with open(f"{current_dir}/answer_grader_template.yaml", "r") as yaml_file:
    answer_grader_templates = yaml.safe_load(yaml_file)

sys_msg_template: str = answer_grader_templates["sys_msg"]
human_msg_template: str = answer_grader_templates["human_msg"]

answer_grader_agent_prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(sys_msg_template),
    HumanMessagePromptTemplate.from_template(human_msg_template),
])

# Output Agent
with open(f"{current_dir}/output_agent_template.yaml", "r") as yaml_file:
    output_agent_templates = yaml.safe_load(yaml_file)

sys_msg_template: str = output_agent_templates["sys_msg"]
human_msg_template: str = output_agent_templates["human_msg"]

output_agent_prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(sys_msg_template),
    HumanMessagePromptTemplate.from_template(human_msg_template),
])