from app.utils.model_factory import get_local_model
from langchain_core.prompts import  PromptTemplate
from app.prompts import QUESTION_WRITER_SYSTEM_PROMPT

model = get_local_model(True)
template = PromptTemplate.from_template(QUESTION_WRITER_SYSTEM_PROMPT)

test_questions = [
    # Department abbreviations - fees related (1-10)
    "ds fees?",
    "cse fees",
    "it fee structure",
    "ce yearly fees",
    "ec admission fees",
    "ic course fees",
    "me semester fees",
    "pe tuition fees",
    "civil engineering fees",
    "ict hostel fees",

    # Department abbreviations - general queries (11-20)
    "cse block?",
    "ict hod",
    "me admission process",
    "ec faculty list",
    "it placement stats",
    "ce lab facilities",
    "ic course duration",
    "pe project guidelines",
    "civil department location",
    "ds hod name",

    # Full department names (21-25)
    "computer science engineering seats",
    "information technology syllabus",
    "mechanical engineering faculty",
    "electronics and communication labs",
    "instrumentation and control projects",

    # Generic queries - no department (26-35)
    "What is the yearly fee?",
    "hostel facilities",
    "admission process",
    "library timing",
    "placement statistics",
    "scholarship available",
    "bus route",
    "college canteen",
    "wifi password",
    "exam schedule",

    # Edge cases - ambiguous/short (36-42)
    "fees",
    "admission",
    "placement",
    "hostel",
    "timing",
    "location",
    "contact",

    # ACPC/STS related (43-45)
    "acpc registration",
    "sts login",
    "acpc merit list",

    # Mixed/Complex queries (46-50)
    "ds vs cse which is better",
    "it or ce for placement",
    "fees for ds and cse both",
    "cse and it difference",
    "mechanical or civil scope"
]

results = []
for question in test_questions:
    prompt = template.invoke({
        "query": question
    })
    result = model.invoke(prompt)
    # print("*"*50)
    # print("Question : ", question)
    # print("Rewrote : ", result.content)
    results.append({
        "question": question,
        "rewrote": result.content
    })

print(results)