from langchain_core.tools import tool
from pydantic import BaseModel, Field
from langchain_core.prompts import ChatPromptTemplate
from typing import Literal
from src.config.llm import llm_2_0 as llm
from typing import Optional
from src.config.vector_store import vector_store_lesson_content
from src.utils.logger import logger


@tool
async def extract_lesson_content(
    query_sentence: str, class_number: int, subject_name: str
) -> str:
    """Call vector store to retrieve documents based on query_sentence, class_number, subject_name

    Args:
        query_sentence (str): Query sentence
        class_number (int): Class number
        subject_name (str): Subject name
    Returns:
        str: Retrieved documents
    """
    # filter = {
    #     "class_number": class_number,
    #     "subject_name": subject_name,
    # }
    # logger.info(f"Filter: {filter}")
    # retriever = vector_store_lesson_content.as_retriever(
    #     search_type="similarity_score_threshold",
    #     search_kwargs={"k": 5, "score_threshold": 0.3},
    # )
    # documents = await retriever.ainvoke(query_sentence, filter=filter)
    # show_doc = " \n =============\n".join([doc.page_content for doc in documents])
    # logger.info(f"Retrieved documents: {show_doc}")
    # return {"documents": documents}
    return "Không có tài liệu"


@tool
def ChangeLesson(lesson_name: str, subject_name: str, class_number: int):
    """Khi người dùng đề soạn giáo án đề cập đến môn, bài học, lớp khác thì call ChangeLesson tool.
    Luôn hỏi confirm thông tin từ giáo viên trước khi call tool

    Args:
        lesson_name (str): Tên bài học
        subject_name (str): Tên môn học
        class_number (int): Số lớp học
    """


@tool
def EntryExtractor(class_number: int, subject_name: str, lesson_name: str):
    """Khi thu thập đủ thông tin class_number, subject, lesson thì call EntryExtractor tool. Không cần hỏi confirm

    Args:
        class_number (int): Số lớp học
        subject_name (str): Môn học
        lesson_name (str): Bài học
    """