import sqlite3
from pydantic import BaseModel, Field
from llama_index.core.tools import FunctionTool


db_path = "./database/mock_qna.db"
description = """
      Use this tool to extract the chapter information from the body of the input text, 
      when user wants to learn more about a particular chapter and requested to be asked 
      with a question to test his/her understanding.
      The format of the function argument looks as follow:
        It should be in the format with `Chapter_` as prefix.
        Example 1: `Chapter_1` for first chapter
        Example 2: For chapter 12 of the textbook, you should return `Chapter_12`
        Example 3: `Chapter_5` for fifth chapter
        Thereafter, the chapter_n argument will be passed to the function for Q&A question retrieval.
"""

class QnA_Model(BaseModel):
    chapter_n: str = Field(..., 
                           pattern=r'^Chapter_\d*$',
                           description=(
                                "which chapter to extract, the format of this function argumet"
                                "is with `Chapter_` as prefix concatenated with chapter number"
                                "in integer. For example, `Chapter_2`, `Chapter_10`."
                                "if no chapter number specified or user requested for random question"
                                "or user has no preference over which chapter of textbook to be tested"
                                "return `Chapter_0`"
                                )
                          )

def get_qna_question(chapter_n: str) -> str:
    """
      Use this tool to extract the chapter information from the body of the input text, 
      the format looks as follow:
        The output should be in the format with `Chapter_` as prefix.
        Example 1: `Chapter_1` for first chapter
        Example 2: For chapter 12 of the textbook, you should return `Chapter_12`
        Example 3: `Chapter_5` for fifth chapter
        Thereafter, the chapter_n argument will be passed to the function for Q&A question retrieval.
    """
    con = sqlite3.connect(db_path)
    cur = con.cursor()
    
    filter_clause = "" if chapter_n == "Chapter_0" else f"WHERE chapter='{chapter_n}'"
    sql_string = """SELECT id, question, option_1, option_2, option_3, option_4, correct_answer
                    FROM qna_tbl
                 """ + filter_clause
    
    res = cur.execute(sql_string)
    result = res.fetchone()

    id       = result[0]
    question = result[1]
    option_1 = result[2]
    option_2 = result[3]
    option_3 = result[4]
    option_4 = result[5]
    c_answer = result[6]

    qna_str  = "Question: \n" + \
               "========= \n" + \
                question.replace("\\n", "\n") + "\n" + \
               "A) " + option_1 + "\n" + \
               "B) " + option_2 + "\n" + \
               "C) " + option_3 + "\n" + \
               "D) " + option_4
    
    con.close()
    
    return qna_str

get_qna_question_tool = FunctionTool.from_defaults(
                            fn=get_qna_question,
                            name="Extract_Question",
                            description=description,
                            fn_schema=QnA_Model
)